Algorithm Foundations

A profile showing low utilization is often misdiagnosed as a hardware limit when the cause is algorithmic: a matrix multiply shape that wastes Tensor Cores, a layout that breaks contiguity, or activation memory held for backpropagation. This appendix collects the computational complexity models, memory footprint derivations, numerical-stability rules, and proofs that the book’s performance chapters lean on, enabling readers to evaluate algorithmic trade-offs before committing code to accelerators. It assumes familiarity with linear algebra and with the iron law introduced early in the book.

How to Use This Appendix

This appendix is designed as a reference. Reach for it when translating a profiler symptom (“slow matmul,” “shape mismatch,” “OOM during training”) into a concrete computational or memory cause. Each symptom or need maps to one section (table 1), and conventions throughout follow the book-wide notation (for example, \(\BatchSize\) is reserved for batch size and \(\Bandwidth\) denotes bandwidth). Every section names the chapter that teaches its concept in context; the chapter carries the argument, and this appendix carries the derivation or the reference formula.

Table 1: Where to Look: Each row pairs a profiler symptom or a design question with the section that carries its derivation or reference formula.

Symptom or need	Section
A general matrix multiply (GEMM) runs far below peak throughput	section 1.1.3
Sparsity is proposed to save memory or compute	section 1.1.4
A layer’s parameters and FLOPs are needed before hardware is chosen	section 1.1.5
Tensor code fails on shape or contiguity, or a copy dominates a loop	section 1.2.1
A broadcast allocates a much larger tensor than intended	section 1.2.2
Gradients look wrong, or the backward pass costs more than expected	section 1.3.2
Training runs out of memory	section 1.3.3
A chain of elementwise operations is memory bound	section 1.3.4.2
Softmax overflows, or the loss becomes NaN in reduced precision	section 1.3.5
A bound on gradient cost relative to the forward pass	section 1.4.1
Attention memory or traffic grows with sequence length	section 1.4.2
A compute budget must be split between parameters and training tokens	section 1.4.3

Neural networks turn linear algebra into learned behavior. Matrices transform activations, tensor layouts determine how data move through memory, and backpropagation carries error signals through the resulting computation. These foundations support the deep learning treatment in Neural Computation, the framework internals in ML Frameworks, and the training strategies in Model Training.

Linear Algebra

Shrinking a square matrix multiply from dimension 4,096 to 64 drops its arithmetic intensity more than tenfold below an H100’s ridge point, turning a compute-bound layer into a bandwidth-bound stall without changing a line of model logic. Most of the compute in a deep learning system is matrix multiplication, and although frameworks such as PyTorch hide the raw math, performance engineering still depends on the underlying linear algebra. Building on how numbers are stored in Numerical Representations, this section focuses on how they are manipulated.

Systems Perspective 1.1: Why this matters

Many dense neural networks, especially transformers and large convolutional neural networks, spend most of their compute in matrix multiplication or GEMM-like kernels: a self-attention block executes query, key, value, and output projection GEMMs alongside two attention-matrix products for scores and weighted values, and the layer’s feed-forward block adds two or three more. Because those kernels dominate execution, GEMM performance characteristics explain why batch size governs throughput, why aligned layer dimensions run faster than unaligned ones, and whether a dense workload will be compute bound or bandwidth bound before any profiler trace runs.

Tensor operations and einsum notation

Tensor operations in deep learning contract across batch, head, sequence, and hidden dimensions at once, and Einstein summation¹ notation writes those contractions as a single index string rather than a chain of nested loops and transposes (implemented as torch.einsum in PyTorch and np.einsum in NumPy). Matrix multiplication \(\GemmOutputMatrix= \GemmLeftOperand\GemmRightOperand\) becomes: \[ \GemmOutputEntry = \sum_k \GemmLeftOperandEntry \GemmRightOperandEntry \]

¹ Einstein summation convention: Repeated indices in a product are implicitly summed over, eliminating explicit summation signs. ML frameworks adopted the convention because it concisely expresses arbitrary tensor contractions in a single string.

In einsum notation, this is ik,kj->ij. The same notation handles batched multi-head attention, bhid,bhjd->bhij (batch \(\BatchAxisLabel\), head \(\HeadAxisLabel\), query sequence \(\QueryAxisLabel\), key sequence \(\KeyAxisLabel\), and head dimension \(\HeadDimensionAxisLabel\)). The head dimension \(\HeadDimensionAxisLabel\) appears in both inputs but not in the output, so it is the index that is summed.

The dot product as similarity

Comparing two token embeddings or feature vectors requires a scalar score that grows when the vectors point in the same direction and shrinks when they are orthogonal. The dot product \(\InnerProductLeftVector\cdot \InnerProductRightVector= \sum \InnerProductLeftComponent_i \InnerProductRightComponent_i\) provides that score because it equals \(\|\InnerProductLeftVector\| \|\InnerProductRightVector\| \cos \VectorAngle\): for nonzero vectors, its sign distinguishes acute, right, and obtuse angles, while its magnitude also depends on both vector norms.

In attention mechanisms, query (\(\AttentionQueryMatrix\)) and key (\(\AttentionKeyMatrix\)) vectors are dot-produced to compute a similarity score that determines how much each token attends to every other token. The resulting attention weights are then used to form a weighted combination of value (\(\AttentionValuesMatrix\)) vectors, so every token can draw on every other token. Attention: Dynamic Processing develops attention itself, including the cost of computing one score for every query-key pair.

General matrix multiply (GEMM)

Whether a dense layer saturates an accelerator’s Tensor Cores or stalls on memory bandwidth depends on the three dimensions of its matrix multiply. Most training and inference FLOPs execute as general matrix multiplies (GEMMs):² for matrices of size \(\GemmRowDimension{\times}\GemmInnerDimension\) and \(\GemmInnerDimension{\times}\GemmColumnDimension\), a GEMM performs \(2\GemmRowDimension \GemmColumnDimension \GemmInnerDimension\) floating-point operations (counting multiply-accumulate as two operations).

² General matrix multiply (GEMM): The BLAS (Basic Linear Algebra Subprograms) family began with vector routines standardized in 1979 (Lawson et al. 1979); GEMM was standardized later as a Level 3 routine (Dongarra et al. 1988). The “GE” prefix stands for “general,” as opposed to symmetric or triangular forms. GEMM computes \(\GemmOutputMatrix= \GemmProductScale\GemmLeftOperand\GemmRightOperand+ \GemmAccumulatorScale\GemmOutputMatrix\) and is a performance-critical routine in deep learning. Matrix operations explains how GEMM shapes determine training throughput.

Lawson, Charles L., Richard J. Hanson, David R. Kincaid, and Fred T. Krogh. 1979. “Basic Linear Algebra Subprograms for Fortran Usage.” ACM Transactions on Mathematical Software 5 (3): 308–23. https://doi.org/10.1145/355841.355847.

Dongarra, Jack J., Jeremy Du Croz, Sven Hammarling, and Richard J. Hanson. 1988. “An Extended Set of FORTRAN Basic Linear Algebra Subprograms.” ACM Transactions on Mathematical Software 14 (1): 1–17. https://doi.org/10.1145/42288.42291.

The arithmetic intensity of GEMM scales linearly with matrix dimension. For square \(\SquareGemmDimension{\times}\SquareGemmDimension\) matrices in FP16 (2 bytes/element), the ideal \(\GemmAccumulatorScale=0\) bound reads \(\GemmLeftOperand\) and \(\GemmRightOperand\) once and writes \(\GemmOutputMatrix\) once: \[\text{Intensity} = \frac{\Operations}{\DataVolume} = \frac{2\SquareGemmDimension^3}{3\SquareGemmDimension^2 \times 2} = \frac{\SquareGemmDimension}{3}\text{ FLOP/byte}\]

Reading the old \(\GemmOutputMatrix\) when \(\GemmAccumulatorScale\ne 0\) changes the intensity to \(\SquareGemmDimension/4\) FLOP/byte. Three consequences follow for GEMM execution:

• Larger batches can improve efficiency: Batching increases effective matrix dimensions when it forms larger GEMMs or enables reuse; merely queuing independent small GEMMs does not raise intensity.
• Aligned dimensions help: Hardware Tensor Cores are optimized for precision- and architecture-specific tile multiples. Dimensions that align with these multiples avoid padding overhead and improve kernel efficiency, but they do not need to be powers of two.
• Small matrices are inefficient: A square GEMM with \(\SquareGemmDimension =\) 64 has intensity 64/3 ≈ 21.3 FLOP/byte, well below the ridge point of the book’s reference accelerator, the H100 (295.2 FLOP/byte). Its roofline cap is ~7.2 percent of peak; launch and tiling overhead can lower realized throughput. Performance Bounds derives the ridge point and the compute-bound and bandwidth-bound regimes on either side of it, which is the comparison each of these consequences relies on.

Sparse matrix formats

When most elements in a matrix are zero, specialized storage formats avoid wasting memory on zeros and enable computations that skip them entirely. The compressed sparse row (CSR) format uses three arrays:

• Values: The nonzero elements, stored in row order
• Col_Idx: The column index of each nonzero element
• Row_Ptr: The starting position in Values for each row (length = num_rows + 1)

CSR is useful for sparse feature matrices in recommendation pipelines and for pruned model weights. For a matrix with \(\SparseMatrixEntryCount\) elements, \(\SparseMatrixRowCount\) rows, and \(\NonzeroCount\) nonzeros, CSR uses \(\mathcal{O}(\NonzeroCount+ \SparseMatrixRowCount)\) storage instead of \(\mathcal{O}(\SparseMatrixEntryCount)\); when \(\NonzeroCount\) is large relative to \(\SparseMatrixRowCount\), this is often summarized as \(\mathcal{O}(\NonzeroCount)\).

To see the trade-off concretely, consider a vocabulary embedding matrix with 100,000 rows and 10,000 columns (1B parameters):

• Dense (FP32): 1B parameters \(\times\) 4 bytes each = 4 GB.
• Sparse (1 percent density): CSR stores roughly 10M entries \(\times\) (4-byte value + 4-byte column index), plus one row pointer per row and a final pointer, totaling ≈ 80.4 MB.
• Result: A 50× reduction in memory footprint, fitting a model that would otherwise trigger an out-of-memory (OOM) error.

The saving holds only at low density. Each stored nonzero carries its index, so at high density CSR occupies more memory than the dense tensor it replaces. Storage and execution break-even finds that break-even density for pruned weights, and N:M structured sparsity mechanics shows the fixed pattern of two nonzeros in every four that an accelerator datapath decodes without a general index.

Operation cost reference

Each common layer type reduces to the GEMMs and elementwise operations above, so its parameter count and forward-pass FLOPs follow from its shapes (table 2). Given layer dimensions and input shapes, the formulas give first-order estimates of model size and compute before hardware is provisioned. The attention row is the one to watch. Sequence length contributes an \(\SequenceLength^2\) term, while hidden dimension contributes \(\HiddenDimension^2\) projection work, so long-context models can shift the dominant cost without changing the parameter count.

Table 2: Deep Learning Tensor Primitives: Summary of shapes, parameters, and FLOP counts. Note: \(\BatchSize\) is batch size, \(\SequenceLength\) is sequence length, and \(\KernelSize\) is kernel size. The attention FLOPs include QKV projections and score interactions. The \(\SequenceLength^2\) term applies to full-sequence training and prefill, while KV-cached generation is linear per token in cached context; \(\HiddenDimension^2\) projection work dominates at large hidden dimensions.

Layer type	Output shape	Parameters (\(\Parameters\))	FLOPs (per forward pass)
Linear	\((\BatchSize, \LayerOutputWidth)\)	\(\LayerInputWidth \times \LayerOutputWidth + \LayerOutputWidth\) if bias is enabled	\(2 \times \BatchSize\times \LayerInputWidth \times \LayerOutputWidth\)
Conv2D	\((\BatchSize, \LayerOutputWidth, \OutputFeatureMapHeight, \OutputFeatureMapWidth)\)	\(\KernelSize^2 \times \LayerInputWidth \times \LayerOutputWidth + \LayerOutputWidth\) if bias is enabled	\(2 \times \BatchSize\times \OutputFeatureMapHeight \times \OutputFeatureMapWidth \times \KernelSize^2 \times \LayerInputWidth \times \LayerOutputWidth\)
Multi-head self-attention	\((\BatchSize, \SequenceLength, \HiddenDimension)\)	\(4 \times \HiddenDimension^2\), plus \(4 \times \HiddenDimension\) if projection biases are enabled	\(\BatchSize\times (4 \SequenceLength^2 \HiddenDimension+ 8 \SequenceLength \HiddenDimension^2)\)
LayerNorm	\((\BatchSize, \SequenceLength, \HiddenDimension)\)	\(2 \times \HiddenDimension\) if affine scale and bias are enabled	\(\mathcal{O}(\BatchSize\times \SequenceLength\times \HiddenDimension)\)

The formulas in table 2 specify what each layer computes and how many operations it performs. Expressing those computations in code raises a second question: how frameworks represent multidimensional arrays in flat memory without copying data unnecessarily.

Tensor Programming Primitives

A shape mismatch crash, a silently wrong broadcast, and a kernel running at 5 percent of peak because of a noncontiguous tensor are common ML engineering failures, and all three trace back to the same layer of abstraction. Tensor programming translates the abstract math of linear algebra into concrete array manipulations that run on hardware.

Shapes, strides, and memory layout

Frameworks represent a tensor as a view over a flat storage buffer, described by shape, stride, dtype, and offset metadata rather than as a multidimensional array of pointers. Its shape specifies the logical dimensions of the tensor (for example, (3, 4)), while its stride defines how many elements to skip in storage to advance one step along each dimension, which means only contiguous tensors store their logical elements in a single adjacent block.

Operations like transpose() or view() often just change the strides, not the data in memory. This is fast (\(\mathcal{O}(1)\)) but can lead to noncontiguous tensors that fail in kernels requiring contiguous data. In such cases, calling contiguous() on a tensor of \(\TensorElementCount\) elements forces an \(\mathcal{O}(\TensorElementCount)\) memory copy that can dominate runtime if triggered repeatedly inside a loop.

The order of the strides is the tensor’s memory layout. Row-major and column-major layouts of the same matrix differ only in which dimension has stride 1, and that choice directly affects cache efficiency: iterating along the stride-1 dimension touches contiguous memory locations, which is dramatically faster than strided access.

A common optimization pattern is to materialize a transposed tensor once before repeated operations to ensure contiguous access in the hot loop. The one-time copy cost is amortized across many subsequent operations. Rank, layout, and data type introduces strides and layouts in the framework that manages them, and Memory-efficient tensor layouts shows how the layout choice (channels-last versus channels-first) changes accelerator throughput for convolutions.

Broadcasting

Zero-copy stride manipulation also powers broadcasting, the rule that lets arithmetic operators combine tensors of different shapes without materializing duplicate copies of the smaller input. Frameworks compare dimensions from right to left and treat two dimensions as compatible when they are equal or when one of them is 1; a singleton dimension is then stretched virtually by setting its stride to zero, though an unintended singleton alignment, such as adding a \((\BatchSize, 1)\) column vector of predictions to a \((\BatchSize,)\) row vector of targets, silently allocates a \((\BatchSize, \BatchSize)\) output matrix that can exhaust memory or corrupt a loss calculation.

Adding a \((3, 1)\) column tensor (\(\BroadcastColumnTensor\)) to a \((1, 4)\) row tensor (\(\BroadcastRowTensor\)) expands both virtually across their complementary singleton dimensions and generates a \((3, 4)\) matrix of pairwise sums \(\BroadcastColumnElement_i + \BroadcastRowElement_j\) (figure 1). Dashed cell boundaries mark result elements whose operands are read through stride-0 views, so neither input is copied; only the \((3, 4)\) output is allocated, and hardware repeatedly accesses identical input buffer addresses.

Figure 1: Tensor Broadcasting Rules: A (3,1) column and a (1,4) row broadcast across complementary dimensions, producing the pairwise sums in a (3,4) result.

Consider a concrete case: tensor A has shape (32, 1, 64) and tensor B has shape (1, 128, 64). Comparing dimensions right to left, 64 matches 64, then 1 stretches to 128, then 1 stretches to 32, yielding result shape (32, 128, 64). Writing out the expanded shape before running catches silent logic bugs that accidentally allocate a large tensor (for example, a (Batch, Batch) matrix instead of an element-wise (Batch) vector).

Shapes, strides, and broadcasting govern how tensors flow through a model’s forward pass. Training adds a second requirement, computing how much each parameter contributed to the error. Backpropagation meets that requirement and imposes the memory costs quantified in the training-memory decomposition below.

Mechanics of Learning

Training a neural network takes roughly three times the compute of a forward pass and holds gigabytes of intermediate activations in memory until the backward sweep consumes them. This section traces how the chain rule and reverse-mode automatic differentiation create that compute and memory footprint, how compilers fuse the resulting computational graph, and how the log-sum-exp shift prevents the softmax at the loss from overflowing in reduced precision.

Systems Perspective 1.2: Why this matters

When training fails (loss goes to NaN, gradients explode, or memory runs out), understanding what backpropagation actually does is essential for diagnosing the problem. The backward graph supplies the mental model for reasoning about gradient flow and memory usage during training.

The chain rule and automatic differentiation

For a composed function \(\ComposedFunctionOutput= \LayerFunction(\InnerLayerFunction(\ComposedFunctionInput))\), the derivative is \(\frac{d\ComposedFunctionOutput}{d\ComposedFunctionInput} = \frac{d\ComposedFunctionOutput}{d\InnerLayerFunction} \cdot \frac{d\InnerLayerFunction}{d\ComposedFunctionInput}\). In a neural network, \(\LayerFunction\) and \(\InnerLayerFunction\) are layers, and the composition can be many levels deep. For a three-layer network \(\ComposedFunctionOutput= \LayerFunction_3(\LayerFunction_2(\LayerFunction_1(\ComposedFunctionInput)))\), the chain rule extends to: \[ \frac{\partial \ComposedFunctionOutput}{\partial \ComposedFunctionInput} = \frac{\partial \LayerFunction_3}{\partial \LayerFunction_2} \cdot \frac{\partial \LayerFunction_2}{\partial \LayerFunction_1} \cdot \frac{\partial \LayerFunction_1}{\partial \ComposedFunctionInput} \]

Each factor in this product is a local derivative, computed at one layer using only that layer’s inputs and outputs. This locality keeps the bookkeeping cheap: the entire network never needs to be differentiated as a monolithic function. Instead, each layer computes its own local derivative during the backward pass and multiplies it by the gradient flowing in from the layer above.

Modern frameworks use reverse-mode automatic differentiation, which computes gradients for all \(\Parameters\) parameters in a single backward pass. Starting from the output and working backward (reverse mode) requires one pass regardless of the number of parameters, whereas starting from each input and working forward (forward mode) would require \(\Parameters\) passes, one per parameter. This is why a training step is a small constant multiple of inference, commonly about 2–3\(\times\) a forward pass for dense networks, rather than \(\Parameters\) passes. Forward and reverse mode shows how frameworks implement both modes and why training uses reverse mode.

The backpropagation algorithm

Backpropagation³ implements the chain rule efficiently through two passes: forward to compute outputs, backward to compute gradients. The gradient edges retrace forward dependencies, meaning each backward step requires intermediate activations retained from the corresponding forward operation. Figure 2 traces these dual dataflow passes across a two-layer network: solid gray arrows stage forward activations (\(\InputFeatureVector\to \HiddenActivation\to \PredictedOutput\)) through weight parameters \(\WeightMatrix_1\) and \(\WeightMatrix_2\) to evaluate loss \(\LossFunction\), while curved dashed red arrows propagate intermediate adjoints (\(\frac{\partial \LossFunction}{\partial \PredictedOutput}, \frac{\partial \LossFunction}{\partial \HiddenActivation}, \frac{\partial \LossFunction}{\partial \InputFeatureVector}\)) in reverse order.

³ Backpropagation: Short for “backward propagation of errors.” The algorithm was independently discovered multiple times, by Werbos (1974) and Linnainmaa (1970) for reverse-mode differentiation and by Rumelhart et al. (1986) for neural network training. Its key insight is that computing gradients for all parameters requires one backward pass through the graph rather than one pass per parameter. Section 1.4.1 derives the constant-factor work bound and a representative dense-layer cost.

Werbos, Paul. 1974. “Beyond Regression: New Tools for Prediction and Analysis in the Behavioral Sciences.” PhD thesis, Harvard University.

Linnainmaa, Seppo. 1970. “The Representation of the Cumulative Rounding Error of an Algorithm as a Taylor Expansion of the Local Rounding Errors.” Master's thesis, University of Helsinki.

Rumelhart, David E., Geoffrey E. Hinton, and Ronald J. Williams. 1986. “Learning Representations by Back-Propagating Errors.” Nature 323 (6088): 533–36. https://doi.org/10.1038/323533a0.

Figure 2: Backpropagation Computational Graph: Solid gray arrows carry the forward computation from input \(\InputFeatureVector\) to loss \(\LossFunction\), while dashed red arrows carry gradients backward along the same dependencies.

Forward pass

Using row-vector activations, omitting biases, and retaining intermediate activations for backward execution, computation begins at input \(\InputFeatureVector\). Multiplying \(\InputFeatureVector\) by weight matrix \(\WeightMatrix_1\) yields hidden activation \(\HiddenActivation\), which the execution runtime caches for backward differentiation. Multiplying \(\HiddenActivation\) by \(\WeightMatrix_2\) produces network output \(\PredictedOutput\), which is similarly cached. Finally, comparing \(\PredictedOutput\) to the target label produces scalar loss \(\LossFunction\).

At this point, forward evaluation is complete, and memory holds input \(\InputFeatureVector\), cached activation \(\HiddenActivation\), cached prediction \(\PredictedOutput\), and loss \(\LossFunction\). For large models, these retained activation tensors frequently dominate overall memory consumption.

Backward pass

Gradient propagation begins backward from scalar loss \(\LossFunction\). The loss function evaluates \(\frac{\partial \LossFunction}{\partial \PredictedOutput}\), the error gradient with respect to the network prediction, which starts the backward pass.

Because \(\PredictedOutput = \HiddenActivation\cdot \WeightMatrix_2\), the matrix chain rule yields two gradients at this layer: \[ \frac{\partial \LossFunction}{\partial \WeightMatrix_2} = \HiddenActivation^T \cdot \frac{\partial \LossFunction}{\partial \PredictedOutput} \qquad \text{(weight gradient—used to update } \WeightMatrix_2\text{)} \] \[ \frac{\partial \LossFunction}{\partial \HiddenActivation} = \frac{\partial \LossFunction}{\partial \PredictedOutput} \cdot \WeightMatrix_2^T \qquad \text{(input gradient—passed backward to the previous layer)} \]

Computing \(\frac{\partial \LossFunction}{\partial \WeightMatrix_2}\) requires cached activation \(\HiddenActivation\) from the forward pass. Activations must remain resident in memory because every layer’s weight gradient depends directly on that layer’s input.

Continuing backward to the first layer, the linear relationship \(\HiddenActivation= \InputFeatureVector\cdot \WeightMatrix_1\) produces: \[ \frac{\partial \LossFunction}{\partial \WeightMatrix_1} = \InputFeatureVector^T \cdot \frac{\partial \LossFunction}{\partial \HiddenActivation} \]

Each backward step requires two operands: the gradient flowing in from downstream layers \(\left(\frac{\partial \LossFunction}{\partial \HiddenActivation}\right)\) and that layer’s cached input (\(\InputFeatureVector\)). Consequently, the backward pass requires roughly 2\(\times\) the forward pass FLOPs: each layer executes two matrix multiplications (one for the parameter gradient, one for the upstream activation gradient) rather than one. Gradient computation and backpropagation teaches the algorithm step by step on a running network. The cached inputs that every backward step needs are the reason training memory exceeds model size, which the next section quantifies.

Training memory

Because the backward pass holds those intermediate activations alongside weights, gradients, and optimizer buffers, training memory is the sum of four components rather than model size alone: \[ \TotalMemoryFootprint = \WeightMemory + \GradientMemory + \OptimizerStateMemory + \ActivationMemory \]

For a standard Adaptive Moment Estimation (Adam) optimizer (Kingma and Ba 2015) in mixed precision (Micikevicius et al. 2018):

Kingma, Diederik P., and Jimmy Ba. 2015. “Adam: A Method for Stochastic Optimization.” 3rd International Conference on Learning Representations (ICLR).

Micikevicius, Paulius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, et al. 2018. “Mixed Precision Training.” International Conference on Learning Representations (ICLR).

• Weights: 2 bytes (FP16/BF16) or 4 bytes (FP32).
• Gradients: Same size as weights.
• Optimizer State: Adam moments require 8 bytes per parameter; an FP32 master copy adds another 4 bytes, giving 8–12 bytes per parameter.
• Activations: Activations typically dominate the dynamic memory footprint during training. A recomputation- or tiled-attention-friendly estimate scales as \(\mathcal{O}(\BatchSize\times \SequenceLength\times \LayerCount \times \HiddenDimension)\); a full no-recompute attention implementation also materializes attention-score tensors that scale with \(\SequenceLength^2\).

To see how these four components interact on a single accelerator, consider a concrete language model. Once the static model state is allocated, the remaining device capacity sets a hard ceiling on the batch size that can run without recomputation.

Napkin Math 1.1: Worked example: GPT-2 (1.5B) training memory

The model: GPT-2 XL has \(\Parameters=\) \(1.5 \times 10^{9}\) parameters, 48 layers, hidden dimension \(\HiddenDimension=\) 1600.

Model state (fixed per step):

• Weights (BF16): \(1.5 \times 10^{9}\) \(\times\) 2 bytes = 3 GB
• Gradients (BF16): \(1.5 \times 10^{9}\) \(\times\) 2 bytes = 3 GB
• Optimizer (FP32 master + momentum + variance): \(1.5 \times 10^{9}\) \(\times\) 12 bytes = 18 GB
• Total model state: 24 GB, which fits on an 80 GB-class H100 (85.9 GB in decimal units) but leaves only 61.9 GB for activations.

Activations (scale with batch):

Per-layer retained activations for this quick estimate are approximately \(12 \times \BatchSize\times \SequenceLength\times \HiddenDimension\) BF16 elements, or \(12 \times \BatchSize\times \SequenceLength\times \HiddenDimension\times 2\) bytes, where \(\BatchSize\) is batch size and \(\SequenceLength\) is sequence length. The factor twelve accounts for the major intermediate tensors retained for backpropagation: input activations, QKV projections (\(3\HiddenDimension\)), attention output, FFN intermediate (\(4\HiddenDimension\)), and layer norm/dropout masks. This estimate assumes the implementation does not retain the full \(\BatchSize\times \AttentionHeadCount \times \SequenceLength^2\) attention-score tensor, as in recomputation- or tiled-attention-friendly accounting. With 48 layers, batch size 8, and sequence length 1024:

\[48 \times 12 \times 8 \times 1024 \times 1600 \times 2\text{ bytes } \approx 15.1\text{ GB}\]

Systems insight: Total memory reaches ~39.1 GB under this activation-accounting convention, which fits on one 85.9 GB accelerator. Each additional sample adds ~1.9 GB of activations, so under this quick estimate the remaining capacity fills near batch 32; at batch 64, activations grow to ~120.8 GB, well beyond it. The estimate omits smaller retained tensors such as dropout masks and softmax inputs. The full per-layer count of 34 bytes per token per hidden unit (Korthikanti et al. 2023) fills the same room near batch 23, the limit GPT-2 optimization walkthrough works from. That crossover is the threshold where gradient checkpointing or another memory-saving method becomes necessary; its savings and recomputation overhead depend on checkpoint placement and implementation.

Korthikanti, Vijay Anand, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. 2023. “Reducing Activation Recomputation in Large Transformer Models.” Proceedings of Machine Learning and Systems (MLSys) 5: 341–53.

Activation explosion

While weights are fixed (\(\mathcal{O}(\Parameters)\)), activations grow linearly with batch size and at least linearly with sequence length; implementations that materialize attention scores add an \(\mathcal{O}(\SequenceLength^2)\) component. As the worked example above shows, activation memory can quickly exceed the room left after model state. Gradient checkpointing⁴ reduces this pressure by storing only selected activations and recomputing the rest during backpropagation (Chen et al. 2016), and FlashAttention removes the quadratic term by tiling attention (Dao et al. 2022). The same decomposition provides a quick feasibility check: sum the per-parameter model state, add the activation term implied by batch size and sequence length, and compare the total against the accelerator’s capacity. Optimizers and their memory applies the model-state terms to optimizer choice, and Backpropagation and activation memory develops the activation term, with Activation checkpointing teaching the recomputation trade-off for the training loop.

⁴ Gradient checkpointing: Trades compute for memory. Instead of storing all activations, only a subset (checkpoints) is kept, and the missing ones are recomputed during the backward pass. This reduces memory usage from storing every layer activation to roughly the square root of the layer count at the cost of ~33 percent more compute (Chen et al. 2016).

Chen, Tianqi, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. 2016. “Training Deep Nets with Sublinear Memory Cost.” arXiv Preprint arXiv:1604.06174.

Computational graphs and optimization

The dependency structure exposed by backpropagation also gives compilers something to optimize. ML compilers represent models as directed acyclic graphs (DAGs), and that representation enables hardware-independent transformations.

Static single assignment

To transform a DAG safely, a compiler must know which write each read depends on. Static single-assignment (SSA) form makes those dependencies explicit by assigning each variable exactly once rather than overwriting tensor variables in place, enabling safe transformations such as operator fusion.

Operator fusion

Without fusion, each operation in a chain like MatMul → Add (bias) → ReLU produces an intermediate tensor that is written to high-bandwidth memory (HBM) and then read back for the next operation. For elementwise operations like Add and ReLU, the compute is trivial (one FLOP per element) but the memory traffic is not (read the tensor, write it back). The arithmetic intensity of unfused elementwise operations is therefore close to zero, which makes them strongly memory bound.

Fusion combines consecutive operations into a single kernel that reads the input once, applies all operations in registers or shared memory, and writes the final result once. For a sequence of \(\ChainedOperationCount\) elementwise operations on a tensor of \(\TensorElementCount\) elements of \(\ElementStorageSize\) bytes each, fusion reduces memory traffic from \(2\ChainedOperationCount \TensorElementCount\ElementStorageSize\) bytes (each operation reads and writes) to \(2\TensorElementCount\ElementStorageSize\) bytes (one read, one write), a \(\ChainedOperationCount\times\) reduction. The memory wall and kernel fusion develops fusion as a framework decision, and Kernel fusion on the accelerator shows the accelerator-side limits on how much a single kernel can fuse.

FlashAttention is fusion applied to attention: it computes the score, softmax, and value-product stages in SRAM tiles without materializing the full attention matrix in HBM, reducing auxiliary attention memory from \(\mathcal{O}(\SequenceLength^2)\) to \(\mathcal{O}(\SequenceLength)\) for a fixed head dimension (Dao et al. 2022). Section 1.4.2 derives the traffic bound step by step, and FlashAttention: IO-aware attention optimization applies the algorithm to training.

Fused and tiled kernels keep results exact while moving fewer bytes. The softmax at the end of the forward pass raises a different question, whether its results stay finite at all.

Logits and numerical stability

A classifier’s forward pass ends in a softmax over the final layer’s outputs, and the cross-entropy loss is computed from that softmax. The softmax exponentiates those outputs, and the exponential of a large output exceeds the range of every reduced-precision format, so the softmax is a common overflow site in reduced-precision training.

Neural-network classifiers commonly output a vector of logits⁵ \(\LogitVector\) (unnormalized scores). For class index \(\ClassIndex\), softmax normalizes \(\OutputLogit_{\ClassIndex}\) over all class indices \(\NormalizingClassIndex\) to produce a probability:

⁵ Logit: Joseph Berkson coined the term in 1944 by analogy with probit. The logit function is the inverse of the logistic (sigmoid) function: \(\text{logit}(\LogisticProbability) = \log(\LogisticProbability/(1-\LogisticProbability))\). In deep learning, “logits” refers more loosely to the raw, unnormalized output of the final linear layer before any activation function is applied.

\[ \text{Softmax}(\OutputLogit_{\ClassIndex}) = \frac{e^{\OutputLogit_{\ClassIndex}}}{\sum_{\NormalizingClassIndex} e^{\OutputLogit_{\NormalizingClassIndex}}} \]

The problem is that if \(\OutputLogit_{\ClassIndex}\) is large (for example, 100), the exponential \(e^{\OutputLogit_{\ClassIndex}}\) overflows common training and inference formats such as FP32, BF16, and FP16. FP64 can represent this specific value, but production ML kernels rarely use FP64 for softmax. The fix is to subtract the largest logit before exponentiating, the log-sum-exp trick. The shift leaves the result mathematically unchanged and keeps every exponent at or below zero, as the worked example below illustrates for realistic logit values.⁶

⁶ Log-sum-exp: Implemented as torch.logsumexp in PyTorch and scipy.special.logsumexp in SciPy. It relies on the identity \(\log\left(\sum e^{\OutputLogit_{\NormalizingClassIndex}}\right) = \SoftmaxRowMaximum + \log\left(\sum e^{\OutputLogit_{\NormalizingClassIndex} - \SoftmaxRowMaximum}\right)\), where \(\SoftmaxRowMaximum = \max(\OutputLogit_{\NormalizingClassIndex})\). Shifted values \(\OutputLogit_{\NormalizingClassIndex} - \SoftmaxRowMaximum\) are \(\le 0\), ensuring exponentials never overflow.

Napkin Math 1.2: Log-sum-exp in action

Setup: A three-class classifier outputs logits \(\LogitVector=\) \([100, 101, 102]\).

Without the trick (naive softmax): \[\text{ $\exp(100) \approx 2.7 \times 10^{43}$, $\exp(101) \approx 7.3 \times 10^{43}$, $\exp(102) \approx 2.0 \times 10^{44}$}\]

These numbers are representable in FP64 but overflow FP32 (max \(\approx 3.4 \times 10^{38}\)). With FP16 (max \(\approx 65{,}504\)), even \(e^{12}\) overflows. Raw logits rarely reach magnitude 100, but large intermediate values do arise (for example, unscaled attention scores or overconfident outputs), and because FP16 and BF16 are the standard training precisions, a naive softmax risks overflow often enough that the stable form is used by default.

With the trick: Subtract \(\SoftmaxRowMaximum = \max(\LogitVector) =\) 102:

Shifted logits: \(\LogitVector- \SoftmaxRowMaximum =\) \([-2, -1, 0]\)

Exponentials: \(\exp(-2) \approx 0.135\), \(\exp(-1) \approx 0.368\), \(\exp(0) = 1.0\)

Sum = 1.503. LogSumExp = 102 \(+ \log(1.503) = 102.408\).

Softmax: \([0.135/1.503,\; 0.368/1.503,\; 1.0/1.503] = [0.090,\; 0.245,\; 0.665]\)

Systems insight: The exponentiated shifted values and the resulting softmax probabilities are in \([0, 1]\), so nothing overflows, even in FP16.

Cross-entropy and classification loss functions builds the cross-entropy loss on this stable softmax. The same maximum subtraction reappears, generalized to tiles of a long sequence, in the tiled-attention derivation of section 1.4.2, one of the three proofs that close this appendix.

Derivations

Without a constant-factor bound on reverse-mode differentiation, training a 175-billion-parameter model would require 175 billion forward passes per step; without an online softmax recurrence, attention would have to materialize an \(\SequenceLength \times \SequenceLength\) matrix in HBM; and without a constrained-optimum derivation, a fixed compute budget could be wasted on an oversized model starved of tokens. Each subsection below states one of these three results, proves it, and closes on the systems consequence that the chapters rely on.

Cost of reverse-mode differentiation

Evaluating the full gradient vector \(\nabla \ScalarGraphFunction(\FunctionInputVector)\) of a scalar loss function \(\ScalarGraphFunction: \mathbb{R}^\FunctionInputDimension\to \mathbb{R}\) requires compute proportional to evaluating \(\ScalarGraphFunction(\FunctionInputVector)\) itself, with a constant-factor bound independent of the input dimension \(\FunctionInputDimension\). Baur and Strassen proved this complexity result in 1983. Computing a function together with its full gradient takes at most four times the operations of the function alone when every operation counts, and at most three times when only multiplications and divisions count (Baur and Strassen 1983). The per-operation accounting below is simpler than their construction and proves the looser constant of five.

Baur, Walter, and Volker Strassen. 1983. “The Complexity of Partial Derivatives.” Theoretical Computer Science 22 (3): 317–30. https://doi.org/10.1016/0304-3975(83)90110-x.

Statement

Let \(\ScalarGraphFunction: \mathbb{R}^\FunctionInputDimension\to \mathbb{R}\) be a multivariate scalar-valued function represented as a computational graph \(\ComputationalGraph= (\GraphNodeSet, \GraphEdgeSet)\) consisting of \(\ElementaryOperationCount(\ScalarGraphFunction)\) elementary operations (additions, subtractions, multiplications, divisions, and unary functions such as \(\exp, \ln, \sin, \sqrt{\cdot}\)).

Theorem 1.1: Cheap-gradient bound

If a scalar function \(\ScalarGraphFunction: \mathbb{R}^\FunctionInputDimension\to \mathbb{R}\) can be computed in \(\ElementaryOperationCount(\ScalarGraphFunction)\) elementary operations, then the full gradient vector \(\nabla \ScalarGraphFunction(\FunctionInputVector) = \left(\frac{\partial \ScalarGraphFunction}{\partial \FunctionInputElement_1}, \frac{\partial \ScalarGraphFunction}{\partial \FunctionInputElement_2}, \dots, \frac{\partial \ScalarGraphFunction}{\partial \LastFunctionInputComponent}\right)^T \in \mathbb{R}^\FunctionInputDimension\) can be evaluated in total work \(\text{Work}_{\text{Reverse}}(\ScalarGraphFunction, \nabla \ScalarGraphFunction)\) satisfying: \[ \text{Work}_{\text{Reverse}}(\ScalarGraphFunction, \nabla \ScalarGraphFunction) \le \ReverseModeWorkConstant\cdot \ElementaryOperationCount(\ScalarGraphFunction) \] where \(\ReverseModeWorkConstant\le 5\) is a universal constant independent of the input dimension \(\FunctionInputDimension\). Furthermore, the backward pass alone requires work \(\text{Work}_{\text{Backward}}(\nabla \ScalarGraphFunction) \le 4 \cdot \ElementaryOperationCount(\ScalarGraphFunction)\).

Proof

Let the computational graph \(\ComputationalGraph= (\GraphNodeSet, \GraphEdgeSet)\) be topologically ordered with nodes \(\FirstInputNodePrimal, \dots, \TraceIntermediateVariable_0, \TraceIntermediateVariable_1\), \(\TraceIntermediateVariable_2, \dots, \OutputNodePrimal\), where input variables are \(\FunctionInputElement_{\InputVariableIndex} = \InputNodePrimal\) for \(\InputVariableIndex\in \{1, \dots, \FunctionInputDimension\}\) and the final output is \(\ScalarGraphFunction(\FunctionInputVector) = \OutputNodePrimal\). Each intermediate node \(\TraceIntermediateVariable_j\) is evaluated via a primitive operation: \[ \TraceIntermediateVariable_j = \PrimitiveOperation_j\left( \{ \TraceIntermediateVariable_k \mid \ParentNodeIndex\in \text{Parents}(\CurrentNodeIndex) \} \right) \]

Define the adjoint variable \(\TraceIntermediateAdjoint_i\) for every node \(\TraceIntermediateVariable_i\) as the partial derivative of the scalar output \(\OutputNodePrimal\) with respect to \(\TraceIntermediateVariable_i\): \[ \TraceIntermediateAdjoint_i \triangleq \frac{\partial \ScalarGraphFunction}{\partial \TraceIntermediateVariable_i} = \frac{\partial \OutputNodePrimal}{\partial \TraceIntermediateVariable_i} \]

By the multivariate chain rule, the adjoint \(\TraceIntermediateAdjoint_i\) is the sum of partial derivatives propagated from all successor nodes (children) that consume \(\TraceIntermediateVariable_i\): \[ \TraceIntermediateAdjoint_i = \sum_{j \in \text{Children}(i)} \TraceIntermediateAdjoint_j \cdot \frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_i} \]

Reverse-mode automatic differentiation initializes \(\OutputNodeAdjoint = 1\) and \(\TraceIntermediateAdjoint_i = 0\) for all \(\GraphNodeIndex< \ElementaryOperationCount\), then traverses the nodes in reverse topological order from \(\OutputNodePrimal\) down to \(\TraceIntermediateVariable_1\). For each primitive node \(\TraceIntermediateVariable_j\), its fully accumulated adjoint \(\TraceIntermediateAdjoint_j\) propagates error signals to its parent nodes \(\TraceIntermediateVariable_k\) (\(\ParentNodeIndex\in \text{Parents}(\CurrentNodeIndex)\)) via: \[ \TraceIntermediateAdjoint_k \gets \TraceIntermediateAdjoint_k + \TraceIntermediateAdjoint_j \cdot \frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_k} \]

We bound the backward cost one operation type at a time:

1. Addition/Subtraction (\(\TraceIntermediateVariable_j = \TraceIntermediateVariable_i \pm \TraceIntermediateVariable_k\)):
   • Local derivatives: \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_i} = 1\), \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_k} = \pm 1\).
   • Backward updates: \(\TraceIntermediateAdjoint_i \gets \TraceIntermediateAdjoint_i + \TraceIntermediateAdjoint_j\) and \(\TraceIntermediateAdjoint_k \gets \TraceIntermediateAdjoint_k \pm \TraceIntermediateAdjoint_j\).
   • Cost: 2 additions in backward pass for 1 forward operation. Work ratio: \(\text{Fwd}(1) + \text{Bwd}(2) = 3 \text{ OPs} \le 5 \cdot 1\).
2. Multiplication (\(\TraceIntermediateVariable_j = \TraceIntermediateVariable_i \cdot \TraceIntermediateVariable_k\)):
   • Local derivatives: \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_i} = \TraceIntermediateVariable_k\), \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_k} = \TraceIntermediateVariable_i\).
   • Backward updates: \(\TraceIntermediateAdjoint_i \gets \TraceIntermediateAdjoint_i + \TraceIntermediateAdjoint_j \cdot \TraceIntermediateVariable_k\) and \(\TraceIntermediateAdjoint_k \gets \TraceIntermediateAdjoint_k + \TraceIntermediateAdjoint_j \cdot \TraceIntermediateVariable_i\).
   • Cost: 2 multiplications + 2 additions = 4 operations in backward pass for 1 forward operation. Work ratio: \(\text{Fwd}(1) + \text{Bwd}(4) = 5 \text{ OPs} \le 5 \cdot 1\).
3. Division (\(\TraceIntermediateVariable_j = \TraceIntermediateVariable_i/\TraceIntermediateVariable_k\)):
   • Local derivatives: \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_i} = \frac{1}{\TraceIntermediateVariable_k}\), \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_k} = -\frac{\TraceIntermediateVariable_i}{\TraceIntermediateVariable_k^2} = -\frac{\TraceIntermediateVariable_j}{\TraceIntermediateVariable_k}\).
   • Backward updates: \(\TraceIntermediateAdjoint_i \gets \TraceIntermediateAdjoint_i + \frac{\TraceIntermediateAdjoint_j}{\TraceIntermediateVariable_k}\) and \(\TraceIntermediateAdjoint_k \gets \TraceIntermediateAdjoint_k - \frac{\TraceIntermediateAdjoint_j \cdot \TraceIntermediateVariable_j}{\TraceIntermediateVariable_k}\).
   • Cost: Compute the quotient of the incoming adjoint and the denominator once, then reuse it in both updates: add it to the numerator adjoint, and multiply it by the cached forward output before subtracting from the denominator adjoint. This requires 2 multiplications/divisions + 2 additions = 4 operations in backward pass for 1 forward operation. Work ratio: \(\text{Fwd}(1) + \text{Bwd}(4) = 5 \text{ OPs} \le 5 \cdot 1\).
4. Unary transcendental operations (\(\TraceIntermediateVariable_j = \PrimitiveOperation(\TraceIntermediateVariable_i)\)):
   • Local derivative: \(\frac{\partial \TraceIntermediateVariable_j}{\partial \TraceIntermediateVariable_i} = \PrimitiveOperation'(\TraceIntermediateVariable_i)\).
   • Backward update: \(\TraceIntermediateAdjoint_i \gets \TraceIntermediateAdjoint_i + \TraceIntermediateAdjoint_j \cdot \PrimitiveOperation'(\TraceIntermediateVariable_i)\).
   • For \(\TraceIntermediateVariable_j = \exp(\TraceIntermediateVariable_i)\), \(\PrimitiveOperation'(\TraceIntermediateVariable_i) = \TraceIntermediateVariable_j\), requiring 1 multiplication + 1 addition = 2 OPs.
   • For \(\TraceIntermediateVariable_j = \ln(\TraceIntermediateVariable_i)\), \(\PrimitiveOperation'(\TraceIntermediateVariable_i) = 1/\TraceIntermediateVariable_i\), requiring 1 division + 1 addition = 2 OPs.

Summing across all primitive nodes in \(\ComputationalGraph\), every forward operation generates at most 4 backward operations: \[ \text{Work}_{\text{Backward}}(\nabla \ScalarGraphFunction) \le 4 \cdot \ElementaryOperationCount(\ScalarGraphFunction) \]

Adding the forward pass work \(\ElementaryOperationCount(\ScalarGraphFunction)\) yields the total work bound: \[ \text{Work}_{\text{Reverse}}(\ScalarGraphFunction, \nabla \ScalarGraphFunction) = \ElementaryOperationCount(\ScalarGraphFunction) + \text{Work}_{\text{Backward}}(\nabla \ScalarGraphFunction) \le 5 \cdot \ElementaryOperationCount(\ScalarGraphFunction) \qquad\blacksquare\]

Systems Perspective 1.3: Why 175B-parameter large language models can be trained

The independence of \(\ReverseModeWorkConstant\le 5\) from input dimension \(\FunctionInputDimension\) is the mathematical reason deep learning works. Consider training a transformer model with \(\Parameters= 175 \times 10^9\) parameters:

• Forward pass: Matrix multiplications \(\LayerOutputMatrix= \LayerInputMatrix\WeightMatrix\) require \(2\Parameters\) FLOPs per token.
• Reverse-mode backward pass: Computes both weight gradients \(\frac{\partial \LossFunction}{\partial \WeightMatrix} = \LayerInputMatrix^T \frac{\partial \LossFunction}{\partial \LayerOutputMatrix}\) (\(2\Parameters\) FLOPs) and input activation gradients \(\frac{\partial \LossFunction}{\partial \LayerInputMatrix} = \frac{\partial \LossFunction}{\partial \LayerOutputMatrix} \WeightMatrix^T\) (\(2\Parameters\) FLOPs), totaling \(4\Parameters\) FLOPs per token.
• Total training work: \(2\Parameters(\text{fwd}) + 4\Parameters(\text{bwd}) = 6\Parameters\) FLOPs per token. The backward compute factor is exactly \(\frac{4\Parameters}{2\Parameters} = 2\), so forward plus backward costs three times the forward pass, within both the bound of \(\ReverseModeWorkConstant\le 5\) proved here and Baur and Strassen’s sharper constant of four.

If we instead used forward-mode automatic differentiation (AD) or finite differences to compute \(\nabla \ScalarGraphFunction\in \mathbb{R}^\Parameters\), evaluating each partial derivative separately would require \(\Parameters\) forward passes: \(\text{Work}_{\text{ForwardMode}}(\nabla \ScalarGraphFunction) = \mathcal{O}(\Parameters\cdot \ElementaryOperationCount) = \mathcal{O}(\Parameters^2)\) FLOPs. For \(\Parameters= 175 \times 10^9\), that method would require \(2 \times (175 \times 10^9)^2 \approx 6.1 \times 10^{22}\) FLOPs per token, \(\Parameters\) times the forward-pass work and therefore prohibitive at this scale. Reverse-mode AD replaces this \(\mathcal{O}(\Parameters)\) factor with a constant multiplier of 2 for the dense-layer accounting above.

Reverse mode achieves this constant factor by keeping the forward pass’s intermediate values for the backward sweep, which is why the training-memory decomposition in section 1.3.3 carries an activation term. Attention adds an intermediate whose size grows with the square of the sequence length, the subject of the next derivation.

Tiled attention and online softmax

Standard self-attention in transformers requires materializing an \(\SequenceLength\times \SequenceLength\) attention score matrix in HBM, creating an \(\mathcal{O}(\SequenceLength^2)\) memory bandwidth bottleneck. FlashAttention (Dao et al. 2022) reduces this bottleneck by tiling query, key, and value matrices and recomputing attention using an online softmax recurrence without materializing intermediate scores to HBM. FlashAttention: IO-aware attention optimization applies the algorithm to training; this subsection proves the online softmax recurrence it relies on and then its memory-traffic bound.

Memory traffic of standard attention

Given queries \(\AttentionQueryMatrix\in \mathbb{R}^{\SequenceLength\times \AttentionHeadDimension}\), keys \(\AttentionKeyMatrix\in \mathbb{R}^{\SequenceLength\times \AttentionHeadDimension}\), and values \(\AttentionValuesMatrix\in \mathbb{R}^{\SequenceLength\times \AttentionHeadDimension}\) for sequence length \(\SequenceLength\) and head dimension \(\AttentionHeadDimension\), standard attention evaluates: \[ \AttentionScoreMatrix= \AttentionQueryMatrix\AttentionKeyMatrix^T \in \mathbb{R}^{\SequenceLength\times \SequenceLength}, \qquad \AttentionWeightMatrix= \text{softmax}(\AttentionScoreMatrix) \in \mathbb{R}^{\SequenceLength\times \SequenceLength}, \qquad \AttentionHeadOutput= \AttentionWeightMatrix\AttentionValuesMatrix\in \mathbb{R}^{\SequenceLength\times \AttentionHeadDimension} \]

To prevent numerical overflow, row-wise softmax subtracts the row maximum (the log-sum-exp shift of section 1.3.5) \(\SoftmaxRowMaximum_i = \max_{1 \le j \le \SequenceLength} \AttentionScoreEntry\): \[ \AttentionWeightEntry = \frac{\exp(\AttentionScoreEntry - \SoftmaxRowMaximum_i)}{\sum_{k=1}^\SequenceLength\exp(\AttentionScoreSummand - \SoftmaxRowMaximum_i)} \]

Standard GPU implementations execute this in three discrete kernel launches:

1. Compute \(\AttentionScoreMatrix= \AttentionQueryMatrix\AttentionKeyMatrix^T\) in SRAM, then write \(\AttentionScoreMatrix\in \mathbb{R}^{\SequenceLength\times \SequenceLength}\) to HBM (\(\AsymptoticTightBound(\SequenceLength^2)\) writes).
2. Read \(\AttentionScoreMatrix\) from HBM, compute row max \(\SoftmaxRowMaximum\) and row sum \(\SoftmaxNormalizer\), compute \(\AttentionWeightMatrix= \text{softmax}(\AttentionScoreMatrix)\), write \(\AttentionWeightMatrix\in \mathbb{R}^{\SequenceLength\times \SequenceLength}\) to HBM (\(\AsymptoticTightBound(\SequenceLength^2)\) reads + writes).
3. Read \(\AttentionWeightMatrix\) and \(\AttentionValuesMatrix\) from HBM, compute \(\AttentionHeadOutput= \AttentionWeightMatrix\AttentionValuesMatrix\), write \(\AttentionHeadOutput\in \mathbb{R}^{\SequenceLength\times \AttentionHeadDimension}\) to HBM (\(\AsymptoticTightBound(\SequenceLength^2 + \SequenceLength\AttentionHeadDimension)\) reads + writes).

Total HBM memory traffic is \(\AsymptoticTightBound(\SequenceLength^2 + \SequenceLength\AttentionHeadDimension)\) elements. For \(\SequenceLength= 8192\), \(\AttentionHeadDimension = 128\) in FP16, \(\AttentionScoreMatrix\) requires \(8192^2 \times 2 \text{ B} = 134 \text{ MB}\) per head per layer. Because modern GPUs feature massive compute throughput relative to HBM bandwidth, repeatedly reading and writing intermediate matrices \(\AttentionScoreMatrix\) and \(\AttentionWeightMatrix\) leaves GPU execution units memory-bound.

Online softmax recurrence

Tiling softmax is difficult because computing \(\SoftmaxNormalizer_i = \sum_{k=1}^{\SequenceLength} \exp(\AttentionScoreSummand - \SoftmaxRowMaximum_i)\) requires the global maximum \(\SoftmaxRowMaximum_i\) across all \(\SequenceLength\) keys. The online softmax recurrence solves this by incrementally updating normalization statistics as sub-blocks arrive.

Theorem 1.2: Online softmax recurrence theorem

Let a vector \(\SoftmaxInputVector\in \mathbb{R}^\SoftmaxVectorLength\) be partitioned into two sub-vectors \(\SoftmaxInputVector^{(1)} \in \mathbb{R}^{\SoftmaxVectorLength_1}\) and \(\SoftmaxInputVector^{(2)} \in \mathbb{R}^{\SoftmaxVectorLength_2}\) with \(\SoftmaxVectorLength= \SoftmaxVectorLength_1 + \SoftmaxVectorLength_2\). Let the local statistics for block 1 be: \[ \SoftmaxRowMaximum^{(1)} = \max_{1 \le j \le \SoftmaxVectorLength_1} \SoftmaxInputScore_j^{(1)}, \qquad \SoftmaxNormalizer^{(1)} = \sum_{j=1}^{\SoftmaxVectorLength_1} \exp\left(\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}\right) \] and for block 2 be: \[ \SoftmaxRowMaximum^{(2)} = \max_{1 \le j \le \SoftmaxVectorLength_2} \SoftmaxInputScore_j^{(2)}, \qquad \SoftmaxNormalizer^{(2)} = \sum_{j=1}^{\SoftmaxVectorLength_2} \exp\left(\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}\right) \]

Then the combined maximum \(\MergedSoftmaxMaximum\), normalization sum \(\MergedSoftmaxNormalizer\), and accumulated output vector \(\MergedAttentionOutput\) satisfy the exact recurrence relations: \[ \begin{aligned} \MergedSoftmaxMaximum &= \max\left(\SoftmaxRowMaximum^{(1)}, \SoftmaxRowMaximum^{(2)}\right), \\ \MergedSoftmaxNormalizer &= \SoftmaxNormalizer^{(1)} e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum} + \SoftmaxNormalizer^{(2)} e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum}, \\ \MergedAttentionOutput &= \frac{\SoftmaxNormalizer^{(1)} e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum}}{\MergedSoftmaxNormalizer} \AttentionHeadOutput^{(1)} + \frac{e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum}}{\MergedSoftmaxNormalizer} \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}} \AttentionValuesMatrix_j^{(2)}, \end{aligned} \] where \(\AttentionHeadOutput^{(1)} = \frac{1}{\SoftmaxNormalizer^{(1)}} \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}} \AttentionValuesMatrix_j^{(1)}\) is the partial output accumulated from block 1.

Denominator recurrence

Expanding the combined sum \(\MergedSoftmaxNormalizer = \sum_{j=1}^\SoftmaxVectorLength e^{\SoftmaxInputScore_j - \MergedSoftmaxMaximum}\): \[ \begin{aligned} \MergedSoftmaxNormalizer &= \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \MergedSoftmaxMaximum} + \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \MergedSoftmaxMaximum} \\ &= \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\left(\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}\right) + \left(\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum\right)} + \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\left(\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}\right) + \left(\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum\right)} \\ &= e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum} \underbrace{\sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}}}_{\SoftmaxNormalizer^{(1)}} + e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum} \underbrace{\sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}}}_{\SoftmaxNormalizer^{(2)}} \\ &= \SoftmaxNormalizer^{(1)} e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum} + \SoftmaxNormalizer^{(2)} e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum}. \quad \blacksquare \end{aligned} \]

Output rescaling recurrence

The true normalized output over the concatenated sequence is \(\MergedAttentionOutput = \frac{1}{\MergedSoftmaxNormalizer} \sum_{j=1}^\SoftmaxVectorLength e^{\SoftmaxInputScore_j - \MergedSoftmaxMaximum} \AttentionValuesMatrix_j\): \[ \begin{aligned} \MergedAttentionOutput &= \frac{1}{\MergedSoftmaxNormalizer} \left[ \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \MergedSoftmaxMaximum} \AttentionValuesMatrix_j^{(1)} + \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \MergedSoftmaxMaximum} \AttentionValuesMatrix_j^{(2)} \right] \\ &= \frac{1}{\MergedSoftmaxNormalizer} \left[ e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum} \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}} \AttentionValuesMatrix_j^{(1)} + e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum} \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}} \AttentionValuesMatrix_j^{(2)} \right] \end{aligned} \]

Noting that \(\SoftmaxNormalizer^{(1)} \AttentionHeadOutput^{(1)} = \sum_{j=1}^{\SoftmaxVectorLength_1} e^{\SoftmaxInputScore_j^{(1)} - \SoftmaxRowMaximum^{(1)}} \AttentionValuesMatrix_j^{(1)}\), we substitute this directly: \[ \MergedAttentionOutput = \frac{\SoftmaxNormalizer^{(1)} e^{\SoftmaxRowMaximum^{(1)} - \MergedSoftmaxMaximum}}{\MergedSoftmaxNormalizer} \AttentionHeadOutput^{(1)} + \frac{e^{\SoftmaxRowMaximum^{(2)} - \MergedSoftmaxMaximum}}{\MergedSoftmaxNormalizer} \sum_{j=1}^{\SoftmaxVectorLength_2} e^{\SoftmaxInputScore_j^{(2)} - \SoftmaxRowMaximum^{(2)}} \AttentionValuesMatrix_j^{(2)}. \quad \blacksquare \]

Memory traffic of tiled attention

The recurrence lets each query tile be combined with one key-value tile at a time, so score tiles never need to leave SRAM. The remaining question is how many bytes then cross the HBM boundary.

Let \(\SramCapacity\) be the fast SRAM capacity (in words). Following Algorithm 1 of Dao et al. (2022), FlashAttention sets \(\KvTileSize = \lfloor \SramCapacity/(4\AttentionHeadDimension) \rfloor\) and \(\QueryBlockSize = \min(\lfloor \SramCapacity/(4\AttentionHeadDimension) \rfloor,\AttentionHeadDimension)\). The proof abstracts constant factors and buffer reuse, using \(\KvTileSize=\AsymptoticTightBound(\SramCapacity/\AttentionHeadDimension)\), \(\QueryBlockSize=\AsymptoticTightBound(\min(\SramCapacity/\AttentionHeadDimension,\AttentionHeadDimension))\), and \(\QueryBlockSize\KvTileSize=\mathcal{O}(\SramCapacity)\) for the on-chip working set.

The algorithm partitions \(\AttentionQueryMatrix\) into \(\QueryBlockCount = \lceil \SequenceLength/\QueryBlockSize \rceil\) blocks of size \(\QueryBlockSize \times \AttentionHeadDimension\), and partitions \(\AttentionKeyMatrix, \AttentionValuesMatrix\) into \(\KvTileCount = \lceil \SequenceLength/\KvTileSize \rceil\) blocks of size \(\KvTileSize \times \AttentionHeadDimension\). It executes two nested loops:

• Outer loop over \(\KvTileIndex= 1, \dots, \KvTileCount\): Load \(\AttentionKeyMatrix_j\) and \(\AttentionValuesMatrix_j\) into SRAM.
• Inner loop over \(\QueryBlockIndex= 1, \dots, \QueryBlockCount\): Load \(\AttentionQueryMatrix_i\), \(\AttentionHeadOutput_i\), and its softmax statistics; compute \(\AttentionScoreTile = \AttentionQueryMatrix_i \AttentionKeyMatrix_j^T\) in SRAM; then update and write \(\AttentionHeadOutput_i\) and the statistics without materializing \(\AttentionScoreTile\) in HBM.

This loop order is the source of the I/O reduction. A key-value tile remains resident while the algorithm visits every query tile, so each element of \(\AttentionKeyMatrix\) and \(\AttentionValuesMatrix\) crosses the HBM boundary once. The query and output tiles must be revisited for each key-value block, but the \(\QueryBlockSize\times \KvTileSize\) score tile is consumed on chip and discarded. Standard attention instead writes and rereads the full \(\SequenceLength\times \SequenceLength\) score and probability matrices.

Total memory access accounting

Total HBM accesses comprise two components: key and value reads, where each element of \(\AttentionKeyMatrix\) and \(\AttentionValuesMatrix\) is loaded once (contributing \(2\SequenceLength\AttentionHeadDimension\) accesses); and query, output, and statistics passes, where each of the \(\KvTileCount\) column blocks reads all of \(\AttentionQueryMatrix\), reads and writes \(\AttentionHeadOutput\), and updates \(\mathcal{O}(\SequenceLength)\) softmax statistics. The matrices dominate the statistics for head dimension \(\AttentionHeadDimension\geq 1\), so each pass contributes \(\AsymptoticTightBound(\SequenceLength\AttentionHeadDimension)\) accesses and all passes contribute \(\AsymptoticTightBound(\SequenceLength\AttentionHeadDimension \KvTileCount)\).

Because \(\KvTileCount=\lceil \SequenceLength/\KvTileSize\rceil=\AsymptoticTightBound(\SequenceLength\AttentionHeadDimension/\SramCapacity)\), the total is \[ \text{HBM}_{\text{Flash}} = \AsymptoticTightBound\left(\SequenceLength\AttentionHeadDimension + \SequenceLength\AttentionHeadDimension \KvTileCount\right) = \AsymptoticTightBound\left(\SequenceLength\AttentionHeadDimension + \frac{\SequenceLength^2 \AttentionHeadDimension^2}{\SramCapacity}\right) = \AsymptoticTightBound\left(\frac{\SequenceLength^2 \AttentionHeadDimension^2}{\SramCapacity}\right), \] where the final equality follows for \(\AttentionHeadDimension \leq \SramCapacity\leq \SequenceLength\AttentionHeadDimension\).

The bound counts communication rather than arithmetic. FlashAttention still evaluates the same dense score blocks and may recompute intermediate quantities during backpropagation; its advantage comes from keeping those intermediates out of HBM.

I/O reduction ratio

Comparing standard attention I/O to FlashAttention I/O: \[ \frac{\text{HBM}_{\text{Standard}}}{\text{HBM}_{\text{Flash}}} = \frac{\AsymptoticTightBound(\SequenceLength^2 + \SequenceLength\AttentionHeadDimension)}{\AsymptoticTightBound\left(\SequenceLength\AttentionHeadDimension + \frac{\SequenceLength^2 \AttentionHeadDimension^2}{\SramCapacity}\right)} = \AsymptoticTightBound\left(\frac{\SramCapacity}{\AttentionHeadDimension^2}\right) \quad \text{for } \SequenceLength\gg \max\left(\AttentionHeadDimension,\frac{\SramCapacity}{\AttentionHeadDimension}\right) \]

Systems Perspective 1.4: A100 I/O reduction

The asymptotic ratio \(\SramCapacity/\AttentionHeadDimension^2\) predicts how greater SRAM capacity can reduce HBM traffic, but hidden constants and kernel details prevent it from serving as an exact transfer count. In an NVIDIA A100 measurement, Dao et al. (2022) report the effect for a forward-and-backward attention workload with sequence length 1,024, head dimension 64, 16 heads, batch size 64, a key-padding mask, and no dropout:

• Standard attention: 35.3 GB of HBM reads and writes; 35.1 ms runtime.
• FlashAttention: 4.4 GB of HBM reads and writes; 11.7 ms runtime.

The measurements show the systems consequence of the I/O bound: substantially less HBM traffic coincides with lower runtime even though FlashAttention performs additional recomputation. They do not imply that every configuration becomes compute bound; tile size, sequence length, and kernel implementation determine the realized regime.

Dao, T., D. Y. Fu, S. Ermon, A. Rudra, and C. Ré. 2022. “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” Advances in Neural Information Processing Systems (NeurIPS) 35: 16344–59. https://doi.org/10.52202/068431-1189.

Tiling changes how many bytes cross the memory boundary for a fixed model. The last derivation applies the same accounting one level up and asks how large the model and its training set should be for a fixed compute budget.

Compute-optimal scaling laws

Determining the compute-optimal allocation between model parameters and training dataset tokens for a given floating-point budget \(\Operations\) sets how large a language model and its training set should be. This subsection derives the compute-optimal allocation using Lagrange multipliers and contrasts the original findings of Kaplan et al. (2020) with the compute-optimal corrections of Hoffmann et al. (2022). The compute-optimal frontier uses the result to test whether a training strategy is limited by data or by compute.

Notation note: This derivation writes model parameters as \(\Parameters\), dataset size in tokens as \(\DatasetSize\), and the training compute budget as \(\Operations\) total FLOPs; Kaplan et al. (2020) and Hoffmann et al. (2022) write \(\Glyph{N}\) for parameters and \(\Glyph{C}\) for compute. The fitted constants \(\ChinchillaIrreducibleLoss\), \(\ParameterScaleCoefficient\), \(\DataScaleCoefficient\), \(\ParameterScalingExponent\), and \(\DataScalingExponent\) and the Lagrange multiplier \(\ComputeConstraintMultiplier\) keep the papers’ letters and are local to this section, so \(\ChinchillaIrreducibleLoss\) is a loss rather than an energy and \(\DataScaleCoefficient\) is a fitted constant rather than batch size.

Parametric loss surface

The cross-entropy evaluation loss \(\LossFunction(\Parameters, \DatasetSize)\) of an autoregressive transformer language model trained with \(\Parameters\) non-embedding parameters on \(\DatasetSize\) tokens is modeled by the power-law formulation: \[ \LossFunction(\Parameters, \DatasetSize) = \ChinchillaIrreducibleLoss+ \frac{\ParameterScaleCoefficient}{\Parameters^\ParameterScalingExponent} + \frac{\DataScaleCoefficient}{\DatasetSize^\DataScalingExponent} \] where:

• \(\ChinchillaIrreducibleLoss\): Irreducible loss representing the entropy of natural language data.
• \(\frac{\ParameterScaleCoefficient}{\Parameters^\ParameterScalingExponent}\): Capacity bottleneck loss resulting from finite model parameters \(\Parameters\).
• \(\frac{\DataScaleCoefficient}{\DatasetSize^\DataScalingExponent}\): Dataset size bottleneck loss resulting from finite training tokens \(\DatasetSize\).
• \(\ParameterScaleCoefficient, \DataScaleCoefficient> 0\): Empirical scaling constants.
• \(\ParameterScalingExponent, \DataScalingExponent> 0\): Power-law scaling exponents for model size and dataset size, respectively.

Constrained optimization

The total floating-point operations required to train a transformer with \(\Parameters\) non-embedding parameters for \(\DatasetSize\) tokens is given by \[ \Operations= 6 \Parameters \DatasetSize\text{ FLOPs}, \] where \(2 \Parameters \DatasetSize\) accounts for the forward pass and \(4 \Parameters \DatasetSize\) accounts for the backward pass.

Given a fixed training compute budget \(\Operations_0\), we formulate the constrained optimization problem to minimize the reducible loss \(\ReducibleLoss(\Parameters, \DatasetSize) = \frac{\ParameterScaleCoefficient}{\Parameters^\ParameterScalingExponent} + \frac{\DataScaleCoefficient}{\DatasetSize^\DataScalingExponent}\): \[ \min_{\Parameters, \DatasetSize} \left( \frac{\ParameterScaleCoefficient}{\Parameters^\ParameterScalingExponent} + \frac{\DataScaleCoefficient}{\DatasetSize^\DataScalingExponent} \right) \quad \text{subject to} \quad \ComputeBudgetConstraint(\Parameters, \DatasetSize) = 6 \Parameters \DatasetSize- \Operations_0 = 0 \]

We construct the Lagrangian function \(\ComputeConstrainedLagrangian(\Parameters, \DatasetSize, \ComputeConstraintMultiplier)\): \[ \ComputeConstrainedLagrangian(\Parameters, \DatasetSize, \ComputeConstraintMultiplier) = \frac{\ParameterScaleCoefficient}{\Parameters^\ParameterScalingExponent} + \frac{\DataScaleCoefficient}{\DatasetSize^\DataScalingExponent} + \ComputeConstraintMultiplier(6 \Parameters \DatasetSize- \Operations_0) \]

Taking partial derivatives with respect to \(\Parameters, \DatasetSize\), and \(\ComputeConstraintMultiplier\) and setting them to zero yields: \[ \frac{\partial \ComputeConstrainedLagrangian}{\partial \Parameters} = -\ParameterScalingExponent \ParameterScaleCoefficient \Parameters^{-\ParameterScalingExponent- 1} + 6 \ComputeConstraintMultiplier \DatasetSize= 0 \implies \ParameterScalingExponent \ParameterScaleCoefficient \Parameters^{-\ParameterScalingExponent- 1} = 6 \ComputeConstraintMultiplier \DatasetSize, \tag{1} \] \[ \frac{\partial \ComputeConstrainedLagrangian}{\partial \DatasetSize} = -\DataScalingExponent \DataScaleCoefficient \DatasetSize^{-\DataScalingExponent- 1} + 6 \ComputeConstraintMultiplier \Parameters= 0 \implies \DataScalingExponent \DataScaleCoefficient \DatasetSize^{-\DataScalingExponent- 1} = 6 \ComputeConstraintMultiplier \Parameters, \tag{2} \] \[ \frac{\partial \ComputeConstrainedLagrangian}{\partial \ComputeConstraintMultiplier} = 6 \Parameters \DatasetSize- \Operations_0 = 0. \tag{3} \]

Multiplying Equation (1) by \(\Parameters\) and Equation (2) by \(\DatasetSize\) gives: \[ \ParameterScalingExponent \ParameterScaleCoefficient \Parameters^{-\ParameterScalingExponent} = 6 \ComputeConstraintMultiplier \Parameters \DatasetSize, \] \[ \DataScalingExponent \DataScaleCoefficient \DatasetSize^{-\DataScalingExponent} = 6 \ComputeConstraintMultiplier \Parameters \DatasetSize. \]

Equating these expressions yields the power-law optimality condition: \[ \ParameterScalingExponent \ParameterScaleCoefficient \Parameters^{-\ParameterScalingExponent} = \DataScalingExponent \DataScaleCoefficient \DatasetSize^{-\DataScalingExponent}. \] This condition states that at the compute-optimal allocation, the marginal loss reduction per fractional increase in parameter count must equal the marginal loss reduction per fractional increase in dataset tokens.

Compute-optimal exponents

From \(\ParameterScalingExponent \ParameterScaleCoefficient \Parameters^{-\ParameterScalingExponent} = \DataScalingExponent \DataScaleCoefficient \DatasetSize^{-\DataScalingExponent}\), solving for \(\DatasetSize\) in terms of \(\Parameters\) yields: \[ \DatasetSize^\DataScalingExponent= \left(\frac{\DataScalingExponent \DataScaleCoefficient}{\ParameterScalingExponent \ParameterScaleCoefficient}\right) \Parameters^\ParameterScalingExponent\implies \DatasetSize= \left(\frac{\DataScalingExponent \DataScaleCoefficient}{\ParameterScalingExponent \ParameterScaleCoefficient}\right)^{1/\DataScalingExponent} \Parameters^{\ParameterScalingExponent/\DataScalingExponent}. \]

Substituting \(\DatasetSize\) into the compute constraint \(\Operations= 6 \Parameters \DatasetSize\) gives: \[ \Operations= 6 \Parameters\left(\frac{\DataScalingExponent \DataScaleCoefficient}{\ParameterScalingExponent \ParameterScaleCoefficient}\right)^{1/\DataScalingExponent} \Parameters^{\ParameterScalingExponent/\DataScalingExponent} = 6 \left(\frac{\DataScalingExponent \DataScaleCoefficient}{\ParameterScalingExponent \ParameterScaleCoefficient}\right)^{1/\DataScalingExponent} \Parameters^{\frac{\ParameterScalingExponent+ \DataScalingExponent}{\DataScalingExponent}}. \]

Solving explicitly for the compute-optimal parameter count \(\Parameters^*(\Operations)\): \[ \Parameters^*(\Operations) = \left( \frac{\ParameterScalingExponent \ParameterScaleCoefficient}{\DataScalingExponent \DataScaleCoefficient} \right)^{\frac{1}{\ParameterScalingExponent+ \DataScalingExponent}} \left( \frac{\Operations}{6} \right)^{\frac{\DataScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent}} \propto \Operations^{\frac{\DataScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent}}. \]

By symmetry, solving for the compute-optimal token count \(\DatasetSize^*(\Operations)\) yields: \[ \DatasetSize^*(\Operations) = \left( \frac{\DataScalingExponent \DataScaleCoefficient}{\ParameterScalingExponent \ParameterScaleCoefficient} \right)^{\frac{1}{\ParameterScalingExponent+ \DataScalingExponent}} \left( \frac{\Operations}{6} \right)^{\frac{\ParameterScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent}} \propto \Operations^{\frac{\ParameterScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent}}. \]

Empirical exponents and the \(\DatasetSize/\Parameters\approx 20\) rule

Fitting the parametric loss surface to over 400 experimental runs, Hoffmann et al. (2022) estimated \(\ParameterScalingExponent\approx 0.34\) and \(\DataScalingExponent\approx 0.28\). Substituting these values gives exponents close to the paper’s parametric-model result:

\[ \frac{\DataScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent} \approx 0.45, \qquad \frac{\ParameterScalingExponent}{\ParameterScalingExponent+ \DataScalingExponent} \approx 0.55 \]

Thus the parametric model predicts: \[ \Parameters^*(\Operations) \propto \Operations^{0.45}, \qquad \DatasetSize^*(\Operations) \propto \Operations^{0.55} \]

The paper’s two direct empirical approaches report near-balanced exponents of \((0.50, 0.50)\) and \((0.49, 0.51)\), which motivate the common square-root shorthand. Under the parametric fit, however, \(\DatasetSize^*/\Parameters^*\) grows slowly with compute rather than remaining strictly constant. Across the regime studied, roughly 20 training tokens per parameter is a useful Chinchilla heuristic, not a universal invariant.

Systems Perspective 1.5: Kaplan vs. Chinchilla scaling

The difference between Kaplan et al. (2020) and Hoffmann et al. (2022) shows how the training schedule used in empirical fits can alter system design conclusions:

• Kaplan scaling model: Kaplan et al. (2020) fit a different loss form with its own exponents and reported compute-optimal growth of \(\Parameters\propto \Operations^{0.73}\) and \(\DatasetSize\propto \Operations^{0.27}\); this split does not follow from substituting exponents into the two-term surface above. This led the field to scale model sizes aggressively while holding training dataset sizes relatively small (e.g., GPT-3 175B trained on 300B tokens, giving \(\DatasetSize/\Parameters= 1.71\)).
• Methodological divergence: Kaplan et al. (2020) used a fixed token budget and learning-rate schedule across models. Hoffmann et al. (2022) matched the schedule to each training horizon and found that the earlier setup understated the value of training smaller models on more data.
• Chinchilla revision: Hoffmann et al. (2022) found near-balanced parameter and token scaling: its three approaches report exponent pairs of \((0.50, 0.50)\), \((0.49, 0.51)\), and \((0.46, 0.54)\) rather than one exact universal ratio.
• Practical impact: Chinchilla (70B parameters trained on 1.4T tokens, \(\DatasetSize/\Parameters= 20\)) outperformed GPT-3 (175B parameters trained on 300B tokens) across downstream benchmarks while requiring \(2.5\times\) less compute during inference (inference cost scales with parameter count) and \(2.5\times\) lower weight memory at the same precision (Hoffmann et al. 2022).

Hoffmann, Jordan, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, et al. 2022. “Training Compute-Optimal Large Language Models.” Advances in Neural Information Processing Systems (NeurIPS) 35: 30016–30. https://doi.org/10.52202/068431-2176.

Kaplan, Jared, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. “Scaling Laws for Neural Language Models.” ArXiv Preprint abs/2001.08361.

Summary

An algorithm becomes a systems workload through its shapes, operation graph, data dependencies, and numerical representation. Linear-algebra primitives, memory layouts, broadcasting, and sparse formats determine work, traffic, locality, and usable parallelism. Automatic differentiation and backpropagation add saved activations and backward computation; computational graphs then expose opportunities for fusion and recomputation, and the log-sum-exp shift keeps the softmax at the loss finite in reduced precision. Among the derivations, the Baur–Strassen result bounds reverse-mode overhead, while FlashAttention shows how tiling and an online recurrence can reduce HBM traffic without changing attention semantics.

The same accounting connects kernels to scaling decisions. Compute-optimal scaling turns a fixed budget into a model-data allocation, but the fitted exponents remain empirical rather than universal. Before selecting an optimization, count the operations, the bytes moved, and the state saved for the backward pass.