Gemstones: A Model Suite for Multi-Faceted Scaling Laws

Motivation

• Scaling laws can change substantially with width-depth ratios, training length, cooldown schedules, and other design choices, limiting transfer across implementations.Practitioners may apply laws derived from substantially different setups without knowing how sensitive the prescription is to those differences.Source: §1 · p. 1

Contributions

• Gemstones enables analysis of how architectural shape affects loss, benchmark performance, and training time.The model suite is designed to support open, granular studies beyond a single parameter-count dimension.Source: §1 · p. 2
• The authors release over 4,000 checkpoints trained on more than 10 trillion tokens across diverse architectural and training-hyperparameter axes.The suite spans model shapes, parameter counts, training tokens, learning rates, and cooldown schedules.Source: §1 · pp. 1–2

Main findings

• The paper highlights that fitting decisions can significantly change the slope of a scaling law, exposing fragility and common pitfalls in prior practice.The authors frame model selection and fitting procedure as central sources of prescription sensitivity.Source: §1 · pp. 1–2
• Deep models achieve lower loss and benchmark error at fixed FLOPs but require significantly more training time under typical open-source parallelism frameworks.This comparison is reported as the paper’s main model-shape result.Source: §1 · p. 2

Scaling laws

• Scaling laws seek the minimum language-model loss achievable under a constrained FLOP budget by balancing parameter count against training tokens.This frames the central optimization problem that the paper’s related work addresses.Source: §2 · p. 2

Scaling-law design

• Prior studies show that fitted laws can differ because of design choices such as parameter-count definitions, training duration, model selection, and excluded early-training data.These studies motivate treating scaling-law construction as sensitive to experimental setup rather than as a single fixed procedure.Source: §2 · p. 2

Downstream performance

• Research has modeled downstream task performance from task loss, perplexity, or overtraining, but predictability varies across tasks and compute scales.ARC, HellaSwag, and MMLU are highlighted elsewhere in the supplied material as the most predictable subset at smaller compute scales.Source: §2 · p. 2

Model shape

• Width-depth studies report different patterns, including logarithmic depth growth, modality-specific optimal aspect ratios, deeper-model advantages, diminishing depth benefits, and cases favoring shallower models.The literature therefore provides no single settled conclusion about the best aspect ratio.Source: §2 · pp. 2–3

Positioning

• The paper describes the aspect-ratio literature as extensive but inconclusive and seeks progress through fully open, reproducible, and extensible experiments.Gemstones is presented as a way to interpret prior findings while testing architecture choices directly.Source: §2 · p. 3

Architecture

• The study searches constrained model families from 50M to 2B parameters, varying width and depth while selecting standard, wide, and deep shapes.At smaller scales up to five shapes are trained, while larger scales use three representative aspect ratios.Source: §3 · p. 3, Figure 2 · p. 3

Training suite

• The main training runs use 350B tokens, while the suite also varies learning rates and cooldown schedules and saves checkpoints every 2B tokens.These choices create more than 4,000 checkpoints and support fitting-law sensitivity studies.Source: §3 · pp. 3–4

Fitting setup

• Scaling laws are fit from validation log perplexity, and the study also compares predictions with widely used open-source models on the same validation data.The fitting procedure uses the recorded model losses and a fixed held-out validation sample.Source: §3.1 · p. 4

Measurement

• The authors use a specialized FLOP counter because the rule FLOPs per token = 6 × parameters does not accurately capture extremely wide and narrow architectures.This makes the compute axis responsive to architectural shape rather than relying on a width-insensitive approximation.Source: §3.1 · p. 4

3.1 Fitting Scaling Laws

• Standard Approach 1 is poorly suited to the sparse two-dimensional model set, while isoFLOP binning is excluded because the models do not finely cluster into bins.The authors instead use the convex hull to accommodate lower data density in the sampled architecture space.Source: §3.1 · p. 4, Figure 3 · p. 5
• The convex-hull method fits only lower-hull vertices, naturally excluding suboptimal models and improving robustness when optimal points are sparse.The method addresses the two-dimensional width-depth search, where optimal models do not densely populate the lower envelope.Source: §3.1 · p. 4

4.1 Sizing Up Our Scaling Laws Against Prior Laws and Industry Models

• Convex-hull fitting omits strictly suboptimal points and can produce more reliable scaling-law fits when optimal points are sparse.It is designed to reduce variance from the ill-conditioned optimization created by asymptotically flat power-law curves.Source: §4.1 · pp. 5–6

4.2 Fragility and Pitfalls of Scaling Laws

• Reduced sampling, embedding-count choices, learning-rate schedules, cooldown schedules, and token coverage all noticeably change prescribed scaling laws.Reduced sampling has a particularly large impact on fitted slopes, while removing late or early checkpoints can matter more than fitting on the ablation datasets.Source: §4.2 · pp. 6–7, Table 1 · p. 6

4.3 Modeling Performance on Different Validation Sets and Downstream Benchmarks

• Validation losses change across Dolma, FineWeb, FineWeb-Edu, and DCLM, but model rankings remain stable and deeper models consistently achieve lower loss.The cross-dataset difference is described as approximately a y-axis shift, suggesting scaling laws may be fit on varied validation sets without retraining models.Source: §4.3 · p. 7
• Benchmark error fits tightly, whereas benchmark accuracy fits noisily because Gemstones vary substantially in width and depth.Both laws use benchmark results sampled every 10 billion training tokens; accuracy is modeled for ARC, HellaSwag, and MMLU.Source: §4.3 · pp. 7–8

4.4 The Width/Depth vs. Compute/Time Continuum

• Deeper models achieve lower loss and benchmark error per FLOP, while wider models can reach lower loss faster in GPU hours under the reported parallelism setup.The time-based preference for wider models may not generalize because the training scheme uses tensor and data parallelism without pipeline parallelism.Source: §4.4 · pp. 8–9
• At comparable FLOP budgets, benchmark accuracy generally increases with depth across 1B- and 2B-scale Gemstones.The pattern appears at 200, 250, 300, and 350 billion tokens and supports benefits from additional depth despite evidence that some deeper layers can be pruned.Source: §4.4 · p. 9

• Scaling-law outcomes depend strongly on often-overlooked design choices and on whether performance is measured by computation or training time.Source: §5 · p. 9
• Applicability is expected to decline in training setups very different from those used for the Gemstones.Source: §5 · p. 9
• The study keeps the attention head size fixed at 128 and constrains embedding widths to be divisible by 128.Source: A Model Implementation Details · p. 16

Unified Learning Rate Scaling Rule

• The proposed rule sets lreff = lrbase/(width × √depth), combining separate width and depth scaling factors.The rule is intended to balance update magnitudes across models varying in both dimensions.Source: §A.1 · p. 16

Empirical Validation

• Rescaling the learning-rate sweep by width × √depth collapses extreme-shape training curves onto a shared trend, supporting width-depth transfer.Source: §A.1 · p. 16, Figure 9 · p. 17

Flaws in the Gemstones

• Wider architectures occasionally show loss spikes, but rollbacks and minor learning-rate changes allow all models to train to 350B tokens without divergence.Source: §A.1 · p. 16

Ablation Study

• Reducing lrbase from 5 to 2.5 marginally increases losses but leaves scaling-law fits robust.Source: §A.1 · p. 16

Scalable Parameter Initialization Rules

• The Gemstones use a parameter-type- and layer-dependent truncated-normal initialization to promote stable training and learning-rate transfer across scales.Attention and MLP projection variances depend on width and the layer index, with the MLP rule incorporating the 4× width-to-intermediate-dimension factor.Source: §A.1 · p. 16

• The displayed prescription uses coefficients a = β/(α + β) and b = α/(α + β).Source: C Data Sampling · p. 18
• The appendix adds width-512 models at depths 11 and 14 alongside the originally selected depths 12 and 13.These additions create four closely spaced width-512 models for analyzing finer-grained depth differences.Source: C Data Sampling · p. 18

• The method represents each model as an (x, L) pair, where x is measured in FLOPs or GPU hours and L is loss.Source: D Mathematical Definition of Our… · p. 18
• Figure 10 distinguishes all feasible model shapes within ±5% as circles from selected models as stars, including extra width-512 points.Source: Figure 10 · p. 19
• The convex hull consists of linear interpolations between dataset points, and the lower convex hull defines the fitted trade-off function.Source: D Mathematical Definition of Our… · p. 19
• Figure 7 visualizes the convex hull as a red line through colored vertices, with gray lines representing other dataset points.Source: D Mathematical Definition of Our… · p. 19

Benchmark extrapolation

• Error extrapolates tightly from models below 1.8B to models above 1.8B, whereas accuracy extrapolation for ARC, HellaSwag, and MMLU fits poorly.The error result is shown in Figure 11, while the accuracy result is shown in Figure 12.Source: Figure 11 · p. 20, Figure 12 · p. 20

Extrapolation robustness

• 0.63% ARE results when fitting on smaller models and testing above 1.8B parameters, remaining within the typical 4% threshold for distinguishing modeling choices.Fitting on all checkpoints gives 0.68% ARE; the reported threshold comes from Choshen et al. [2024].Source: E Extrapolation to Larger Models · pp. 19–20

Model-selection sensitivity

• Leave-one-out refitting measures how omitting each model shape changes predicted tokens-per-parameter ratios under Approaches 1 and 3.The analysis reports minimum and maximum ratios across trials and compares them with the all-data fit.Source: F Leave-One-model-Out Analysis · p. 20, Figure 13 · p. 21

Sampling sensitivity

• Sampling every 2B rather than 10B tokens changes Approach 1 more than Approach 3 because the former relies on a sparser Pareto hull.The paper therefore uses 2B-token sampling for Approach 1 and 10B-token sampling for Approach 3 in later plots.Source: §G.1 · pp. 20–21

Dataset sensitivity

• Laws fitted on DCLM and FineWeb-Edu differ minimally from the main laws fitted on Dolma.The comparison uses data recorded every 10B tokens for DCLM and FineWeb-Edu.Source: §G.1 · p. 21

Scaling-law comparisons

• Benchmark prescriptions differ little across FineWeb-Edu, DCLM, and Dolma, while sampling frequency matters more for Approach 1 than Approach 3.The dataset comparison is shown in Figure 15 alongside fits using 2B- and 10B-token sampling.Source: Figure 15 · p. 22

Architectural inefficiency

• Skinny architectures can overspend about 50% in FLOPs but more than 200% in GPU hours relative to the best width-depth configurations.Under the reported setup, the time spent training one highly suboptimal model could train three optimal-width-depth models.Source: I The Price of Stepping Off the… · p. 22, Figure 17 · p. 24

Benchmark versus training time

• Depth often improves benchmark accuracy at similar FLOP budgets, but this does not eliminate the training-time advantage observed for wider models.The benchmark trend is reported for many tasks, while the time comparison is based on the reported training setup.Source: I The Price of Stepping Off the… · p. 22, Figure 16 · p. 23

Overtraining

• Overtraining increases predicted loss, but the loss curves are fairly flat, especially at high FLOP counts where smaller models can be trained on more tokens efficiently.The analysis fixes FLOP budgets, varies the overtraining factor, and optimizes model shape at each budget.Source: I The Price of Stepping Off the… · p. 22, J Scaling Laws Predict That… · p. 23, Figure 18 · p. 24

Practical trade-offs

• The combined results suggest that trading some FLOP optimality for time optimality can produce only a marginal loss difference.This connects the GPU-hour advantage of wider models with the small loss penalty predicted for overtraining smaller models.Source: I The Price of Stepping Off the… · p. 22, J Scaling Laws Predict That… · pp. 23–24

Measurement choice

• Validation loss is used for fitting because global training loss shows slight jumps when training switches between data chunks.The validation-loss curves are described as smoother for fitting all laws.Source: K Training · p. 25

Training instabilities

• Extreme patched models still mostly diverge after 350B tokens, limiting conclusions for the widest architectures.The authors could not rerun all experiments because the original experiments consumed substantial compute resources.Source: §K.2 · p. 25

Incomplete ablations

• The cooldown ablation differs from the main runs for the widest models and omits some models and a 100B-token checkpoint.The discrepancies arise because the ablation predates the patched learning-rate scaling rules.Source: §K.2 · pp. 25–26

Supplementary analyses

• The appendix includes extended Approach 1 plots and fits on learning-rate and cooldown ablation datasets.These supplementary figures extend the main prescriptions and show how the ablations are incorporated into fitting analyses.Source: Figure 20 · p. 26, §L.1 · p. 26, §L.2 · p. 26

Optimization stability

• A Huber-loss delta of 10^-5 is unstable for smaller grid searches, while adding more fitting tokens generally stabilizes the optimized exponents.The relationship is examined for data restricted to 100B tokens and for data extending to 350B tokens.Source: §L.4 · pp. 26, 28, Figure 23 · p. 28

FLOP accounting

• The common FLOPs-per-token approximation, 6 × parameters, can significantly miscount computation for architectures with nonstandard aspect ratios.The paper introduces a specialized counting function because wide and narrow architectures have different operation counts.Source: M FLOP counting matters · p. 28, Figure 24 · p. 28