Scaling Laws, Fine-Tuning and In-Context Learning

The compute-optimal allocation computation, LoRA's rank argument, and what is actually known about how in-context learning works.

LLMsResearch⏱ 3 hours9 min readscaling lawschinchillacompute optimallorafine-tuning

5 problems →

AssumesLanguage Modelling, Cross Entropy and Perplexity · Why Does Deep Learning Optimization Work?

Compute accounting

PropositionThe 6ND FLOP estimate

Training a dense transformer with $N$ non-embedding parameters on $D$ tokens costs approximately

C \approx 6ND \text{ FLOPs}.

Compute-optimal scaling

PropositionCompute-optimal allocation under a power-law loss

Suppose the fitted loss is

\mathcal{L}(N,D) = \mathcal{L}_\infty + \frac{A}{N^\alpha} + \frac{B}{D^\beta},

and the compute budget is $C = 6ND$ . Then the minimiser of $\mathcal{L}$ subject to fixed $C$ satisfies

N^\star \propto C^{\frac{\beta}{\alpha+\beta}}, \qquad D^\star \propto C^{\frac{\alpha}{\alpha+\beta}} .

In particular, if $\alpha = \beta$ then $N^\star$ and $D^\star$ both scale as $\sqrt C$ — that is, in fixed proportion to each other.

Key idea

The fitted exponents decide the allocation, and the empirical finding of Hoffmann et al. is that $\alpha\approx\beta\approx0.34$ , so $N^\star$ and $D^\star$ scale together as $C^{1/2}$ — roughly 20 tokens per parameter. The earlier Kaplan et al. analysis, fitting the same functional family with different methodology (notably a fixed learning-rate schedule that disadvantaged the longer-trained runs), inferred $N^\star \propto C^{0.73}$ and therefore recommended models much larger and data much smaller.

The lesson is worth stating plainly: the form of the law is empirical, the exponents are fitted, the allocation rule is a Lagrange computation on top of the fit, and the whole chain is only as reliable as the fit. A methodological change to the fit moved the recommendation by more than an order of magnitude. Compute-optimality is an engineering estimate, not a theorem.

What this does not say

The computation optimises training loss at fixed training compute. It says nothing about inference cost, and inference is where most lifetime compute is spent for a deployed model. For a model that will serve many queries, a smaller model trained on more tokens than the compute-optimal point is often the better system-level choice, even though it is suboptimal by this criterion. Recent practice reflects this, and the divergence is not a contradiction of the analysis but a different objective.

Parameter-efficient fine-tuning

PropositionLoRA parameter count and the rank argument

Replace a frozen weight $W_0\in\R^{d\times k}$ by

W = W_0 + \frac{\alpha}{r}BA, \qquad B\in\R^{d\times r},\ A\in\R^{r\times k},

with $A$ initialised Gaussian and $B$ initialised to zero. Then

the number of trainable parameters is $r(d+k)$ instead of $dk$ , a factor $dk/(r(d+k))$ reduction, which for $d=k$ is $d/(2r)$ ;
at initialisation $W = W_0$ , so training starts from the pretrained function exactly;
after training, $BA$ can be folded into $W_0$ , so inference cost is unchanged.

Key idea

Why should a low-rank update suffice? The multi-head decomposition gives a partial answer: the operative objects in an attention layer are already the low-rank matrices $W_Q(W_K)^\top$ and $W_VW_O$ , each of rank at most $d/h$ . An update that changes attention behaviour therefore need not be full rank to be expressive in the relevant sense.

That is an argument, not a theorem. The empirical finding that $r$ as small as $4$ – $16$ recovers most of full fine-tuning's gain on many tasks is stronger evidence, and the accompanying observation that the gradient of the loss with respect to $W$ is low rank early in fine-tuning (being a sum of a few outer products per batch) is suggestive. There is no result establishing that the optimal fine-tuning update is approximately low rank in general, and tasks requiring substantial new knowledge do show a gap.

PropositionMixture of experts decouples parameters from FLOPs

Replace an MLP block of $P$ parameters by $E$ expert blocks of $P$ parameters each, with a router sending every token to the top $k$ experts. Then the parameter count is $EP$ while the per-token FLOPs are proportional to $kP$ . The ratio $E/k$ is the sparsity factor.

What this does not say

The decoupling is real but the costs move rather than vanish. All $EP$ parameters must be resident in memory (or fetched), so memory bandwidth rather than FLOPs becomes the constraint. Routing is discrete and therefore non-differentiable, so training needs auxiliary load-balancing losses to prevent expert collapse — a single expert receiving all tokens is a stable and useless equilibrium. And the scaling laws fitted for dense models do not transfer directly: a sparse model's loss depends on $E$ , $k$ and $P$ separately, and the appropriate scaling law has to be refit.

In-context learning

Empirical observationIn-context learning

Empirical observation — a regularity seen in experiments. No theorem on this page establishes it in general.

A pretrained language model, with no parameter updates, performs tasks specified only by examples in its prompt — including tasks it was not explicitly trained on — with accuracy improving in the number of in-context examples.

TheoremTransformers can implement gradient descent in context

There exist explicit weight settings for a linear-attention transformer such that the forward pass on a prompt containing $(x_1,y_1),\dots,(x_m,y_m),x_{\text{query}}$ computes the prediction of $\ell$ steps of gradient descent on the least-squares objective over those examples, where $\ell$ is the number of layers.

Quoted from von Oswald et al.; Garg et al. independently showed that transformers trained on in-context regression tasks match or exceed the performance of least squares and of specialised estimators on those tasks. The construction is explicit and checkable, which makes it a genuine theorem about expressivity.

What this does not say

What these results establish and what they do not:

Established. A transformer can implement an optimisation algorithm in its forward pass, and transformers trained specifically on function-class regression tasks do behave like good estimators for those classes. In-context learning is therefore not mysterious in principle — there is a mechanism a transformer can implement.
Not established. That this is what a language model trained on text is doing. The constructions use linear attention, specific weight patterns, and a prompt format that encodes a clean regression problem; a language model's in-context behaviour on, say, translation or arithmetic has not been shown to be gradient descent on anything.

A separate and complementary line of evidence is mechanistic. Olsson et al. identify induction heads — attention heads that implement the rule "find the previous occurrence of the current token and copy what followed it" — and show that their formation during training coincides with a sharp improvement in in-context learning ability, across model scales. This is strong evidence about a mechanism contributing to in-context learning, established by circuit-level analysis and ablation.

The honest summary: the phenomenon is real and robust, two distinct partial mechanisms are documented (implicit optimisation in constructed settings, induction heads in real models), and there is no account of the general case. This sits in the same category as the other entries in the ledger.

Problems

Derive the compute-optimal allocationProof · Intermediate · 30 min The 6ND FLOP estimateNumerical · Beginner · 15 min LoRA parameter accountingNumerical · Beginner · 15 min Mixture-of-experts sparsity accountingNumerical · Intermediate · 30 min

References

PaperJared Kaplan and collaborators (2020). Scaling Laws for Neural Language Models
PaperJordan Hoffmann and collaborators (2022). Training Compute-Optimal Large Language Models
PaperEdward J. Hu and collaborators (2021). LoRA: Low-Rank Adaptation of Large Language Models
SurveyCatherine Olsson and collaborators (2022). In-context Learning and Induction Heads
PaperJohannes von Oswald and collaborators (2023). Transformers learn in-context by gradient descent
PaperShivam Garg, Dimitris Tsipras, Percy Liang and Gregory Valiant (2022). What Can Transformers Learn In-Context? A Case Study of Simple Function Classes
PaperNoam Shazeer and collaborators (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

These works are cited as sources to read, not reproduced. All exposition, examples and problems on this site are written independently; where an argument is standard, the citation points to a canonical treatment.

Results stated on this page

PropositionCompute-optimal allocation under a power-law lossproof includedMinimising the fitted loss at fixed compute gives parameter and token budgets that scale as powers of the compute, with the exponents determined by the loss exponents.
PropositionThe 6ND FLOP estimateproof includedTraining a dense transformer costs about six FLOPs per parameter per token, two for the forward pass and four for the backward.
PropositionLoRA parameter count and the rank argumentproof includedA rank-r update to a d-by-k matrix costs r(d+k) parameters instead of dk, and composes with the existing weight at no inference cost.
Empirical observationIn-context learningModels perform tasks specified only in the prompt; several mechanisms are demonstrated in restricted settings and none accounts for the general phenomenon.
TheoremTransformers can implement gradient descent in contextThere are explicit weight constructions under which a transformer's forward pass performs steps of gradient descent on an in-context regression problem.
PropositionMixture of experts decouples parameters from FLOPsproof includedRouting each token to k of E experts makes the compute proportional to k rather than E, at the cost of memory and load-balancing.

Problems (5)

All problems

Compute accounting

Compute-optimal scaling

Parameter-efficient fine-tuning

In-context learning

Problems

References

Results stated on this page

Problems (5)

Papers that use this mathematics