Scaling Laws Meet Model Architecture: Toward Inference-Efficient LLMs

AI summary

Overview

Research area: Machine learning — large language model efficiency, scaling laws, and transformer architecture design.

Technical level: Intermediate. The paper assumes familiarity with transformer basics (attention, MLPs, KV cache) and with scaling-law formulations, but its argument is accessible to anyone who works on model deployment.

Scope: This paper augments classical scaling laws with architectural variables to identify transformer designs that are simultaneously accurate and cheap to run at inference time.

What This Paper Is About

Conventional scaling laws predict how a language model's training loss improves as you add parameters and training data, but they say nothing about how much the resulting model costs to serve. This omission is a problem because inference — not training — dominates the lifetime expense of a deployed model. The authors ask whether architectural choices such as hidden size, the balance between MLP and attention parameters, and grouped-query attention can be folded directly into a scaling law, so that practitioners can predict both accuracy and efficiency before spending compute.

Key Contributions

1. A conditional scaling law that incorporates architecture. The authors extend the Chinchilla formulation with two additional variables — normalized hidden size (d_model / √N) and the MLP-to-attention parameter ratio — using a two-step framework: first establish the Chinchilla-optimal loss as a reference point, then calibrate the loss of architectural variants relative to it.

2. Empirical characterization of architectural effects on inference cost. Through controlled ablations they show that larger hidden sizes (with fewer attention heads) and higher MLP-to-attention ratios reduce inference FLOPs and shrink the KV cache, improving throughput at a fixed parameter budget. GQA further improves throughput but does not behave as a smooth, continuous variable with respect to loss.

3. A search framework for Pareto-optimal architectures. Algorithm 1 combines the fitted scaling law with a constrained optimization (maximize inference efficiency subject to a loss ceiling) plus a bounded local search over feasible GQA values, which must divide the attention-head count.

4. Large-scale validation with trained models. Over 200 models from 80M to 3B parameters and 8B to 100B tokens were trained. The resulting designs (Panda-1B/3B, Surefire-1B/3B) beat LLaMA-3.2 baselines on nine downstream tasks while delivering up to 42% higher inference throughput (47% with SGLang on H200).

Main Findings

• U-shaped loss curves are consistent across model sizes. When other factors are held fixed, training loss as a function of normalized hidden size follows a consistent U-shape with nearly the same optimum at 80M, 145M, and 297M parameters. The same holds for the MLP-to-attention ratio. This means both factors have an interior optimum — pushing more parameters into attention or away from it both hurt.

• Bigger hidden size and higher MLP-to-attention ratios are more inference-friendly. Under a fixed non-embedding parameter budget, both choices reduce FLOPs and shrink the KV cache, improving tokens/second across batch sizes. This contradicts the recent tendency of open-weight models to allocate progressively fewer parameters to attention as they scale.

• The conditional scaling law predicts well out of distribution. Fitting on 80M models and evaluating on 145M, then on progressively larger data, yields low mean squared error and high Spearman rank correlation between predicted and actual loss.

• Simple separable calibrations are sufficient. Multiplicative and additive calibration forms perform comparably, and more complex non-separable joint formulations do not improve prediction. The two-step reference-and-calibration structure appears robust.

• Extreme ratios should be excluded from fitting. The MLP-to-attention ratios of real open-weight models fall between roughly 0.5 and 5. Including outliers such as 0.1 or 12.6 degrades rank correlation noticeably.

• Fit the law on models near the target scale. Fitting with 80M/145M/297M/1B data to predict 3B behavior worked worse than fitting with 1B data alone, suggesting the law's coefficients shift with model size. Fitting on models around one third of the target scale is often sufficient and sometimes preferable.

• The optimized models win on both axes. Panda-1B reaches 57.0 average downstream accuracy versus 54.9 for LLaMA-3.2-1B (a 2.1% gain), and Panda-3B reaches 62.5 versus 61.9. Surefire-1B and Surefire-3B satisfy a loss ceiling matched to the LLaMA baselines while achieving up to 42% higher throughput. Gains held across servers (vLLM and SGLang) and GPUs (A100 and H200).

Methodology in Plain English

The authors fix the number of layers and vary everything else. They begin with existing open-weight designs (LLaMA-3.2, Qwen3) and generate architectural variants by changing hidden size, the number of attention heads, and the MLP intermediate size, all while holding the total non-embedding parameter count constant. This isolates the effect of how parameters are arranged from how many there are.

They then train more than 200 small models (80M to 297M parameters) on up to 30B tokens to map out how architecture affects loss, and separately measure inference throughput with vLLM on A100 GPUs. Each architectural factor is varied one at a time to produce clean, conditional curves.

For the scaling law, they take the Chinchilla equation as a baseline "best possible loss" for a given model size and token count, then multiply or add a correction term that depends on hidden size and the MLP-to-attention ratio. The correction has a fixed functional form (a constant plus a log term plus an inverse term) chosen because it naturally produces U-shaped curves with sublinear growth. Coefficients are fit with the Levenberg-Marquardt least-squares algorithm.

Finally, to pick an architecture, they solve a constrained optimization: maximize measured inference throughput subject to predicted loss staying below a target. GQA is handled separately because it is discrete and erratic with respect to loss, so they enumerate the feasible values (it must divide the head count) and stop early when performance drops below a baseline.

Why This Matters

Impact on research. Scaling laws have long been treated as a tool for allocating training compute. This paper shows they can also be a tool for allocating inference cost, and that architecture is a first-class variable rather than a nuisance parameter. It also provides evidence against the common heuristic that attention should receive a shrinking share of parameters as models grow.

Real-world applications:

• LLM serving at scale. Cloud providers pay for every token generated; a 42% throughput gain translates directly into lower serving cost per request.
• Reasoning systems. Models that generate long chains of thought or use test-time search issue far more tokens than they consume in training, making inference efficiency disproportionately valuable.
• On-device and edge deployment. Small models at 1B–3B scale are exactly the regime studied here, and throughput and KV-cache footprint are the binding constraints on phones and laptops.
• Compute-constrained fine-tuning and research. The two-step framework lets a lab without frontier-scale hardware predict good architectures from cheap small-scale sweeps.

Industry relevance. Open-weight model families (LLaMA, Qwen, Gemma, Phi) converge on similar parameter counts but adopt very different internal shapes. This paper gives a principled way to choose among them, and demonstrates measurable gains against a widely deployed baseline, which is the kind of evidence that shapes architecture decisions in practice.

Future Directions

• Extend to Mixture-of-Experts. The work covers only dense models. MoE architectures change both the parameter accounting and the inference cost profile, and whether the same conditional framework applies is unresolved. The authors report preliminary MoE throughput measurements but no scaling law.

• Validate at 7B and beyond. Resource limits kept the largest validated model at 3B parameters. Whether the fitted coefficients and U-shaped optima persist at frontier scale is untested.

• Account for post-training. All analysis is on pretraining loss and zero-shot evaluation. Instruction tuning and reinforcement learning may shift which architectures are optimal, and the paper explicitly leaves this open.

• Make inference efficiency an analytical term. The current search evaluates throughput empirically on specific hardware and serving stacks. A hardware-aware analytical model of throughput as a function of architecture would remove the need for repeated measurement and make the optimization fully closed-form.

Target Audience

Machine learning engineers and researchers who train or serve language models and care about the cost-accuracy frontier — particularly those working on model architecture design, inference infrastructure, or deployment at the 1B–10B parameter scale. It is also useful for scaling-law researchers interested in extending the Chinchilla framework beyond raw parameter and token counts, and for practitioners choosing among open-weight model families who want a principled basis for architectural decisions rather than convention.

Authors’ abstract

Scaling the number of parameters and the size of training data has proven to be an effective strategy for improving large language model (LLM) performance. Yet, as these models grow increasingly powerful and widely deployed, the cost of inference has become a pressing concern. Despite its importance, the trade-off between model accuracy and inference efficiency remains underexplored. In this work, we examine how key architectural factors, hidden size, the allocation of parameters between MLP and attention (mlp-to-attention ratio), and grouped-query attention (GQA), influence both inference cost and accuracy. We introduce a conditional scaling law that augments the Chinchilla framework with architectural information, along with a search framework for identifying architectures that are simultaneously inference-efficient and accurate. To validate our approach, we train more than 200 models spanning 80M to 3B parameters and 8B to 100B training tokens, and fit the proposed conditional scaling law. Our results show that the conditional scaling law reliably predicts optimal architectural choices and that the resulting models outperform existing open-source baselines. Under the same training budget, optimized architectures achieve up to 2.1% higher accuracy and 42% greater inference throughput compared to LLaMA-3.2.