Is More Data Worth the Cost? Dataset Scaling Laws in a Tiny Attention-Only Decoder

AI summary

Overview

Research area: Machine learning / large language model pretraining, specifically empirical scaling laws and data-efficiency for decoder-only Transformers.

Technical level: Intermediate. Readers need basic familiarity with Transformer decoders, tokenization, cross-entropy loss, and the idea of power-law scaling laws.

Scope: A controlled, small-scale empirical study of how validation performance changes as training-dataset size grows in a fixed-capacity, attention-only decoder derived from the 117M GPT-2 design.

What This Paper Is About

Scaling laws predict that language model performance improves smoothly as data and compute grow, but they were derived at industrial scale where model size, dataset size, and compute are scaled together. That makes it hard to attribute performance trends to any single factor, and it leaves open whether the same behavior holds for very small models. This paper freezes model capacity almost entirely and varies only dataset size, asking whether the diminishing returns predicted by scaling laws still appear in a tiny attention-only decoder, and how much data is actually needed to approach full-data performance.

Key Contributions

1. A fixed-capacity attention-only architecture. A decoder-only model built on the 117M GPT-2 design with pretrained token embeddings and the output projection frozen, and the MLP sublayers removed entirely, so only self-attention parameters are trained.
2. An analysis of dataset representativeness and robustness. Token-distribution similarity between subsets and the full corpus is measured with Jensen–Shannon divergence, and training outcomes are checked across multiple random seeds.
3. A controlled dataset-scaling study. Training on nested power-of-two subsets ranging from 2^7 to 2^17 sequences, under both a fixed 150-epoch schedule and a fixed number of optimization steps.
4. A characterization of the cost–performance trade-off. Quantifying how much training time and what fraction of the data are needed to reach given fractions of full-data accuracy.

Main Findings

• Diminishing returns appear even in tiny models. Validation token-level accuracy improves smoothly with subset size and then saturates, matching the qualitative shape of scaling-law behavior despite the model being far smaller than industrial-scale LLMs.
• Roughly 30% of the data reaches roughly 90% of full-data accuracy. In the paper's abstract-level summary, using about 30% of the training data is sufficient to reach approximately 90% of the full-data validation token-level accuracy.
• The detailed cost–performance table. Under the fixed training schedule, 80% of full-data accuracy is reached at subset 2^13 (7% of data, 8 hours, 7% of cost); 90% at 2^15 (29% of data, 32 hours, 26% of cost); 95% at 2^16 (56% of data, 60 hours, 49% of cost); and 100% at 2^17 (100% of data, 123 hours, 100% of cost).
• Distributional convergence is fast. Jensen–Shannon divergence between subset and full-corpus token distributions falls below 0.003 for subsets of 2^12 sequences and above, meaning moderate subsets are effectively indistinguishable from the full dataset. Only the smallest subsets (2^7–2^8 sequences) show measurable distributional deviation.
• Results are stable across seeds. Across five seeds at epoch 1, validation loss, perplexity, and accuracy show small standard deviations that shrink as subset size grows (for example, 2^16 gives validation loss 6.34 ± 0.03, perplexity 567 ± 16, accuracy 11.87 ± 0.20, compared with 2^7 at 11.24 ± 0.07, 76,320 ± 4973, and 0.40 ± 0.12).
• Under fixed compute, extra data alone is not enough. When all configurations are trained with the same number of optimization steps, the same early improvement and saturation pattern appears, and the largest subsets no longer provide consistent improvements.
• The pattern replicates on a second dataset. A repeat of the subset-scaling experiment on WikiText-103 shows similar scaling dynamics.
• The setting is data-limited, not model-limited. Applying the Kaplan et al. (2020) joint scaling law with α_N = 0.076, α_D = 0.103, N_c = 6.4 × 10^13, and D_c = 1.8 × 10^13 to a model of N = 4.0 × 10^7 (40 million) parameters and D ≈ 1.34 × 10^8 tokens (134 million) yields a predicted test loss of approximately 3.46 nats/token. This is much closer to the infinite-model limit of approximately 2.96 than to the infinite-data limit of approximately 3.37, indicating performance is constrained mainly by data.
• Empirical saturation is faster than the Kaplan prediction. The empirical loss saturates substantially earlier than the scaling-law curve, which the authors attribute to fixed-epoch training where not all subsets are trained to convergence.

Methodology in Plain English

The researchers took the standard 117M GPT-2 architecture and stripped it down. They kept the pretrained token embeddings and the tied output projection but froze both, so those weights never change during training. They then deleted the feed-forward (MLP) sublayers from the decoder block. What remains is a single decoder layer with a single attention head and an embedding dimension of 768, operating on sequences of 1024 tokens. The model holds 79,556,352 parameters in total, of which only 2,361,600 are trainable.

Frozen embeddings mean the model cannot relearn word representations, so any performance differences across experiments come from learning contextual relationships inside self-attention rather than from representation learning. The researchers note that in preliminary runs where the embedding layer was trained jointly, they saw rapid overfitting and unstable training with apparent gains driven mainly by embedding adaptation.

Data came from the All the News 2.0 corpus, roughly 2.7 million news articles from 27 U.S. publications covering 2016 to 2020, chosen because it was not part of the original GPT-2 training corpus. Articles shorter than 500 tokens were filtered out, which reduced padding to approximately 28% of tokens. The final setup used 131,072 training sequences and a fixed held-out validation set of 20,000 sequences, tokenized with the standard GPT-2 tokenizer. Because GPT-2 has no padding token and adding one would break compatibility with the pretrained embedding matrix, the end-of-sequence token was reused for padding, masked during self-attention, and assigned the ignore label (-100) in the loss and accuracy computation.

Training subsets were nested: a single random permutation of the full training set was generated with a fixed seed, and each subset was defined as the first 2^k sequences for k from 7 to 17. This produces subsets from 128 sequences up to 131,072 sequences, which Table 2 lists as spanning 0.1% to 100% of the full dataset and 0.13 million to 134.22 million tokens, with token-to-parameter ratios from 0.056 to 56.86.

All runs used a single Tesla V100-SXM3-32GB GPU with PyTorch (CUDA), the AdamW optimizer at a learning rate of 3 × 10^-4, a fixed batch size of 16, 150 epochs, no dropout, no learning rate scheduling, and no early stopping. Each optimization step processes 16 × 1024 = 16,384 tokens, and steps per epoch equals the subset size divided by 16. The authors note that a commonly cited compute-optimal heuristic would correspond to approximately 47 million training tokens for 2.36 million trainable parameters, so the smaller subsets operate deep in a data-limited regime. Evaluation used masked cross-entropy loss, token-level accuracy over non-padding positions, and perplexity computed as exp of the token-averaged validation loss, measured before training, every epoch for the first 20 epochs, then every 10 epochs, and at the final epoch, with the best checkpoint selected by minimum validation loss.

Why This Matters

Impact on research. The paper shows that scaling-law-shaped diminishing returns can be reproduced and studied in a tiny, component-isolated model, which gives researchers a cheap experimental substrate for testing scaling hypotheses. It also supports subset-based pretraining with early stopping as a principled strategy rather than a shortcut, since the scaling-law analysis indicates performance here is data-limited rather than capacity-limited.

Real-world applications (as implied by the paper's framing):

• Planning pretraining budgets in compute- and data-restricted environments, where the reported cost–performance table (7% of data for 80% of accuracy, 29% for 90%) guides how much data to acquire or process.
• Rapid, economical base-model training in small research labs, where the paper specifically cites small research labs and exploratory model development as the target setting.
• Deciding whether to invest in more data or more model capacity: the scaling-law comparison suggests that additional data is the more valuable lever in this regime.
• Using distributional diagnostics (Jensen–Shannon divergence, per-token relative probability error) to decide whether a conveniently sized subset is representative enough of a full corpus before committing compute.

Industry relevance. The finding that most achievable accuracy is captured early, while training cost rises near-linearly, directly informs the trade-off between dataset size and training cost that pretraining teams negotiate. The architectural reduction to 2,361,600 trainable parameters also demonstrates that competitive experimental throughput is possible without full-capacity models.

Future Directions

• Test the effect of training to convergence. The paper explicitly states that the fixed-epoch schedule leaves larger subsets untrained to convergence, which it blames for the stronger-than-predicted diminishing returns. Retraining subsets to convergence would separate data-scaling effects from training-budget effects.
• Extend beyond a single architecture. The study isolates one reduced configuration: one decoder layer, one attention head, frozen embeddings and output projection, no MLP. Whether the same subset-scaling curve holds at other depths, head counts, or with partially unfrozen components is not reported.
• Broaden dataset and domain coverage. Results come from All the News 2.0 with a replication on WikiText-103. Behavior on other domains, languages, or corpora with different token distributions is not established.
• Strengthen robustness testing at the largest subsets. The robustness table covers five seeds at epoch 1 for subset levels 2^7 through 2^16; the largest subset and the full 150-epoch schedule rely on a single random seed.

Target Audience

This paper benefits machine learning researchers and engineers working on data-efficient pretraining, particularly those with limited compute budgets. It is also useful for practitioners designing pretraining data pipelines who need to decide how large a training subset must be, for students and smaller labs looking for a reproducible small-scale experimental setup, and for anyone studying scaling laws who wants a controlled, component-isolated complement to large-scale industrial studies.

Authors’ abstract

Training Transformer language models is expensive, as performance typically improves with increasing dataset size and computational budget. Although scaling laws describe this trend at large scale, their implications in controlled, smaller-scale settings remain less explored. In this work, we isolate dataset-size effects using a strongly reduced attention-only decoder architecture. By training on progressively larger power-of-two subsets, we observe smooth performance improvements accompanied by clear diminishing returns, consistent with scaling-law behavior. Using only about 30% of the training data is sufficient to reach approximately 90% of the full-data validation token-level accuracy. These results provide actionable insights into dataset scaling in a controlled, component-isolated setting and offer practical guidance for balancing dataset size and computational cost in compute- and data-restricted environments, such as small research labs and exploratory model development.