Examines efficient transformer architectures, hardware co-design, inference optimizations like speculative decoding.
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Modern language models have achieved remarkable capabilities, but they come with a price tag that makes deployment difficult for most organizations. GPT-3 reportedly cost over $4 million to train a single time. Serving a 175B-parameter model requires dozens of high-end GPUs just to handle inference. A smartphone cannot run these models at all. These constraints have pushed researchers to ask a question that turns out to be surprisingly deep: can we get the same capabilities for a fraction of the compute?
The answer, increasingly, is yes. Efficiency research has become one of the most active areas in language AI, and it does not just mean shrinking models. It means redesigning architectures to avoid unnecessary computation, co-designing software and hardware so neither wastes the other's capacity, and rethinking the entire training and inference pipeline from first principles.
This chapter covers four interconnected fronts. First, we examine efficient architectures: changes to the transformer design that reduce the quadratic cost of attention, eliminate redundant computation, and pack more capability into fewer parameters. Second, we look at hardware co-design: how modern AI accelerators work, what their bottlenecks are, and how architectures can be shaped to match the machine. Third, we explore inference efficiency: the techniques that make deployment faster and cheaper once a model is trained, from quantization to speculative decoding. Finally, we look at training efficiency: how recent advances in data quality, mixed-precision arithmetic, and gradient checkpointing have dramatically reduced the cost of training large models from scratch.
As we discussed in the chapter on Model Compression, techniques like knowledge distillation and pruning can reduce a model's size after training. This chapter complements that work by focusing on efficiency baked into the architecture and training process itself. The next chapter, Capability Frontiers, explores whether these efficiency gains are enabling qualitatively new capabilities in small models.
The original transformer architecture was designed for correctness and flexibility, not for computational efficiency. Its self-attention mechanism computes relationships between every pair of tokens, which scales quadratically with sequence length. For a 1,000-token sequence, this means one million attention computations per layer. For a 100,000-token sequence, it means ten billion. The architecture also allocates compute uniformly across all positions and all layers, even when most positions are not informative and most layers are doing similar work.
Efficient architecture research attacks these inefficiencies from several directions.
Standard attention computes a full attention matrix for a sequence of tokens. The quadratic cost arises because every query token attends to every key token:
where , , are the query, key, and value matrices, and is the head dimension. This is in time and in memory just to store the attention matrix.
Sparse attention replaces the full matrix with a structured pattern. Instead of every token attending to every other token, each token attends to a limited neighborhood. The two most influential patterns are:
- Local window attention: each token attends only to the tokens to its left and right. This reduces complexity to , which is linear in sequence length for fixed window size.
- Strided attention: tokens attend locally plus to tokens at fixed stride intervals, allowing global information to propagate over multiple layers.
Longformer and BigBird combine local attention for most tokens with global attention for special tokens like [CLS]. The global tokens can see all positions, giving the model a way to aggregate information across the entire sequence without paying the full cost everywhere.
Linear attention takes a different approach. Recall that the softmax in standard attention creates a non-linear function that cannot be factored:
If we replace softmax with a kernel function such that approximates the original attention score, we can rewrite attention as:
where:
- is the feature map applied to the query vector
- is the feature map applied to the -th key
- is the -th value row
The key insight is that can be computed once and reused for every query, reducing the total cost to , which is linear in sequence length. Performers and Linear Transformers use this idea with random Fourier features as the kernel approximation.
The tradeoff is that linear attention loses some expressiveness. The softmax non-linearity helps attention sharply focus on a few relevant tokens, while kernel approximations tend to produce softer distributions. In practice, linear attention works well on tasks that do not require sharp token selection, but lags on tasks requiring precise retrieval.
A standard transformer layer applies the same feed-forward network to every input token. This means the model uses the same parameters regardless of whether it is processing a Python variable name or a French sentence. Mixture-of-Experts (MoE) architectures replace the single feed-forward network with many expert sub-networks and route each token to only a few of them.
The gating mechanism uses a learned linear transformation followed by softmax:
where:
- is the token embedding
- is the gating weight matrix
- keeps only the top- logits and sets the rest to
- is optional noise for load balancing during training
The output of the MoE layer is:
The efficiency gain is striking. A model with 8 experts and top-2 routing has the same computational cost per token as a model with one feed-forward network, but four times the total parameter count. Mixtral 8x7B, for example, has 46.7B total parameters but only uses about 12.9B active parameters per token, giving it performance comparable to 70B dense models at a fraction of the inference cost.
MoE architectures introduce two key challenges. The first is load balancing: if the gating network routes all tokens to the same few experts, the unused experts contribute nothing. Training typically includes an auxiliary load-balancing loss that encourages uniform expert utilization. The second is communication overhead: in distributed training, different experts may live on different devices. Routing tokens to remote experts requires network communication, which can become a bottleneck at scale.
A different approach to the quadratic attention problem is to abandon the attention mechanism entirely and replace it with architectures that process sequences with linear or sub-linear complexity.
State Space Models (SSMs) model the sequence transformation as a linear dynamical system. The input sequence is processed through a recurrence:
where:
- is the hidden state
- , , , are learned matrices
- The bars indicate that and are discretized from continuous-time parameters
The recurrent form processes each token in time with memory (constant in sequence length). SSMs can also be unrolled into a convolutional form for efficient parallel training:
where is a learned convolution kernel derived from the SSM parameters. The convolutional form enables parallel training on GPUs while the recurrent form enables efficient autoregressive inference.
Mamba is the most influential SSM architecture for language modeling. Its key innovation is selective state spaces: the transition matrices and depend on the input, making the model input-dependent rather than linear and time-invariant. This selective mechanism allows Mamba to focus its hidden state on the most relevant parts of the input history, similar to how attention focuses on relevant tokens. Mamba achieves competitive performance with transformers on many language tasks at complexity, though its performance on tasks requiring precise token retrieval still lags behind full attention.
RetNet takes a hybrid approach, showing that transformers with decay-based attention (rather than softmax) can be computed in three equivalent ways: recurrently for inference (linear complexity), in parallel for training (quadratic complexity), and in chunks for a middle ground. This "train in parallel, infer recurrently" duality was an important insight because it separated the training computation graph from the inference computation graph.
Transformers typically apply the same number of layers to every input token. A trivial question, like "What is 2 + 2?", receives the same number of forward passes as a complex reasoning problem. This is wasteful.
Early exit architectures add classifier heads after intermediate layers. During inference, if the model's confidence at layer exceeds a threshold , it exits early without computing the remaining layers:
This reduces average inference depth on easy inputs while preserving full-depth computation for hard ones. DeeBERT and FastBERT demonstrated that classification tasks can often exit by layer 3-4 out of 12, reducing computation by 50-70% with minimal accuracy loss.
A related idea is mixture-of-depths: routing tokens to skip certain layers. If a token is already well-represented after 6 layers, it may not need to pass through layers 7-12. This allows the model to dynamically allocate depth to the tokens that need it.
Both early exit and mixture-of-depths challenge a core assumption of the standard transformer: that all computation should be homogeneous. In practice, language is heterogeneous. Some tokens are rare and semantically loaded; others are grammatical connectives that carry little information. Some sentences require careful multi-step reasoning; others are simple pattern completions. Depth efficiency exploits this heterogeneity by allocating computation adaptively. The catch is that training these models is harder: the model must learn useful representations and reliable policies for when to stop. Distillation from a full-depth teacher helps, as does using differentiable threshold mechanisms that allow the exit decision to be trained end-to-end.
Algorithm design alone does not determine efficiency. The same algorithm can be ten times faster or slower depending on how well it matches the underlying hardware. Understanding what makes AI accelerators fast is essential for designing architectures that deliver their theoretical efficiency gains.
Modern AI training and inference happens primarily on GPUs (and increasingly on specialized AI chips like TPUs, H100s, and Trainium). Understanding their architecture explains why some operations are fast and others are not.
A GPU contains thousands of CUDA cores organized into streaming multiprocessors (SMs). All these cores share a memory hierarchy with dramatically different bandwidth at each level:
- Registers: fastest, local to each thread, limited size
- Shared memory / L1 cache: on-chip, ~10-50 MB total, ~100x slower than registers
- L2 cache: on-chip, ~40-80 MB, shared across SMs
- HBM (High Bandwidth Memory): off-chip, ~80 GB, bandwidth ~3 TB/s on H100
The critical bottleneck for most transformer operations is memory bandwidth, not compute. Moving data from HBM to the compute cores is slow relative to the arithmetic operations themselves. An operation is compute-bound if the GPU's arithmetic units are the bottleneck, and memory-bound if data movement is the bottleneck. On modern hardware, most attention operations are memory-bound.
This memory-bandwidth bottleneck explains why attention is slower in practice than its FLOP count suggests. The standard attention computation reads the , , and matrices from HBM, writes the intermediate attention matrix back to HBM, then reads it again to multiply with . Every round trip to HBM costs time.
FlashAttention is one of the most impactful efficiency innovations in the transformer era. It computes exact (not approximate) attention using a tiling algorithm that minimizes HBM reads and writes.
The key insight is that the attention computation:
does not need to materialize the full attention matrix in HBM. Instead, FlashAttention tiles the computation into blocks that fit in SRAM (shared memory). It processes one tile at a time, computing a local softmax and updating the output matrix, keeping track of the normalization constants needed for the global softmax using an online algorithm.
The online softmax computation maintains two accumulators per query: the running maximum and the running sum of exponentials . When a new tile of keys arrives with scores , the global softmax is updated without ever materializing the full row:
where:
- is the running maximum of attention logits for query
- is the running sum of exponentiated, shifted logits
- is the partial output for query
By keeping all intermediate values in fast SRAM rather than writing to HBM, FlashAttention reduces HBM reads/writes from to , making it 2-4x faster than standard attention in practice and enabling context lengths that would otherwise cause out-of-memory errors.
FlashAttention-2 extended this by improving work partitioning across thread blocks, reducing non-matrix-multiplication FLOPs, and achieving ~70% GPU utilization compared to ~35% for the original. FlashAttention-3 specifically targets the H100 architecture, exploiting its asynchronous pipeline and FP8 support.
Training very large models requires distributing computation across multiple GPUs. The two main strategies are tensor parallelism and pipeline parallelism.
Tensor parallelism splits individual weight matrices across GPUs. For a feed-forward layer with , we split column-wise across GPUs:
Each GPU computes its portion , and the results are gathered with an all-reduce operation. This allows a single layer to use more GPUs than would fit in a single device's memory. The communication cost is an all-reduce per layer per forward and backward pass.
Pipeline parallelism splits the model's layers across GPUs. GPU 1 holds layers 1-8, GPU 2 holds layers 9-16, etc. The input flows through GPUs in a pipeline. Without careful scheduling, most GPUs sit idle waiting for their predecessor. The GPipe pipeline fills this bubble by processing multiple micro-batches simultaneously, so each GPU always has work to do.
ZeRO (Zero Redundancy Optimizer) is a memory optimization that partitions optimizer states, gradients, and parameters across data-parallel workers instead of replicating them. The basic version reduces optimizer state memory by a factor equal to the number of workers, which can be 64x for large training runs.
Hardware supports different numerical formats with different tradeoffs between precision and memory:
- FP32: 32-bit float, standard training format, 4 bytes per value
- BF16: 16-bit bfloat, same exponent range as FP32, 2 bytes per value, now the standard for training
- FP16: 16-bit float, smaller exponent range, prone to overflow, requires loss scaling
- FP8: 8-bit float, supported on H100, approximately 2x throughput over FP16
- INT8: 8-bit integer, well-supported for inference on many platforms
- INT4: 4-bit integer, increasingly common for inference, requires careful quantization
The NVIDIA H100 GPU can perform FP8 matrix multiplications at approximately twice the throughput of FP16, with 2x more values fitting in the same memory bandwidth. Exploiting these low-precision formats requires that the model's activations and weights stay within the representable range.
Quantization-aware training (QAT) simulates low-precision arithmetic during training by inserting fake quantization operations:
where is the number of bits and is the quantization step size. During the forward pass, values are quantized. During the backward pass, gradients flow through using the straight-through estimator (the gradient of the rounding operation is treated as 1). This allows the model to learn representations that retain accuracy under low-precision arithmetic.
Training a model is a one-time cost. Inference is a continuous cost that scales with usage. For a large language model serving millions of users, inference efficiency directly determines whether a product is economically viable.
Autoregressive generation is inherently sequential: to generate token , the model needs all preceding tokens . Without optimization, this requires recomputing the attention keys and values for all previous tokens at each step, making generation in the total number of generated tokens.
The KV-cache solves this by storing the key and value tensors for all processed tokens and reusing them at each step. Token only needs to compute its own queries and attend to the cached keys and values:
where and are retrieved from cache. This reduces generation from to per token.
The problem is that the KV-cache is large. For a 7B-parameter model with 32 layers, 32 heads, and 128 head dimension, each token's KV-cache occupies:
For a context of 8,192 tokens, that is 4 GB just for the KV-cache. Serving 100 concurrent users requires 400 GB of GPU memory just for KV-caches, far exceeding what a single GPU holds. Managing KV-cache memory efficiently is one of the central challenges in LLM serving.
Grouped Query Attention (GQA) reduces KV-cache size by sharing key and value heads across multiple query heads. Instead of separate KV heads for query heads, GQA uses KV groups where . With , this becomes Multi-Query Attention (MQA), where all query heads share a single KV head. The memory reduction is -fold. Llama 3 uses GQA with 8 key-value heads for 32 query heads, reducing KV-cache by 4x with minimal quality loss.
Post-training quantization (PTQ) reduces model weights and/or activations to lower bit-widths without retraining. For inference on edge devices or in memory-constrained serving environments, this can make the difference between a model fitting or not.
GPTQ is an influential PTQ method for weight-only quantization to INT4 or INT8. It uses second-order gradient information (the inverse Hessian of the weight perturbation) to optimally round each weight:
where is the set of quantized weight matrices, and is a calibration input. By solving this layer-by-layer, GPTQ can quantize 175B-parameter models to INT4 in a few hours on a single GPU, with perplexity degradation of 1-3% compared to FP16.
AWQ (Activation-Aware Weight Quantization) makes a key observation: weights that are multiplied by large-magnitude activations are more important and should receive higher precision. By identifying the "salient" weights based on activation statistics and applying a per-channel scaling before quantization, AWQ achieves similar or better results than GPTQ with a simpler algorithm.
GGUF and llama.cpp have made quantized inference accessible on consumer hardware. A 7B model in 4-bit quantization occupies about 4 GB, allowing it to run on a GPU with 8 GB of VRAM or even on CPU with sufficient RAM. The throughput on CPU is low (a few tokens per second) but the accessibility has enabled a large open-source inference ecosystem.
Speculative decoding is a clever approach to accelerating autoregressive generation by exploiting the observation that a small "draft" model can propose tokens that a large "verifier" model can evaluate in parallel.
The algorithm works as follows. A small draft model generates draft tokens autoregressively:
The large verifier model then evaluates all draft tokens simultaneously in a single forward pass (since the verifier is not autoregressive at this point, just evaluating given tokens):
Each draft token is accepted with probability , where is the draft model's probability. This rejection sampling scheme guarantees that the accepted tokens follow exactly the verifier model's distribution, giving no quality degradation.
When many draft tokens are accepted, the effective throughput of the large model increases because one verifier forward pass generates multiple tokens. With a 7B draft model and a 70B verifier, acceptance rates of 70-80% on typical text produce about 2-3x speedup compared to running the 70B model alone.
The speedup depends on two factors: the acceptance rate and the relative latency of draft versus verifier. If the draft model is much smaller than the verifier (say, 10x smaller in parameters), its contribution to total latency is small, and high acceptance rates translate directly to throughput gains.
Traditional serving systems process one request at a time or batch requests of the same length together. Both are wasteful. A single-request system leaves GPU compute idle whenever the model waits for memory. Fixed-length batching requires padding all sequences to the longest one in the batch, wasting compute on padding tokens.
Continuous batching (also called iteration-level batching) processes requests at the granularity of individual decoding steps. At each step, the serving system looks at all in-flight requests, determines which are currently generating a new token, and batches those tokens together for the forward pass. New requests can be inserted and completed requests can be removed at each step, keeping the GPU fully utilized.
PagedAttention addresses the KV-cache fragmentation problem. The KV-cache for each request grows as tokens are generated, but traditional systems pre-allocate a maximum-length block of contiguous memory for each request. This leads to fragmentation: a request that generates 100 tokens wastes the memory reserved for 1900 more (if the max length is 2000). PagedAttention stores KV-cache in fixed-size pages (similar to virtual memory paging in operating systems), allocating pages as needed and maintaining a page table mapping logical token positions to physical memory locations. This reduces memory waste to near zero and allows the system to serve 2-4x more concurrent requests.
vLLM is the most widely used implementation of these ideas. It combines continuous batching with PagedAttention and has become the de facto standard for high-throughput LLM serving.
The cost of training large models has been a fundamental constraint on AI development. Training GPT-3 once cost around $4 million; training GPT-4 is estimated to have cost over $100 million. Efficiency improvements here compound directly into more training runs, more experimentation, and faster iteration.
The shift from FP32 to mixed-precision training (BF16 or FP16 for forward and backward passes, FP32 for optimizer states) was one of the first major training efficiency wins. It halves memory usage for most tensors and doubles throughput on hardware with tensor cores.
Loss scaling is required for FP16 training because the smaller exponent range causes underflow in gradients. The loss is multiplied by a large scale factor before the backward pass, the gradients are unscaled before the optimizer step, and the scale factor is adjusted dynamically (increased when no overflow occurs, decreased when it does).
BF16 avoids this problem. With the same exponent range as FP32 but only 8 bits of mantissa (versus 24 for FP32), BF16 covers the same dynamic range but at lower precision. The reduced precision in the mantissa is generally acceptable because gradient noise in stochastic optimization already introduces similar-scale errors.
Standard backpropagation requires storing all layer activations computed during the forward pass, because they are needed to compute gradients during the backward pass. For a transformer with layers and batch size , this requires storing activations in memory, which grows linearly with model depth.
Gradient checkpointing (also called activation recomputation) trades compute for memory. During the forward pass, only a subset of layer outputs (checkpoints) are stored. During the backward pass, the activations at non-checkpointed layers are recomputed from the nearest preceding checkpoint. By checkpointing every layers, the memory for activations is reduced from to at the cost of approximately 33% additional compute per training step.
At large scale, this tradeoff is almost always worthwhile. Training a 70B-parameter model with full activation storage would require enormous memory, while checkpointing reduces this to a manageable level at a modest compute overhead.
The datasets used for training have grown alongside models, but data quality matters more than quantity. Recent work on data filtering and dataset curation shows that training on a smaller, higher-quality dataset often outperforms training on a much larger noisy one.
Chinchilla scaling laws (Hoffmann et al., 2022) established that previous large models were undertrained. The optimal training configuration allocates roughly equal compute to model size and training data: for a model with parameters, the compute-optimal number of training tokens is approximately . A 70B-parameter model should train on about 1.4 trillion tokens. This insight shifted the field from "largest possible model" to "right-sized model with enough data."
Curriculum learning pre-trains models in stages, presenting easier or more structured data first and more complex data later. Code training, for example, benefits from showing well-formatted code early before introducing partially broken or obscure examples. The rationale is that the model builds a strong foundation on clear examples before having to handle noisy or ambiguous ones.
Repetition and replay are important considerations for long training runs. Seeing the same data too many times causes memorization and reduces generalization. Many recent training pipelines do multiple passes over high-quality data (e.g., textbooks, scientific papers) but only one pass over noisier sources, explicitly managing the per-source repetition rate.
Standard Adam requires storing first and second moment estimates for every parameter, tripling the optimizer memory footprint beyond the parameters themselves. For a 70B model, Adam adds roughly 420 GB of optimizer state (at FP32) on top of the 140 GB for the parameters.
Adafactor reduces this by factoring the second moment estimate. For a weight matrix , instead of storing a full second moment matrix, Adafactor stores rank-1 factors and such that . This reduces the second moment storage from to , a substantial savings for large weight matrices.
CAME (Confidence-guided Adaptive Memory Efficient optimization) extends this with a confidence matrix that distinguishes reliable gradient directions from noisy ones, improving training stability on large models.
Sophia uses an approximate Hessian diagonal rather than gradient squared as the second moment estimate, which produces better step size calibration and has shown faster convergence on language modeling benchmarks.
Deep learning frameworks like PyTorch execute operations one at a time by default. Each operation involves a kernel launch, memory allocation, and HBM read/write. For many small operations (layer normalizations, activation functions, elementwise additions), these overhead costs dominate the actual computation.
Kernel fusion combines multiple operations into a single GPU kernel. Instead of separate kernels for the bias addition, GELU activation, and dropout in a feed-forward layer, a single fused kernel reads the input once, applies all three operations, and writes the output once. This can reduce memory bandwidth by 3-5x for these operations.
torch.compile (introduced in PyTorch 2.0) automates kernel fusion using the Triton JIT compiler. By tracing the computation graph and identifying fusible operations, it can speed up training loops by 20-30% with a single line of code: model = torch.compile(model). This makes kernel fusion accessible without requiring custom CUDA code.
XLA (Accelerated Linear Algebra), used by JAX and TensorFlow, takes a more aggressive approach by compiling entire training steps into optimized machine code. It performs whole-program optimization including operation reordering, buffer aliasing, and cross-operation fusion, often achieving higher utilization than PyTorch's per-operation approach.
This section implements several key efficiency techniques: quantization, KV-cache management, and a simple demonstration of speculative decoding concepts.
In[3]:
Code
Let's implement a simple uniform quantization scheme to see how bit-width affects representational precision.
In[4]:
Code
Out[5]:
Console
The table shows how quantization error increases as bit-width decreases. At 8 bits, the mean squared error is tiny, and the weight values are nearly indistinguishable from the original. At 4 bits, the error grows but remains small enough for most practical purposes. At 2 bits, reconstruction quality degrades substantially, which is why 2-bit quantization is rarely used in practice without special techniques. The compression ratio grows inversely with bit-width: 8-bit weights occupy 4x less memory than FP32, 4-bit weights occupy 8x less.
Out[6]:
Visualization
Let's simulate the KV-cache to understand how its size grows with context length and how Grouped Query Attention reduces it.
In[7]:
Code
Out[8]:
Console
The table shows how dramatically context length affects KV-cache requirements. At 32,768 tokens, a standard Multi-Head Attention (MHA) model requires several gigabytes just for the KV-cache of a single request. With Grouped Query Attention (GQA, 8 KV heads), this is reduced by 4x. Multi-Query Attention (MQA, 1 KV head) reduces it by 32x but at a higher quality cost. These tradeoffs explain why Llama 3 and Mistral models use GQA by default.
Out[9]:
Visualization
Arithmetic intensity determines whether an operation is compute-bound or memory-bound. Let's compute it for key transformer operations.
In[10]:
Code
Out[11]:
Console
This analysis reveals why FlashAttention matters so much. The matrix multiplication and softmax operations in attention are memory-bound on modern GPUs, while the feed-forward layer matrix multiplications are compute-bound. FlashAttention specifically targets the memory-bound attention operations by reducing HBM accesses, which is why it delivers such substantial speedups.
Out[12]:
Visualization
Let's implement a simplified speculative decoding simulation to illustrate the acceptance-rejection mechanism.
In[13]:
Code
Out[14]:
Console
The simulation shows how speculative decoding speedup scales with acceptance rate. At 50% acceptance, the expected speedup is about 1.4x because draft-model overhead offsets part of the gain. At 85-95% acceptance (typical for a well-matched draft-verifier pair on natural text), the idealized speedup is roughly 2.6-3.2x. The key insight is that the draft model cost is small relative to the verifier, so even moderate acceptance rates produce meaningful gains.
Out[15]:
Visualization
Let's compute theoretical compute requirements for standard transformer versus MoE architectures.
In[16]:
Code
Out[17]:
Console
The computation confirms the MoE efficiency story. With top-2 routing across 8 experts, each MoE layer uses roughly 2x the feed-forward FLOPs of a dense layer (because 2 of 8 experts are active), but the model has 8x more total feed-forward capacity. The practical result is that Mixtral 8x7B competes with Llama 2 70B on many benchmarks while being much faster at inference.
The key parameters for the efficiency techniques covered in this chapter are:
- Quantization bits: Controls the tradeoff between precision and memory. 8-bit is nearly lossless; 4-bit requires careful calibration; 2-bit degrades quality significantly.
- KV heads (GQA): Number of key-value groups. Full MHA uses heads; MQA uses 1; GQA uses an intermediate value (4-8 is common). Reduces KV-cache memory proportionally.
- Speculative decoding k: Number of draft tokens per round. Larger increases potential speedup but also wastes more compute when tokens are rejected.
- Gradient checkpointing interval: Layers between saved checkpoints. Checkpointing every layers is theoretically optimal; in practice, checkpointing every layer or every other layer is common.
- Expert count and top-k (MoE): Total experts and how many are active per token. Common choices are 8 experts with top-2 or 64 experts with top-2. Increasing experts increases total capacity without increasing active FLOPs.
- FlashAttention tile size: Determined automatically based on SRAM size; controlled implicitly by block size hyperparameter in FlashAttention implementations.
Efficiency gains in research settings do not always translate cleanly to production. Several practical limitations deserve careful attention.
Hardware heterogeneity is a constant challenge. FlashAttention is highly optimized for NVIDIA GPUs and exploits specific features of their memory hierarchy. On AMD GPUs, Apple Silicon, or custom AI chips, the same algorithm may have different (sometimes worse) performance characteristics. Teams using diverse hardware stacks must carefully benchmark claimed efficiency gains on their actual infrastructure before committing to architectural choices.
Accuracy-efficiency tradeoffs are task-dependent. A quantization scheme that preserves accuracy on general text benchmarks may degrade significantly on specialized domains like mathematics or code generation, where precise numeric values and low-frequency tokens matter more. Similarly, sparse attention that works well for document classification may struggle on tasks requiring long-range dependency tracking across thousands of tokens.
MoE training instability is a known problem. The gating network can collapse to routing all tokens to a few experts, effectively wasting the capacity of the others. Load balancing losses help but do not fully solve this. At very large scales (hundreds of experts), achieving stable expert utilization requires careful tuning of the auxiliary loss weight, the noise injection strategy, and the capacity factor (maximum fraction of tokens any one expert can receive per batch).
Speculative decoding requires a matched draft model. The draft model must be from the same model family as the verifier, or at minimum trained on the same tokenizer and similar data. A mismatched draft model produces tokens from a different distribution, reducing acceptance rates to near zero. Maintaining a paired draft model adds operational complexity: when the verifier is updated or fine-tuned, the draft model may need to be updated too.
Compilation overhead is a real cost. torch.compile adds 30-120 seconds of compilation time on the first call. For short-running workloads or models that change frequently (as in fine-tuning experiments), the compilation overhead can outweigh the runtime savings. Compilation is most beneficial for long production inference deployments or extended training runs.
Despite these limitations, efficiency research has fundamentally changed what is possible. Models that required clusters of A100 GPUs to run now fit on a laptop. Training runs that took months now take weeks. Inference that cost dollars per query now costs fractions of a cent. The efficiency frontier has moved consistently toward more capable models at lower cost, and this trajectory shows no signs of stopping.
This chapter surveyed the major fronts of efficiency research in language AI:
Efficient architectures reduce the computational cost inherent in the transformer design. Sparse and linear attention address the quadratic scaling of standard attention, either by restricting the attention pattern (local windows, strided patterns) or by approximating softmax with a kernel function. Mixture-of-experts architectures decouple parameter count from compute cost by routing each token to a small subset of specialized sub-networks. Sub-quadratic models like Mamba process sequences in linear time using structured state spaces, trading some expressiveness for dramatically lower complexity.
Hardware co-design aligns algorithms with the realities of GPU memory hierarchies. The key insight is that most transformer operations are memory-bandwidth-bound rather than compute-bound, which is why FlashAttention produces such large speedups despite computing exactly the same result as standard attention. Understanding arithmetic intensity helps predict which operations will benefit from optimization and which are already near the hardware limit.
Inference efficiency makes deployment economically viable. The KV-cache converts quadratic generation into linear generation, but its memory footprint is a new bottleneck. Grouped Query Attention reduces KV-cache by sharing key-value heads. Post-training quantization shrinks model weights to INT4 or INT8, enabling models to run on consumer hardware. Speculative decoding uses a small draft model to propose tokens that a large verifier accepts in parallel, achieving 2-3x throughput gains. Continuous batching and PagedAttention maximize GPU utilization in serving systems.
Training efficiency reduces the cost of building new models. Mixed-precision training cuts memory usage and increases throughput. Gradient checkpointing trades compute for memory, enabling larger batch sizes. Chinchilla scaling laws revealed that optimal training requires matching compute allocation between model size and training tokens. Compiler-based kernel fusion, implemented in torch.compile and XLA, automatically merges small operations to minimize memory round-trips.
These efficiency advances are not independent. An MoE model benefits from FlashAttention in its attention layers, quantization in its weights, and speculative decoding at inference time. The cumulative effect of stacking multiple efficiency techniques is multiplicative: a model that is 2x faster from architecture changes, 2x faster from quantization, and 2x faster from speculative decoding runs 8x faster than the naive baseline. This stacking is how modern systems achieve inference on large models at costs that would have seemed implausible just a few years ago.
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about efficiency frontiers in language AI.
Comments
Nothing yet. Say the first thing.
Sign in to join the conversation.