Chizoba Obasi blog

FlashAttention: Part 2

Thu, 17 Sep 2026 00:00:00 +0000

This is Part 2 of a two-part FlashAttention series. If you haven’t read Part 1 yet, start there. Part 1 covers the fundamental memory problem, the mathematical trick, and the FA1 implementation.

From FA1 to FA4: The Evolution of FlashAttention

The reference handbook for this blog post was inspired by this tweet and is titled: Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention.

In Part 1, we covered the fundamental problem FlashAttention-1 (FA1) addresses, the mathematical tricks it uses, the GPU implementation of those tricks, and the architectural compatibility (MHA/MQA/GQA) of FlashAttention.

In this blog post, we continue the story through the evolution of the FlashAttention family (FA2, FA3, FA4), its relationship to other efficient attention techniques (PagedAttention, sparse, linear), the training vs. inference distinction, current PyTorch integration, and the practical mental model to walk away with.

1. FlashAttention Evolution

1.1 FlashAttention-2: what changed
1.2 FA2 parallelism across sequence tiles
1.3 FA2 work partitioning and non-matmul FLOPs
1.4 FlashAttention-3: the Hopper generation
1.5 FA3 asynchrony: overlap data movement, GEMM, and softmax
1.6 FA3 FP8: performance without pretending precision is free
1.7 FlashAttention-4: the Blackwell generation
1.8 FA4 and asymmetric hardware scaling
1.9 FA4 implementation and current status
1.10 FlashAttention-1 through FA-4 compared

2. Using FlashAttention in Frameworks

2.1 PyTorch scaled-dot-product attention today
2.2 Exactness is not bitwise identity

3. FlashAttention vs. Other Techniques

3.1 FlashAttention vs. PagedAttention
3.2 FlashAttention vs. sparse and linear attention

4. Training vs. Inference

4.1 Training, prefill, and decode are different regimes

5. Practical Engineering for FlashAttention

5.1 Common implementation mistakes
5.2 Common Misconceptions

6. Conclusion

6.1 The mental model to have
6.2 FlashAttention Summary

1. FlashAttention Evolution

1.1 FlashAttention-2: What Changed

FA2 didn’t change the goal. It still computes exact attention without materializing quadratic intermediates. However, it made much better use of the GPU.¹ The FA2 paper identifies three main changes:

Reduce the number of non-matmul FLOPs.
Parallelize attention across sequence tiles, even for a single head, to improve occupancy.
Repartition work among warps within a thread block to reduce shared-memory communication.

The roughly 2x speedup of FA2 over FA1 on A100 in the original paper’s experiments is real,¹ but it’s a benchmark- and GPU-specific number that depends on dtype, shapes, causal mode, and software versions. It is not a universal speed guarantee.

The question FA2 is really asking: “How can we utilize the GPU more effectively?”

The FlashAttention pattern across generations

FA1’s breakthrough was IO-aware tiled exact attention.
FA2’s breakthrough was better parallelism and work partitioning around the same high-level algorithmic idea.
Every subsequent FlashAttention evolution since has followed the same pattern: keep the semantics, and adapt the kernel pipeline to the resources and bottlenecks of newer GPU architectures.

1.2 FA2 Parallelism Across Sequence Tiles

The practical scenario where this matters:

batch size is small,
number of heads is small,
sequence length is huge.

In that case, FA1 might not expose enough independent work. A GPU has many processing clusters known as streaming multiprocessors (SMs) that need enough independent work to be kept busy. If we only parallelize over batch times heads, some of those units sit idle. FA2 adds parallelism along the sequence dimension, so different query-row tiles can be processed independently. That gives the GPU more work to schedule, since more resident blocks means more available parallel work.

Let’s consider a simplified mental picture:

Q rows [ block 0 | block 1 | block 2 | block 3 | block 4 | block 5 ]
           ↓                   ↓                   ↓
          CTA                 CTA                 CTA

Query rows get split into blocks (block 0, block 1, …, block 5), and each block gets assigned to a CTA.
A CTA, Cooperative Thread Array also known as a thread block, is a logical grouping of threads that execute together on a single SM.

More parallel blocks don’t automatically mean better performance. We still have constraints: occupancy, registers, shared memory, tile size, data reuse. Kernel design is an optimization problem over hardware resources, not just a request for maximum thread count.²

Key insight

Long sequence length itself can become a source of parallelism.

1.3 FA2 Work Partitioning and Non-Matmul FLOPs

FA2 also changes how work is divided among warps in a thread block. In FA1, warp partitioning could require partial results to be written to shared memory and then combined.¹ That creates communication overhead.

Picture it: a warp writes a partial result to shared memory, another warp reads it, combines it, and writes again. Every round trip through shared memory costs cycles.

[warp 0] --> shared memory --> [warp 1] --> shared memory --> ...

FA2 repartitions the work so warps cooperate in a way that cuts down those shared-memory reads and writes, thereby reducing communication. It also reduces non-matmul operations via algebraic manipulation around softmax rescaling and backward computations.¹ This matters because:

A FLOP isn’t a universal unit of performance across all operation types.
Matrix multiply-accumulate (MMA) has significantly higher throughput than exponentials, scalar reductions, conversions, and shared-memory synchronization.
A tensor-core matmul and an exponential aren’t equivalent from a hardware-throughput perspective.

We can track the bottlenecks addressed by each FlashAttention generation so far:

FA1 attacked HBM traffic.
FA2 exposed the following as the next important bottlenecks: occupancy, warp communication, non-matmul operations.

FA2 also simplifies what gets saved for backward. Instead of keeping probabilities around, it stores a single row-wise log-sum-exp value per query row and recomputes the tiles, keeping the same linear-memory footprint.

Bottlenecks shift

A useful optimization principle from the FlashAttention evolution so far: after fixing one bottleneck, another bottleneck emerges.

1.4 FlashAttention-3: The Hopper Generation

Why does FA3 exist at all? GPU hardware changed. Hopper introduced capabilities that changed the best way to schedule attention, and FA2 didn’t fully exploit what the H100 could do. The FA3 paper measured FA2 at only around 35% utilization on H100, which leaves a lot of headroom.

FA3 was designed around NVIDIA Hopper GPUs and attacks that utilization gap with four broad ideas:

Asynchronous computation and data movement.
Warp specialization.
Interleaving GEMM (General Matrix Multiplication) and softmax work.
FP8 support.

The important point: The mathematical algorithm didn’t suddenly change. The execution workflow did, to exploit Hopper’s hardware architecture.

On H100, the NeurIPS 2024 publication reports a 1.5 to 2.0x speedup over FA2 in its benchmark suite. BF16 throughput reaches up to 840 TFLOPs/s (85% utilization), and FP8 reaches 1.3 PFLOPs/s.³ (The earlier arXiv preprint reported 740 TFLOPs/s and close to 1.2 PFLOPs/s; the figures here are from the final NeurIPS version.) As before, these are hardware- and benchmark-specific empirical results, not universal across all models.

Exactness vs. precision

FA3 is not “a more approximate FlashAttention.” The FP16/BF16 path preserves the exact-attention algorithmic goal. The FP8 path intentionally introduces lower-precision arithmetic and therefore needs its own numerical-accuracy discussion.

1.5 FA3 Asynchrony: Overlap Data Movement, GEMM, and Softmax

Picture a naive pipeline:

load tile → GEMM → softmax → PV GEMM → load next tile

Each stage waits for the previous one. That leaves periods where some hardware resources sit idle while others work. Tensor cores go quiet during softmax. Memory buses go quiet during GEMM. FA3 builds a more overlapping pipeline:

Load K/V tile 1
       Compute tile 1
              Load K/V tile 2
                     Compute tile 2
                            Load K/V tile 3
                                   Compute tile 3

More explicitly:

Load K/V tile 1
         QK^T GEMM
            softmax/update
                     PV GEMM
                          Load K/V tile 2
                                  QK^T GEMM
                                     softmax/update
                                              PV GEMM

The exact implementation is more sophisticated than the toy depictions above, but the architectural intuition is what matters. Two Hopper capabilities make this overlapping possible:

Tensor Memory Accelerator (TMA), which moves data asynchronously.
Asynchronous warp-group matrix-multiply-accumulate (WGMMA), which lets tensor-core work proceed without blocking the issuing warp.

FA3 also uses warp specialization to allow data movement and computation to overlap.³ Different warps take responsibility for different stages of the pipeline: some handle data movement, some issue GEMMs, some handle softmax and update. That way a warp that’s waiting on a memory transfer doesn’t stall the whole block.

FA3 also uses a ping-pong style schedule to interleave block matrix multiplication and softmax. The goal is to avoid leaving tensor cores idle while non-matmul operations (scalar and special-function work) are performed, and vice versa.³

The online-softmax recurrence itself hasn’t changed. The normalization and value-accumulation stages still do what Part 1 describes in FlashAttention-1. What changed is the hardware pipeline used to execute them.

Progression: What bottleneck is being addressed?

FA1 asks “How do we tackle HBM traffic?”
FA2 asks “How do we partition & parallelize the work better?”
FA3 asks “How do we maximize the use of distinct execution resources on Hopper?”

1.6 FA3 FP8: Performance Without Pretending Precision Is Free

Now let's focus on precision. Hopper gives extremely high tensor-core throughput for FP8. That’s tempting, but it isn’t free. Simply converting every attention operand to FP8 produces unacceptable numerical error.² Attention contains a chain of operations:

\[QK^T \rightarrow e^{x} \rightarrow \frac{e^{x}}{\sum e^{x}} \rightarrow PV.\]

These operations have different numerical sensitivities. Dot products tolerate some loss. The exponential is much more sensitive. Normalization involves division. The value accumulation sums over long sequences. Each stage has its own error profile.

FA3 introduces an FP8 path with two techniques to control that error:

Block quantization, which scales values relative to their local tile rather than globally.
Incoherent processing, which multiplies Q and K by a random orthogonal matrix to spread outlier values across dimensions before quantization.

In the final NeurIPS 2024 publication, FA3 reaches 1.3 PFLOPs/s on H100 with this path, and reports 2.6x lower numerical error than the baseline FP8 attention method used for comparison.³

So low precision doesn’t automatically buy free performance. It shifts error into specific parts of the pipeline, therefore we have to design around that shift. This is exactly where “exact” starts to mean two different things:

Algorithmic exactness. The algorithm still represents $\mathrm{softmax}(QK^T)V$. No pairs are skipped, no low-rank approximation is made.
Numerical precision. The actual calculation may use FP8, and therefore experiences quantization and rounding error.

An algorithm can be exact in its mathematical formulation while a particular low-precision implementation of that algorithm is numerically less accurate. Both things are true at once, and neither cancels the other. A model’s acceptable precision depends on the full training and inference setup. FP8 support is a hardware-and-numerics feature layered onto the FlashAttention execution strategy, not proof that FP8 is universally interchangeable with BF16 or FP16.²

One more scope note: FA3’s FP8 path is Hopper-specific. The official repository still separates FA3’s implementation from other backends, so these aren’t universal features of “FlashAttention.”⁴

1.7 FlashAttention-4: The Blackwell Generation

In FA4, the core pattern of hardware-aware redesign continues, this time for NVIDIA Blackwell.

Blackwell exhibits asymmetric hardware scaling. Tensor-core throughput roughly doubled relative to Hopper, but several other resources didn’t scale at the same rate. Shared-memory bandwidth and exponential throughput, in particular, lag behind.⁵

This changes which stage is the bottleneck. See the illustration below:

Tensor cores GEMM   [  |  |  |  |  |  |  |  |  |  ]
Softmax/exponential [  |  ]

The matrix multiplication has become extremely fast. Now the secondary softmax/exponential work becomes relatively expensive, and as such becomes the bottleneck. When one subsystem becomes much faster than the others, operations that used to be secondary can become the bottleneck.²

On Hopper, tensor cores were fast, but the gap to softmax and shared-memory bandwidth was narrower.
On Blackwell, the matrix-multiplication bar gets much longer, and the softmax/exponential bar stays roughly the same length, as shown in the illustration above.

FA4 redesigns both the forward and backward pipelines rather than just reusing the Hopper schedule. The major improvement techniques include:

Fully asynchronous MMA pipelines and larger tiles.
Software-emulated exponential plus conditional online-softmax rescaling in the forward pass.
Tensor Memory (TMEM) and 2-CTA MMA techniques to reduce shared-memory traffic and atomic additions in the backward pass.

On B200 with BF16, the paper reports up to 1.3x speedup over cuDNN 9.13 and 2.7x over its Triton comparison, reaching 1613 TFLOPs/s (71% utilization) under the evaluated configurations.⁵

Same algorithm, different balance

FA4 does not change the big-O arithmetic of dense attention. It attacks the new bottlenecks created by a GPU generation whose compute, memory, and function-unit balance differs from Hopper’s.

1.8 FA4 and Asymmetric Hardware Scaling

The FlashAttention lineage is a case study in why kernels can’t be optimized once and assumed optimal forever.² This section contains a broader systems lesson. Let’s look at two hypothetical GPU generations (throughput in FLOP/s):

FLOP/s	GPU A	GPU B
GEMM	100	200
Softmax	50	50
Memory	50	50

GPU B makes GEMM 2x faster. But the overall system doesn’t become 2x faster automatically. Now that GEMM is less of a bottleneck, the relative cost of everything else goes up. Softmax and memory traffic, which were already on the critical path, are now even more clearly on it. This is asymmetric hardware scaling. The ratio between subsystems changes, and the optimal algorithm changes with it.

Asymmetric hardware scaling is one of the major systems principles behind the FA1 → FA4 evolution.

Blackwell is a case of this. Tensor-core throughput outran softmax and exponential throughput. So the forward pass becomes constrained by softmax and exponential work rather than GEMM. FA4 responds with a software-emulated exponential and by skipping online-softmax rescaling when the running maximum doesn’t need it. Both moves take pressure off the relatively slower non-matmul units.

The backward pass has a different bottleneck. FA4 uses Blackwell’s Tensor Memory (TMEM) and 2-CTA MMA mode to cut shared-memory traffic and atomic accumulation overhead.

This is why “FA4 is just a faster FA3” misses the point. The algorithm and kernel pipeline are co-designed around the new asymmetries, which is reflected in the paper title: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling.⁵

That is also why performance numbers should always name the GPU generation. A speedup on B200 doesn’t reveal anything about H100, A100, or a non-NVIDIA accelerator. The ratios are the whole story.

The systems principle

A kernel is balanced against a particular hardware ratio. If tensor-core throughput improves faster than memory bandwidth or special-function throughput, the optimal algorithmic pipeline can change even when the mathematical function is identical.

1.9 FA4 Implementation and Current Status

FlashAttention-4 is a CuTe DSL implementation, optimized for Hopper and Blackwell GPUs such as the H100 and B200, and installed via the flash-attn-4 package.⁴ CuTe DSL is a Python-based domain-specific language from NVIDIA’s CUTLASS for writing and just-in-time compiling high-performance GPU kernels.

One engineering motivation behind the FA4 choice is compile-time productivity. The FA4 paper reports substantially faster compilation, roughly 20 to 30 times, compared with the traditional C++ template-based approach used in its comparison, while still retaining the required low-level expressivity.⁵ As of September 2026, PyPI classifies flash-attn-4 as Development Status :: 3 - Alpha, and its latest release (4.0.0b32, uploaded September 23, 2026) is a pre-release.⁶

That’s the caveat. Current implementation status is not the same thing as an evergreen property of FlashAttention. Packages, GPU support, CUDA compatibility, PyTorch, hardware, and framework integration all evolve quickly.

The durable lesson is the algorithmic progression. The current library support layer changes much faster than the underlying algorithmic ideas. Algorithmic knowledge evolves slowly. Library support can change overnight. When deploying FlashAttention, it’s worth separating the two. The algorithm tells us what should be possible. The current library tells us what’s actually available on your specific GPU, CUDA version, and framework.

Implementation snapshot

Treat this as a current implementation snapshot, not an evergreen property. “FA4 exists in the official repository” and “every production environment should replace FA2 with FA4” are very different claims.

1.10 FlashAttention-1 Through FA-4 Compared

Here’s the whole evolution on one page:

Version	Main Problem	Main Solution
FA1	Excessive memory (IO/HBM) traffic	Tiling + online softmax + fused exact attention + IO awareness
FA2	GPU underutilization	More sequence-level parallelism + better warp-level work partitioning + fewer non-matmul FLOPs
FA3	Hopper execution imbalance	Asynchronous TMA/WGMMA pipeline + GEMM-softmax overlap + warp specialization + FP8 path
FA4	Blackwell hardware asymmetry	Fully async MMA pipeline + larger tiles + optimized non-matmul work (software exp / conditional online-softmax rescale, TMEM, 2-CTA backward techniques)

A few things to hold onto when reading those numbers:

These are four successive generations of efficient implementations of attention, not four different attention mechanisms. The mathematical target remains fundamentally the same for dense attention.
The exact feature matrix differs across implementations. FA3 is strongly associated with Hopper-specific optimization, while the current FA4 repository targets both Hopper and Blackwell through its CuTe DSL path.⁴

2. Using FlashAttention in Frameworks

2.1 PyTorch Scaled-Dot-Product Attention Today

Modern PyTorch provides a high-level API:⁷

torch.nn.functional.scaled_dot_product_attention(
    query, key, value,
    attn_mask=None,
    dropout_p=0.0,
    is_causal=False,
    scale=None,
    enable_gqa=False,
)

Conceptually, this function computes $\mathrm{softmax}(QK^T / \sqrt{d})V$. Internally, PyTorch picks an optimized implementation based on the inputs. The available backends include FlashAttention-2, a memory-efficient attention implementation, cuDNN attention, and the C++ math implementation. To restrict which backend runs, use torch.nn.attention.sdpa_kernel:⁸

from torch.nn.attention import SDPBackend, sdpa_kernel
import torch.nn.functional as F

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

If a requested fused kernel can’t run for the given inputs, PyTorch can warn with reasons when fallbacks are disabled. Backend selection and eligibility depend on device, dtype, shape, mask, and other arguments.⁸ The torch.nn.attention docs also expose registration and activation hooks for newer implementations, including FA3 and FA4. This surface moves fast, so treat the installed PyTorch version’s docs as the source of truth.⁹

Which kernel actually ran

Calling the high-level API does not necessarily mean “I know exactly which kernel executed.”
If we need to know which backend is running, restrict it explicitly with sdpa_kernel or check the PyTorch version’s registration APIs. In production and benchmarking, verify rather than assume.

2.2 Exactness Is Not Bitwise Identity

Suppose two implementations, a math kernel and a FlashAttention kernel, compute the same dense SDPA definition. There will still be small numerical differences. The reason is straightforward. Floating-point addition is not associative:

\[(a + b) + c \neq a + (b + c)\]

in finite precision, for some values. Tiling changes the reduction order. Fusion can change when values are rounded. Accumulation precision can differ.

Consider a simple example:

\[f(Q, K, V) = g(Q, K, V) = a + b + c + d\]

But:

Implementation A sums in order $1, 2, 3, 4$: $(((a + b) + c) + d)$
Implementation B sums in order $3, 1, 4, 2$: $(((c + a) + d) + b)$

Same mathematical function, slightly different floating-point bits.

PyTorch’s reproducibility docs note this directly: SDPA backends can produce different results because they accumulate floating-point values in different orders.¹⁰ So the same mathematical function is not necessarily the same floating-point bits. Validation should use numerical tolerances, and those tolerances depend on:

FP32 vs. BF16/FP16/FP8
accumulation precision
sequence length and score distribution
hardware and backend
dropout and randomness

Bitwise equality shouldn’t be used as the definition of exactness.

Determinism is a separate question from exactness. Libraries expose backend- and version-specific controls for it, and those controls can carry performance or memory costs.⁴ $^,$ ¹⁰

Exactness vs Bitwise identity

Exact algorithm means no deliberate mathematical approximation of dense attention. Bitwise identical implementation means every finite-precision operation lands on exactly the same bits. FlashAttention promises the first idea, not the second across arbitrary kernels.

3. FlashAttention vs. Other Techniques

3.1 FlashAttention vs. PagedAttention

The names sound related. They’re not. FlashAttention and PagedAttention answer two completely different questions:

FlashAttention: how do I compute $QK^T$, softmax, and $PV$ efficiently?
PagedAttention: how do I store and retrieve the growing KV cache for many serving requests?

	FlashAttention	PagedAttention
Main problem	Efficient attention computation	KV-cache memory management
Main setting	Attention execution	LLM serving
Mechanism	Tiling, online softmax, fused/hardware-aware kernels	Paging KV cache
Main state	Attention tiles and row-wise normalization/output state	Persistent per-request KV cache
Changes dense attention formula?	No for dense FlashAttention	No, primarily changes cache memory management
Specific task to tackle	$N^2$ score/probability intermediates	fragmented per-request KV-cache allocation

PagedAttention, introduced with vLLM, reduces KV-cache waste from fragmentation and duplication in high-throughput serving.¹¹ FlashAttention doesn’t manage KV cache allocation. It handles how the attention operation consumes $Q/K/V$. How systems access the persistent KV state that lives across requests is a separate concern handled by PagedAttention. That separation is why they can coexist in the same serving stack without stepping on each other. Dao-AILab even ships a KV-cache-oriented FlashAttention interface with optional block tables. Kernel execution and cache paging are two modular concerns that play nicely together.⁴

In practice, the stack looks like this:

Paged KV cache → FlashAttention-style kernel → Attention output

When to use

If the problem statement contains “$N^2$ score/probability intermediates,” think FlashAttention.
If it contains “fragmented per-request KV-cache allocation,” think PagedAttention.

3.2 FlashAttention vs. Sparse and Linear Attention

Three very different strategies, grouped together because they all make attention cheaper in some regime. They are not the same kind of change.

	Dense FlashAttention	Sparse Attention	Linear Attention
What it does	Same dense attention, different execution	Fewer interactions	Different mathematical formulation
Attention pattern	Dense (all pairs)	Sparse or local	Rearranged via feature maps or associativity
Scaling	Still $O(N^2 d)$	Can fall below full $N^2$ work	Can fall below full $N^2$ work

Dense FlashAttention

computes $\mathrm{softmax}(QK^T)V$ for all required Q-K pairs.
changes the execution of dense attention, not the mathematical function.

Sparse / local attention

only computes a selected subset of $Q$-$K$ interactions.
$i \rightarrow j$ only if $\lvert i-j \rvert \leq w$ where $w$ is the window size.
fewer $Q$-$K$ interactions, and the model’s attention pattern changes.

Linear attention

changes the mathematical formula itself.
usually rearranges the attention computation with feature maps or associativity, avoiding an explicit dense softmax score matrix.

The phrase “memory-efficient attention” is too broad to identify a method by itself. Ask two questions instead:²

Does it compute the same dense softmax attention?
Does it reduce arithmetic pairs, or mostly memory traffic/intermediate storage?

The first question splits dense FlashAttention from sparse and linear approaches. The second separates FlashAttention from methods that change the number of interactions that get computed.

The original FlashAttention paper makes this line explicit. It contrasts its dense exact algorithm with approximate methods on one side, and explores block-sparse FlashAttention as an approximate/sparse extension on the other.¹²

FlashAttention reduces memory requirements and IO, but dense attention remains quadratic.

In summary:

FlashAttention: same dense attention, different execution.
Sparse attention: fewer interactions, different attention pattern.
Linear attention: different mathematical formulation.

FlashAttention doesn't make attention linear

“FlashAttention makes attention linear” is wrong. If a system shows near-linear scaling because it uses a local window or another sparse pattern, the sparsity is what changed the number of interactions. FlashAttention may still be the kernel underneath.

4. Training vs. Inference

4.1 Training, Prefill, and Decode Are Different Regimes

FlashAttention shines when both query and key sequence dimensions are large. Plenty of query rows means avoiding the materialized attention matrix actually pays off, and tiled GEMMs have real work to do.

Training. Suppose $N_q = N_k = N$. There are many query and key rows. FlashAttention can exploit large GEMMs, tiling, and massive parallelism, and it avoids $N^2$ intermediates. This is a very favorable regime.

Prefill. Suppose the user sends a prompt of $8000$ tokens. The model processes those tokens together. Again $N_q \approx N_k$, so dense attention is substantial. FlashAttention can be very useful here. Prefill is typically compute-bound.

Decode. Now suppose the model has already generated $8000$ tokens and wants to generate token $8001$. The new query might have $N_q = 1$ while $N_k = 8000$. So the computation is now:

\[1 \times 8000 \text{, not } 8000 \times 8000\]

The bottleneck changes. Decode is heavily interacting with the existing KV cache. The system’s behavior is now driven by how fast it can read that cache, not by how much arithmetic it can do. Decode is typically memory-bound. Serving and batching behavior matter here too.² The official repository exposes flash_attn_with_kvcache, which updates and attends to the cache in one kernel and supports MQA/GQA and optional RoPE handling.⁴

The key distinction in one line:

Prefill: many query rows → large attention computation.
Decode: one query row → large KV-cache memory access.

FlashAttention also cannot eliminate other costs. MLP layers, communication in tensor parallelism, KV-cache capacity, model-weight bandwidth, sampling, and scheduler overhead remain separate concerns.²

So when someone says “FlashAttention makes LLM inference faster,” ask:

Prefill or decode?
Compute-bound or memory-bound?
Sequence length?
Batch size?
KV cache?
GPU?
Which kernel?

Phase and bottleneck matter

This does not mean FlashAttention is irrelevant to inference. Prefill is an attention-heavy dense regime, and decode can still use specialized kernels from the same implementation family. It means “FlashAttention speeds up LLM inference” is incomplete unless you say which phase and bottleneck.

5. Practical Engineering for FlashAttention

5.1 Common Implementation Mistakes

A running list of things that bite people in practice:²

Assuming the flash backend always ran. High-level frameworks may fall back when dtype, device, shape, or mask is unsupported. Use backend controls and profiling when it matters.
Leaving dropout on during evaluation. PyTorch SDPA applies dropout_p exactly as supplied, so pass 0.0 when evaluation requires no dropout.
Using the wrong mask convention. Boolean mask semantics differ across APIs.
Expecting bitwise agreement. Fused kernels can reorder floating-point operations.
Ignoring GQA head constraints. Current interfaces require compatible/divisible query and KV head counts.
Treating current head-dimension/dtype support as timeless. These constraints change across releases and backends.
Benchmarking without synchronization or warmup. Asynchronous GPU execution can make naive wall-clock measurements meaningless.
Calling a sparse or local pattern “dense FlashAttention.” Kernel and attention pattern are separate choices.
Confusing prefill with decode. The shapes and bottlenecks are different.
Assuming the newest generation is always deployable. Current FA4 package metadata is alpha, and hardware/software compatibility must be checked.

5.2 Common Misconceptions

A list of things people say that aren’t quite right about FlashAttention, and what’s actually true.²

“FlashAttention approximates softmax.”
No. Dense FlashAttention evaluates the dense softmax-attention function using tiled online normalization. Nothing is thrown away, nothing is approximated.
“FlashAttention makes attention $O(N)$.”
No. Dense arithmetic remains $O(N^2 d)$. What becomes linear is the large auxiliary attention-memory footprint with respect to sequence length. HBM traffic is reduced, but the underlying number of query-key interactions is unchanged.
“It is just kernel fusion.”
Fusion is part of the implementation story. But online-softmax tiling and IO-aware scheduling are algorithmic, not merely concatenating existing kernels. Fusion alone wouldn’t produce the same memory profile or the same IO complexity bound.
“It is the same as PagedAttention.”
No. PagedAttention manages KV-cache allocation for serving. FlashAttention manages how the attention operation itself is executed. Different problems, different solutions.
“Exact means bitwise identical.”
No. Floating-point operation order can change rounding. Two exact implementations of the same function can produce slightly different bits, and that’s expected.
“Long context becomes free.”
No. Dense pairwise arithmetic is still quadratic, and KV-cache and other model costs remain. What FlashAttention gives us is a smaller memory footprint and less HBM traffic, not a free lunch.
“FA1, FA2, FA3, and FA4 are different attention architectures.”
No. They are successive generations of efficient kernels and algorithms, shaped by different hardware bottlenecks. The mathematical target is fundamentally the same for dense attention.

6. Conclusion

6.1 The Mental Model to Have

Think of FlashAttention as a streaming matrix computation with exact streaming softmax. Picture the full attention matrix, $Q$ on the $y$-axis and $K$ on the $x$-axis. It’s a big grid.

\[Q \quad \begin{array}{|c|c|c|c|} \hline & & & \\ \hline & & & \\ \hline & & & \\ \hline & & & \\ \hline \end{array} \\\] \[\qquad K\]

Naive attention:
- Calculate the entire matrix grid and store it.
Flash attention:
- Bring one small region of the matrix grid close to the compute units. Calculate that region, normalize it using running statistics, immediately use the result to accumulate the output, then discard the region. Then move on to the next region.
- The mathematical interactions remain. The physical representation of the intermediate computation does not. That’s the central insight.

The evolution from FA1 to FA4, in one line each:

FA1 → “Reduce expensive HBM movement.”
FA2 → “Now use the GPU more effectively via better parallelism.”
FA3 → “Now overlap communication and memory movement on Hopper.”
FA4 → “Now redesign the pipeline for Blackwell’s new hardware balance.”

This evolution is why FlashAttention is much more interesting than just “a faster attention kernel.” It’s a case study in algorithm-hardware co-design. The mathematical function can stay identical while the optimal execution strategy changes dramatically as the memory hierarchy, compute throughput, and specialized hardware change.

The final practical mental model:

keep score tiles near compute,
carry only the row statistics needed to merge softmax blocks,
immediately consume probabilities for each tile into the $V$ accumulation, and
avoid sending the full $N^2$ attention matrix through HBM.

6.2 FlashAttention Summary

Let’s walk through the whole FlashAttention story in one derivation, and hit the key points to remember.

Ordinary attention: $QK^T \in \mathbb{R}^{N \times N}$. Materializing it creates an enormous intermediate matrix.¹³

\[O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V.\]

Naive execution: This creates huge memory traffic.

\[QK^T \to \text{store } N \times N \text{ scores} \to \text{softmax} \to \text{store } N \times N \text{ probabilities} \to \text{multiply by } V\]

FlashAttention: Tiling is used. Split $Q \to Q_i$ and $K, V \to K_j, V_j$, then calculate the score tile:

\[S_{ij} = \frac{Q_i K_j^T}{\sqrt{d}}.\]

The problem with tiling: Softmax needs the maximum and the denominator over the whole row.
The solution: online softmax
1. Maintain $(m, \ell, a)$ where:
  - $m$ = running maximum
  - $\ell$ = running sum of exponentials (the softmax denominator)
  - $a$ = running unnormalized value-weighted sum (accumulator)
2. When a new block has a larger maximum $m’$, calculate a rescaling factor $\alpha$ and rescale the old state ($\ell$ and $a$) by it:
\[\alpha = e^{m - m'}\]
Therefore, we can process the full FlashAttention pipeline without ever materializing the $N \times N$ matrix:
\[\text{tile} \to \text{score} \to \text{online softmax} \to V \text{ accumulation} \to \text{discard tile}.\]
Result: The same exact dense attention mathematics, but much better memory behavior: less HBM traffic, much smaller intermediates. Arithmetic remains $O(N^2 d)$.

FlashAttention in one table:

Property	Complexity / Formulation
Memory complexity	$O(Nd)$ auxiliary attention storage
Compute complexity	$O(N^2 d)$ dense attention compute
Attention formulation	Exact dense softmax attention

Key takeaway

Same exact dense attention mathematics, but much better memory behavior. Less HBM traffic, much smaller intermediates, arithmetic unchanged at $O(N^2 d)$.

Appendix

Reference Map

The table below shows what paper to refer to for each question.²

Question	Best starting source
What is the original IO-aware algorithm?	FlashAttention-1 paper ¹²
Why does online normalization work?	Milakov & Gimelshein ¹⁴
Was exact subquadratic-memory attention known earlier?	Rabe & Staats ¹⁵
What specifically changed in FA2?	FlashAttention-2 paper ¹
What is Hopper-specific in FA3?	FlashAttention-3 paper ³
What changes on Blackwell in FA4?	FlashAttention-4 paper ⁵
What does the current official package support?	Dao-AILab repository ⁴
How does current PyTorch select SDPA kernels?	PyTorch SDPA and `sdpa_kernel` docs ⁷ $^,$ ⁸
Why can exact backends differ numerically?	PyTorch reproducibility notes ¹⁰
How is PagedAttention different?	Kwon et al. ¹¹

A note on how to use this map

Implementation support changes faster than the mathematics. For deployed systems, treat the installed framework’s documentation and the exact kernel repository/release as authoritative for supported devices, dtypes, head dimensions, masks, dropout, and GQA behavior.

References

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026FlashAttentionPt2,
  title   = "FlashAttention: Part 2",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Sep",
  url     = "https://chizkidd.github.io/2026/09/17/flashattention-2/"
}

Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. ↩ ↩² ↩³ ↩⁴ ↩⁵
Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention. Google Drive PDF. Accessed September 2026. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹ ↩¹⁰ ↩¹¹
Jay Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024. ↩ ↩² ↩³ ↩⁴ ↩⁵
Dao-AILab. flash-attention official repository and current implementation documentation. GitHub. Accessed September 2026. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷
Ted Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv, 2026. ↩ ↩² ↩³ ↩⁴ ↩⁵
PyPI. flash-attn-4 package metadata and release history. PyPI. Accessed September 2026. ↩
PyTorch. torch.nn.functional.scaled_dot_product_attention documentation. PyTorch Documentation. Accessed September 2026. ↩ ↩²
PyTorch. torch.nn.attention.sdpa_kernel documentation. PyTorch Documentation. Accessed September 2026. ↩ ↩² ↩³
PyTorch. torch.nn.attention module and FlashAttention implementation registration APIs. PyTorch Documentation. Accessed September 2026. ↩
PyTorch. Reproducibility documentation, including scaled-dot-product attention backend differences. PyTorch Documentation. Accessed September 2026. ↩ ↩² ↩³
Woosuk Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. ↩ ↩²
Tri Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. ↩ ↩²
Ashish Vaswani et al. Attention Is All You Need. NeurIPS 2017. ↩
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv, 2018. ↩
Markus N. Rabe and Charles Staats. Self-attention Does Not Need O(n²) Memory. arXiv, 2021. ↩

FlashAttention: Part 1

Sun, 13 Sep 2026 00:00:00 +0000

This is Part 1 of a two-part FlashAttention series. Check out Part 2 next. Part 2 covers the FlashAttention evolution (FA-1/2/3/4).

How IO-Aware Attention Makes Transformers Faster Without Approximating Attention

The reference handbook for this blog post was inspired by this tweet and is titled: Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention.

The mechanism, in three ideas: Tiling + Online Softmax + Recomputation. Everything in this handbook¹ is elaboration on that summary. The handbook is a technical deep dive on exact tiled attention: GPU memory traffic, online softmax, forward and backward passes, IO complexity, the evolution from FlashAttention-1 through FlashAttention-4, and current framework behavior.

In this blog post, we cover the fundamental problem FlashAttention-1 (FA1) addresses, the mathematical tricks it uses, the GPU implementation of those tricks, and the architectural compatibility (MHA/MQA/GQA) of FlashAttention.

0. Introduction

0.1 How to Read This Handbook
0.2 Notation

1. The Fundamental Problem

1.1 What FlashAttention actually optimizes
1.2 The attention equation is not the implementation
1.3 GPU memory hierarchy and why IO matters
1.4 Why materializing the score and probability matrices is expensive
1.5 Dense arithmetic is still quadratic
1.6 Memory-efficient exact attention predates FlashAttention

2. The Mathematical Trick

2.1 Tiling queries, keys, and values
2.2 Softmax is the difficult part of streaming
2.3 Online softmax from first principles
2.4 The blockwise merge recurrence
2.5 A complete numerical example

3. Putting the Mathematics onto the GPU

3.1 FlashAttention forward pass
3.2 Why dense FlashAttention is exact
3.3 IO complexity: what the theorem actually says
3.4 Memory complexity: linear auxiliary state, not linear compute
3.5 Causal masking and tile skipping
3.6 Backward pass: recompute instead of save
3.7 Why more FLOPs can still be faster

4. Architectural Compatibility

4.1 MHA, MQA, and GQA compatibility
4.2 Variable lengths, local attention, and dropout

5. Summary

Appendix

Reference Map
References
Citation

0. Introduction

0.1 How to Read This Handbook

What stands out to me is the number of round trips to HBM in naive standard attention. Every intermediate value $(S$, $P$, $O)$ has to be written out and read back. That is the problem FlashAttention solves. The handbook frames the subject as easiest to understand when three different questions are kept separate:

What mathematical function is being computed? For dense attention, the target remains ordinary scaled dot-product attention.
How much arithmetic does that function require? Dense all-pairs query-key scoring remains quadratic in sequence length.
How does the implementation move data through the GPU memory hierarchy? This is where FlashAttention changes the execution of the algorithm dramatically.

The lesson I keep coming back to: FLOP count alone doesn’t decide wall-clock speed. An algorithm can do essentially the same math, or even recompute intermediates from scratch, and still finish faster because it shuttles far less data to and from high-bandwidth memory (HBM).

Core distinction

Dense FlashAttention is an exact attention algorithm: it does not replace softmax attention with a low-rank, sparse, kernelized, or approximate formula. “Exact” refers to the mathematical attention computation. Floating-point kernels can still differ by small rounding effects because operations are reordered.

The word exact is doing real work here. Exactness is a statement about the mathematical function, not about bitwise reproducibility. The kernel is free to reorder floating-point operations. It is not free to change the function being computed.² $^,$ ³ $^,$ ⁴

0.2 Notation

For one attention head, let

\[Q \in \mathbb{R}^{N_q \times d}, \quad K \in \mathbb{R}^{N_k \times d}, \quad V \in \mathbb{R}^{N_k \times d_v}.\]

For self-attention, typically $N_q = N_k = N$. The scaled score matrix is

\[S = \frac{QK^T}{\sqrt{d}} + B,\]

where $B$ represents an optional additive mask or bias, and

\[P = \mathrm{softmax}_{\mathrm{row}}(S), \quad O = PV.\]

Throughout, HBM refers to large off-chip high-bandwidth GPU memory. On-chip memory is a broad teaching term for much smaller, faster storage such as registers and shared memory/SRAM. Exact hardware details vary by GPU generation.

The contrast in one table:

HBM	SRAM
slower	faster
large	smaller
off-chip	on-chip

System Problem

Where do all those intermediate values $(S, P, O)$ live while the GPU computes them?

That question, not the arithmetic, is what FlashAttention was built to answer.

1. The Fundamental Problem

1.1 What FlashAttention Actually Optimizes

Start with the ordinary attention function:

\[O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.\]

A textbook implementation boils this down to three big steps:

form the score matrix $S$,
apply row-wise softmax to get $P$, then
multiply by $V$.

Mathematically, that’s perfectly fine. On a GPU, though, writing a giant intermediate to HBM and reading it back is often far more expensive than the equation lets on.

FlashAttention’s real contribution is an IO-aware implementation. It splits the computation into tiles that fit in fast on-chip memory, streams blocks of $K$ and $V$ past them, and keeps just enough row-wise softmax state around to produce the exact output. No full $N \times N$ attention matrix ever needs to land in HBM.²

My working definition of IO-awareness:

minimize data movement between the different levels of GPU memory (HBM and SRAM), rather than just trying to reduce the number of mathematical operations (FLOPs).

The reasoning is that the speed bottleneck in modern AI hardware is often not how fast the GPU can compute math, but how fast it can read and write data. This is the memory-compute tradeoff.

The problem, in one line:

compute → HBM write → HBM read → compute

The arithmetic often isn’t the bottleneck. The round trips are. Tiling is the fix. Instead of one big $QK^T$, process it in blocks:

\[Q K^T \rightarrow Q_i K_j^T \qquad (\text{for small blocks } Q_i, K_j)\]

Each score block gets:

computed,
softmax-processed,
immediately multiplied by its corresponding $V_j$, and
discarded.

For $N = 8192$, instead of an $8192 \times 8192$ intermediate matrix, you process blocks like $128 \times 128$. All the $QK$ interactions still happen, just in manageable pieces.

Important Fact

What changes: the execution schedule, memory traffic, and stored intermediates.
What does not change: the dense scaled-dot-product attention function being evaluated.

This framing matters. Calling FlashAttention “a faster kind of attention” misses the point. It’s better understood as an algorithm and kernel family for evaluating attention efficiently on accelerators. A model can use causal masking, RoPE, MQA/GQA, or any other attention feature and still run a FlashAttention kernel underneath it.

The original paper draws a line between this and approximate attention methods, which cut arithmetic by changing the math itself. Dense FlashAttention doesn’t do that. It keeps the math intact. The block-sparse extension from the same paper is a different story, since skipping blocks literally changes which query-key interactions get computed.²

1.2 The Attention Equation Is Not the Implementation

The equation does not tell you where tensors live. Let’s look at a standard naive attention implementation:

Calculate $S$, store $S$.
Read $S$, calculate $P$, store $P$.
Read $P$, calculate $O$.

The problem is the number of HBM round trips. A simple materializing implementation conceptually does:

\[S \leftarrow QK^T / \sqrt{d}, \quad P \leftarrow \mathrm{softmax}(S), \quad O \leftarrow PV.\]

Prior to FlashAttention, here is what the standard attention implementation looks like in more detail in terms of memory communication. Load $Q, K, V \in \mathbb{R}^{N \times d}$ from HBM, then:

Read $Q, K$ from HBM, compute $S$, write $S$ to HBM.
Read $S$ from HBM, compute $P$, write $P$ to HBM.
Read $P, V$ by blocks from HBM, compute $O$, write $O$ to HBM.
Return $O$.

The GPU spends significant time moving an $N^2$ object through memory. FlashAttention avoids repeatedly moving huge intermediates:

Calculate a small $S$ tile, apply softmax to the tile, use the tile with $V_j$, and discard the tile.
Calculate the next tile, and repeat step 1 until all tiles are completed.

The key questions I ask when looking at the attention equation:

How many operations are required?
What data must move between memory levels to perform those operations?

FlashAttention’s move is simple: keep the work close. Blocks of $Q$, $K$, and $V$ sit near the compute units, score tiles get built and consumed on the spot, and only the running row statistics plus the output survive across tiles. Nothing else needs to stick around.

Algorithmic lesson

A computational graph is not a memory schedule. Writing $P = \mathrm{softmax}(QK^T)$ on paper does not require an implementation to store all of $QK^T$ or $P$ in off-chip memory at once.

The pattern isn’t unique to attention. Fused kernels, tiling, recomputation, and operator scheduling all do the same thing: spend a little extra arithmetic to avoid dragging huge intermediates through memory. On modern hardware, that’s usually a winning bet. Matmul throughput has raced ahead of the rest of the memory hierarchy, so the math is cheap and the data movement is what hurts. FlashAttention-2 (FA2), FlashAttention-3 (FA3), and FlashAttention-4 (FA4) take this further, and they’re upfront about it: the algorithm bends to the hardware, not the other way around.⁵ $^,$ ⁶ $^,$ ⁷

1.3 GPU Memory Hierarchy and Why IO Matters

GPUs expose a hierarchy rather than one uniform pool of equally fast memory. The names and capacities vary by architecture, but the mental model is:

flowchart BT HBM[HBM/device memory] SRAM[on-chip shared memory/SRAM] Registers[Registers/very local state] HBM -->|load tiles| SRAM SRAM -->|feed compute| Registers classDef memoryDevice stroke:#818cf8,fill:#eef2ff classDef memorySRAM stroke:#2dd4bf,fill:#f0fdfa classDef memoryRegisters stroke:#a78bfa,fill:#f5f3ff class HBM memoryDevice class SRAM memorySRAM class Registers memoryRegisters

Figure 1: GPU Memory Hierarchy.

FA1 models the HBM/SRAM asymmetry and explicitly optimizes the number of transfers between them.² My breakdown of each GPU memory level in the figure above:

1. HBM:

Relatively large
Relatively slower to access
Physically farther from the compute units, so data must travel further to reach them
Stores model weights, $Q/K/V$, activations, and other large tensors

2. SRAM / shared memory:

Much smaller
Faster
Cheaper reuse cost

3. Registers:

Even smaller
More local

Why tiling helps: Suppose $Q_i$ (a query tile/block) needs to interact with many $K/V$ tiles/blocks prior to writing its output state. Instead of constantly moving $Q_i$ back and forth, we can keep it close to the compute units while processing:

\[K_1, V_1 \to K_2, V_2 \to K_3, V_3 \to \cdots\]

Therefore, one loaded $Q_i$ can participate in lots of computation. This is called reuse. Tiling doesn’t only make tensors smaller, it also increases reuse while a tile is resident on chip. Conversely, $K/V$ blocks can be streamed through query blocks according to the chosen schedule.

IO-Awareness

A kernel becomes IO-aware when the placement and movement of data are part of the algorithm, rather than an afterthought left to a sequence of separately launched tensor operations.

This does not mean attention is “always memory-bound.” FA3 and FA4 exist partly because, as hardware changed, the dominant bottlenecks changed too.⁶ $^,$ ⁷ The bottleneck depends on:

Sequence length, $N$
Head dimension, $d$
dtype
Mask pattern
GPU generation (Ampere/Hopper/Blackwell)
Forward vs. backward pass
Which kernel is running

1.4 Why Materializing the Score and Probability Matrices Is Expensive

The quadratic intermediate becomes concrete very quickly. Suppose a batch contains one sequence, with 32 attention heads, sequence length $N = 8192$, and a two-byte dtype such as FP16 or BF16. Concretely, my arithmetic:

\[\begin{aligned} \text{batch} &= 1 \\ \text{heads} &= 32 \\ N &= 8192 \\ \text{dtype } &= \text{FP16/BF16} \\ \text{tensor shape } &= [1, 32, 8192, 8192] \\ 2 & \text{ bytes/element} \\ \end{aligned}\]

At two bytes per element, the total memory is: $32 \times 8192^2 \times 2 = 4{,}294{,}967{,}296 \text{ bytes} = 4 \text{ GiB}.$

One dense tensor with shape $[1, 32, 8192, 8192]$ contains $32 \times 8192^2$ elements. At two bytes per element, that is $4 \text{ GiB}$. That is the size of one full score- or probability-like tensor for this example. Imagine reading and writing this much data.

Naive pipeline:

A naive decomposition can produce arrays of this size at several stages. That’s not to say every framework keeps $S$ and $P$ alive at once, or that compiler fusion hasn’t already cut some of the traffic. The example is there to make the problem concrete: a dense $N^2$ intermediate is big enough that rewriting and rereading it can dominate memory and bandwidth.

flowchart LR QKT["QKᵀ"] --> S["huge S"] S --> Softmax["softmax"] Softmax --> P["huge P"] P --> PV["PV"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:22px classDef emphasized fill:none,stroke:none,font-size:26px,font-weight:bold class QKT,Softmax,PV boxed class S,P emphasized

Figure 2: Naive Attention Implementation Pipeline.

FlashAttention pipeline:

FlashAttention skips the naive decomposition entirely. Score tiles get formed, softmaxed on chip, used to accumulate their contribution to the output, then thrown away.² The full matrix never touches HBM.

flowchart LR QK["QᵢKⱼᵀ"] --> Softmax["softmax tile"] Softmax --> PV["tile × Vⱼ"] PV --> Discard["discard tile"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:22px classDef plain fill:none,stroke:none,font-size:26px,font-weight:bold class QK,Softmax,PV boxed class Discard plain

Figure 3: FlashAttention Implementation Pipeline.

Important wording

FlashAttention removes the need to materialize the full attention matrix as an off-chip intermediate. It does not remove the logical pairwise interactions required by dense attention.

This matters most during training, where naive autograd would want those large intermediates saved for the backward pass.

Important distinction: FlashAttention acknowledges that $N^2$ interactions exist and says: Do not store the entire result of these interactions as a giant intermediate if we can consume each piece immediately.

1.5 Dense Arithmetic Is Still Quadratic

FlashAttention changes the memory schedule, not the math. The function itself is still dense attention. That means forming all query-key scores for $N$ tokens with head dimension $d$ still costs work proportional to $N^2 d$. There’s no way around that.

Multiplying probabilities by values is another dense pairwise matrix multiplication of the same broad order. FlashAttention reshuffles these operations, but it doesn’t skip the dense set of query-key interactions. Those still happen.² $^,$ ⁵

So three things must not get conflated:

Arithmetic complexity: dense attention stays at $O(N^2 d)$.
Large intermediate storage: FlashAttention avoids an $O(N^2)$ materialized score/probability tensor in HBM.
HBM traffic: the original paper proves a lower IO cost than standard materializing attention under its two-level memory model.²

This is how a method makes much longer sequences practical without making long context “free.” Double $N$, and you roughly quadruple the number of dense query-key pairs. FlashAttention just attacks the data-movement and memory-footprint side of that.

Important Fact

If you actually want to reduce the number of query-key pairs, you need a different mathematical structure: sparsity, a local pattern, or a different attention formulation. Those choices change model behavior, and they’re conceptually separate from dense FlashAttention.

When you report speedups, always separate asymptotic arithmetic from measured runtime. A kernel can run several times faster at the same $O(N^2 d)$ complexity because the constant factors, occupancy, fusion, and memory traffic change dramatically.

Essentially, the key takeaways are summarized below:

Property	Dense FlashAttention
Arithmetic	$O(N^2 d)$
Full attention intermediates	Avoids $O(N^2)$ storage
HBM traffic	Reduced substantially

1.6 Memory-Efficient Exact Attention Predates FlashAttention

It would be historically inaccurate to say FlashAttention first discovered that exact attention can avoid quadratic memory. Earlier work by Rabe and Staats demonstrated exact memory-efficient attention with quadratic computation but subquadratic memory. They showed that attention does not require $O(N^2)$ memory with respect to $N$.⁸ Milakov and Gimelshein also did earlier work on online softmax, showing a memory-efficient, online recurrent way of computing the classical stable softmax normalizer.⁹ FlashAttention builds the same kind of running-max/running-normalizer idea into tiled attention, while also accumulating the value-weighted output.²

The lineage from stable softmax to FlashAttention is shown in Figure 4 below:

flowchart TD Stable["Stable softmax"] --> Online["Online normalization"] Online --> Exact["Exact memory-efficient attention"] Exact --> Flash["FlashAttention"] classDef boxed fill:#e3f2fd,stroke:#1565c0,stroke-width:2px,font-size:16px class Stable,Online,Exact,Flash boxed

Figure 4: Memory-Efficient Exact Attention Lineage.

FlashAttention’s key contribution was to bring together the techniques below into a practical, high-performance algorithm:

Tiling
Online softmax
Fused computation
GPU memory hierarchy awareness
IO complexity analysis

This distinction matters because “memory-efficient” does not automatically mean “IO-optimized for a particular hardware model.”

2. The Mathematical Trick

2.1 Tiling Queries, Keys, and Values

Most important. This is the section where the trick lives. Score tiles are small enough to stay near the compute units. Each tile’s contribution gets folded into the running row statistics and the output accumulator, then the tile is thrown away.

Suppose $Q \in \mathbb{R}^{N_q \times d}$, $K \in \mathbb{R}^{N_k \times d}$, and $V \in \mathbb{R}^{N_k \times d_v}$. Instead of processing all $N_q, N_k$ rows at once, divide them into blocks. For example:
\[Q: \boxed{Q_1} \boxed{Q_2} \boxed{Q_3} \quad K, V: \boxed{K_1, V_1} \boxed{K_2, V_2} \boxed{K_3, V_3}\]
For one query-key/value pair ($Q_i, K_j/V_j$), we process that score tile $S_{ij}$ locally:
\[S_{ij} \doteq \frac{Q_i K_j^T}{\sqrt{d}} + B_{ij},\]
Then, for each query block:
```
  1. Calculate score tile, S_ij
  2. Row-wise softmax and update
  3. Multiply by V_j and accumulate
  4. Update the running output
  5. Discard the tile 
  6. Move to the next tile, never writing a global S or P matrix.
```
Recall that softmax is defined as
\[\mathrm{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}.\]
Matrix multiplication is easy to tile because it is fundamentally accumulation:
\[AB = \sum_j A_j B_j,\]
so we can calculate pieces and add them.
Softmax is harder because every element depends on the entire row. If we process the first block, we do not know the eventual denominator. Even worse, numerical stability requires knowing the maximum.
In the kernel, we choose tile shapes and loop order based on on-chip capacity, head dimension, GPU generation, causal structure, and work partitioning. FA1 uses tile sizes derived from SRAM capacity $M$.²

Mathematical Trick

Tiling alone is not enough. Matrix multiplication tiles compose naturally because sums can be accumulated. Softmax couples every score in a row through a shared maximum and denominator, so we need a way to merge blocks without seeing the full row at once.

2.2 Softmax Is the Difficult Part of Streaming

A numerically stable softmax for one row $x_1, \ldots, x_N$ uses

\[m = \max_j x_j, \qquad \ell = \sum_j e^{x_j - m}, \qquad p_j = \frac{e^{x_j - m}}{\ell}.\]

Subtracting $m$ keeps large logits from blowing up. However, it seems to create a streaming problem: how does one normalize an early block if a later block turns out to have a bigger maximum?

The answer is to keep running statistics and rescale them when the maximum changes. Say the state after the earlier elements is $(m_{\text{old}}, \ell_{\text{old}})$, and the new block’s maximum is $m_b$. Then

\[m_{\text{new}} = \max(m_{\text{old}}, m_b)\]

Anything accumulated under the old max converts to the new one by a single factor:

\[\alpha \doteq e^{m_{\text{old}} - m_{\text{new}}}\]

The new block is then evaluated against the same $m_{\text{new}}$.

None of this is an approximation. It’s just the identity

\[e^{x - m_{\text{old}}}\, e^{m_{\text{old}} - m_{\text{new}}} = e^{x - m_{\text{new}}}\]

Mathematical Insight

The running maximum is a change of numerical reference point. When that reference changes, previously accumulated exponentials can be rescaled exactly in real arithmetic rather than recomputed from scratch.

The online-normalizer recurrence predates FlashAttention and is the mathematical basis for streaming stable softmax.⁹

Worked example 1: why subtract the max

Start with two extreme values to see why the max subtraction matters:

\[\begin{aligned} x &= [1000, 999], \quad \text{Directly computing $e^{1000}$ overflows}\\ m &= 1000 \\ x - m &= [1000 - 1000, \; 999 - 1000] = [0, -1] \\ e^{[0, -1]} &\approx [1, 0.368] \\ \end{aligned}\]

The exponential values $1$ and $0.368$ are perfectly manageable.

Worked example 2: streaming two blocks

Now let’s work through streaming two blocks.

\[\begin{aligned} \text{Block 1}: [2, 1]&, \text{ Block 2}:[4, 3]\\ \text{After Block 1} &: m_{\text{old}} = 2 \\ \text{Evaluating Block 1 at } m_{\text{old}} &: [e^{0}, e^{-1}] \\ \text{When Block 2 arrives} &: m_b = 4 \\ \text{Now in Block 2} &: m_{\text{new}} = \max(2, 4) = 4, \quad \alpha = e^{2 - 4} = e^{-2}\\ \text{Rescaling Block 1 at the new reference} &: [\alpha e^{0}, \; \alpha e^{-1}] = [e^{-2}, \; e^{-3}] \\ \text{Block 2 at the new reference} &: [e^{4 - 4}, \; e^{3 - 4}] = [e^{0}, \; e^{-1}] \\ \end{aligned}\]

Putting it all together, with $x = [2, 1, 4, 3]$:

\[\begin{aligned} \mathrm{softmax}(x) &= \frac{[e^{-2}, \; e^{-3}, \; e^{0}, \; e^{-1}]}{e^{-2} + e^{-3} + e^{0} + e^{-1}}\\ \text{Block 1 streaming computation} &: \ell = \sum_j e^{x_j - m} = e^{0} + e^{-1} = 1 + e^{-1}, \\ \text{Block 1 normalized output} &: \frac{[e^{0}, e^{-1}]}{(1 + e^{-1})} \\ \text{Block 2 running normalizer} &: \ell_{\text{new}} = \alpha\, \ell_{\text{old}} + \sum_{j \in \text{Block 2}} e^{x_j - m_{\text{new}}} = e^{-2}(1 + e^{-1}) + e^{0} + e^{-1} \\ \text{Block 2 unnormalized accumulator} &: a = [\alpha e^{0}, \alpha e^{-1}, e^{0}, e^{-1}] = [e^{-2}, e^{-3}, e^{0}, e^{-1}] \\ \text{Combined result} &: \frac{[e^{-2}, e^{-3}, e^{0}, e^{-1}]}{(e^{-2} + e^{-3} + 1 + e^{-1})} \\ \end{aligned}\]

This matches the full softmax computed in one shot.

2.3 Online Softmax from First Principles

Process scalar logits $x_1, x_2, \ldots$ one at a time. Initialize

\[m_0 = -\infty, \qquad \ell_0 = 0\]

After observing $x_j$:

\[m_j = \max(m_{j-1}, x_j), \qquad \ell_j = \ell_{j-1} e^{m_{j-1} - m_j} + e^{x_j - m_j}\]

Milakov and Gimelshein showed this produces the same stable softmax normalizer as the conventional safe-softmax procedure, with fewer passes over the input.⁹

Attention needs a weighted value sum, so we also maintain a value-weighted, unnormalized accumulator,

\[a_j = \sum_{i \leq j} e^{x_i - m_j} v_i.\]

When the maximum moves from $m_{j-1}$ to $m_j$, the accumulator is rescaled exactly like the normalizer before the new value contribution is folded in:

\[a_j = a_{j-1}\, e^{m_{j-1} - m_j} + e^{x_j - m_j} v_j.\]

After the last element, $x_n$, the output is

$o = \frac{a_n}{\ell_n}.$

Important Fact

For attention, $x_j$ is not a fixed input vector stored in advance. Each block of logits is generated on demand from a matrix product $QK_j^T / \sqrt{d}$ plus mask/bias terms. The online recurrence lets the kernel consume that block immediately.

The same idea works row by row and block by block, which is what makes a tiled exact softmax-attention forward pass possible.

Key Insight. We don’t need the entire probability vector. We only need enough information to reconstruct its contribution to the final output.

2.4 The Blockwise Merge Recurrence

For one query row, suppose the running state after some key blocks is

\[(m, \ell, a),\]

where $m$ is the running maximum score, $\ell$ is the stable softmax normalizer under that maximum, and $a \in \mathbb{R}^{d_v}$ is the unnormalized value accumulator.

A new block lands with scores $s \in \mathbb{R}^b$ and values $V_b \in \mathbb{R}^{b \times d_v}$. First, find the block’s own maximum and update the running one:

\[m_b = \max(s), \qquad m' = \max(m, m_b).\]

The rescale factor and the block’s exponentials follow:

\[\alpha = e^{m - m'}, \qquad p = e^{s - m'}.\]

Then update:

\[\ell' = \alpha \ell + \sum_j p_j, \qquad a' = \alpha a + p^T V_b,\] \[m \leftarrow m', \qquad \ell \leftarrow \ell', \qquad a \leftarrow a'.\]

At the end:

\[o = a / \ell.\]

When you process a block of query rows rather than a single row, the running state changes shape. $m$ and $\ell$ each become vectors with one entry per row, and $a$ becomes a matrix. If you have masks, apply them to the score tile before the exponentials. Masked positions contribute zero probability, so they drop out of the update entirely. FA1 and FA2 express equivalent running-max/running-normalizer/output updates in block form.² $^,$ ⁵

Important Fact

The recurrence is the algebraic reason tile boundaries do not change the dense softmax result. A different tiling changes the order of floating-point operations, but not the intended real-arithmetic function.

This recurrence is what makes tiled exact attention possible.

2.5 A Complete Numerical Example

Take one already-scaled, unmasked attention row

\[s = [2, 1, 4, 3]\]

with two-dimensional values

\[v_1 = [1, 0], \quad v_2 = [0, 1], \quad v_3 = [2, 0], \quad v_4 = [0, 2].\]

The global maximum is 4, so the stable unnormalized weights are

\[[e^{-2}, e^{-3}, e^{0}, e^{-1}] \approx [0.135335, 0.049787, 1, 0.367879].\]

Their sum is

\[\ell \approx 1.55300179,\]

and full softmax gives

\[O \approx [1.37497284, 0.50582424].\]

Now run it through in two blocks, the way a tiled kernel would: $[2, 1]$ and $[4, 3]$.

Block 1: $s = [2, 1]$. Running max $m_1 = 2$, stable exponentials $[e^{0}, e^{-1}] = [1, e^{-1}]$. So

\[\ell_1 = 1 + e^{-1} \approx 1.36788,\]

and the value accumulator comes out to

\[a_1 = 1 \cdot [1, 0] + e^{-1} \cdot [0, 1] = [1, e^{-1}].\]

Block 2: $s = [4, 3]$. The block max is $m_b = 4$, beating the running max, so $m’ = \max(2, 4) = 4$. The rescale factor:

\[\alpha = e^{2 - 4} = e^{-2}.\]

Then

\[\ell_2 = \alpha \ell_1 + e^{0} + e^{-1} = e^{-2}(1 + e^{-1}) + 1 + e^{-1} = e^{-2} + e^{-3} + 1 + e^{-1} \approx 1.553,\]

and

\[a_2 = e^{-2} \cdot [1, e^{-1}] + [2 e^{0}, 2 e^{-1}] = [2 + e^{-2}, 2 e^{-1} + e^{-3}] \approx [2.135, 0.7855].\]

Output:

\[O_2 = a_2 / \ell_2 \approx [1.375, 0.506],\]

matching the full-row computation.

Important Fact

The old block was not revisited. Its contribution was merely rescaled when a larger maximum appeared.

This rescale-instead-of-recompute is what makes streaming possible.

3. Putting the Mathematics onto the GPU

3.1 FlashAttention Forward Pass

The forward pass, conceptually:

Partition $Q$ into query-row tiles, and $K$ and $V$ into key/value tiles.
For each query tile, initialize row-wise running maxima, normalizers, and output accumulators: $m = -\infty$, $\ell = 0$, $a = 0$.
Load a $K_j, V_j$ tile and form the local score tile $S_{ij} = Q_i K_j^T / \sqrt{d}$.
Apply the masks or biases $B_{ij}$ belonging to this tile.
Find the tile maximum, update the running maximum $\max(m, m_b)$, and rescale the previous state.
Exponentiate the current tile relative to the updated maximum.
Update the normalizer $\ell$ and the value-weighted accumulator $a$.
Repeat for every required key tile.
Normalize the accumulator row-wise, $O = a / \ell$, then write the output.

The original algorithm picks block sizes so the tiles and running state fit in on-chip SRAM, cutting trips to HBM.² A quick summary of the forward pass is shown below:

for each K/V block:
    scores = Q @ K_block.T / sqrt(d)
    update the running max m and normalizer l
    update the output accumulator a
O = a / l   # normalize once after the last block

You can reproduce the algebra in a few lines of PyTorch, but that’s not the same as a high-performance FlashAttention kernel. Real production implementations additionally exploit GPU-specific tiling, registers, shared memory, warp scheduling, tensor cores, async copies, specialized intrinsics, and occupancy optimization.

Important Fact

The essential algorithmic idea is independent of one CUDA kernel: generate a score tile, consume it immediately through online softmax and $V$ accumulation, and never materialize the complete score/probability matrix in HBM.

Minimal educational PyTorch.

import math
import torch

# Single query block, tiled over K/V. Production kernels also
# tile Q; the outer loop is omitted here to focus on the recurrence.
def tiled_attention(q, k, v, block=128):
    scale = 1 / math.sqrt(q.shape[-1])
    m = q.new_full((q.shape[0],), -float("inf"))
    l = q.new_zeros(q.shape[0])
    a = q.new_zeros((q.shape[0], v.shape[-1]))

    for j in range(0, k.shape[0], block):
        kj, vj = k[j:j+block], v[j:j+block]
        s = (q @ kj.T) * scale
        m_new = torch.maximum(m, s.max(dim=-1).values)
        alpha = torch.exp(m - m_new)
        p = torch.exp(s - m_new[:, None])
        a = a * alpha[:, None] + p @ vj
        l = l * alpha + p.sum(dim=-1)
        m = m_new

    return a / l[:, None]

The gap between this educational code and a production kernel matters a lot for FA2, FA3, and FA4. Every generation keeps the same semantic structure, but schedules the work differently for newer hardware.⁵ $^,$ ⁶ $^,$ ⁷ Same math, different schedule, each tuned to the hardware of its generation.

3.2 Why Dense FlashAttention Is Exact

Does FlashAttention approximate attention? For dense FlashAttention, NO. The target function is untouched:

\[\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.\]

A short checklist for why this counts as exact:

No Q-K pairs are intentionally removed.
No low-rank approximation is introduced.
No alternative kernelized attention function replaces softmax.
Tiling just processes the same interactions in blocks.

Online softmax doesn’t swap the exponential or normalization for something else either. The block recurrence only changes the order in which sufficient statistics get accumulated.

But exact needs a qualification. Suppose two implementations compute $a + b + c$:

A: $(a + b) + c$
B: $a + (b + c)$

In exact mathematics,

\[(a + b) + c = a + (b + c).\]

In floating point, for some values,

\[(a + b) + c \neq a + (b + c).\]

Tiling changes the reduction order. Fusion can change how rounding shakes out. So:

\[\text{exact algorithm} \neq \text{bitwise identical output}.\]

Three caveats keep the word exact precise:

Floating point has finite precision. Reorder the additions and reductions and you’ll get small numerical differences from another implementation. PyTorch says so directly: the output of SDPA can differ depending on which backend kernel is chosen, because each backend accumulates floating-point values in a different order.³ $^,$ ⁴
Dropout is random during training. If you want two runs to match bit-for-bit, you have to line up the random behavior too. Attention semantics alone won’t be enough.
Sparse variants aren’t dense. The block-sparse variant in the original paper skips blocks entirely, so it approximates full dense attention rather than reproducing it.²

Important Fact

“Exact” doesn’t mean “bitwise identical to every reference kernel.” It means the algorithm isn’t quietly changing dense softmax attention to save work.

This distinction becomes important in numerical testing. A sensible tolerance depends on dtype, accumulation order, sequence length, and backend, rather than requiring binary identity.

3.3 IO Complexity: What the Theorem Actually Says

The original FlashAttention paper works in a two-level memory model: HBM plus on-chip SRAM of size $M$. Under its assumptions, which include head dimension $d$ and an SRAM size in the range $d \leq M \leq Nd$, the two attention implementations require:

Standard materializing attention: $\Theta(Nd + N^2)$ HBM accesses
FlashAttention: $\Theta\left(\frac{N^2 d^2}{M}\right)$ HBM accesses

The second quantity is a communication-complexity bound, and a few things are worth keeping straight:

It holds for a specific memory model, not for every GPU.
It counts scalar elements moved between the modeled memory levels, up to asymptotic factors. It is not a byte count.
It is not the arithmetic complexity. Dense attention still performs $O(N^2 d)$ work.
More usable on-chip memory $M$ means more reuse and less modeled HBM traffic.

The intuition: bigger on-chip memory lets larger working tiles stay resident. Bigger tiles mean more reuse and fewer HBM reloads. In the theorem’s regime, that reuse is exactly why the HBM-access bound shrinks as $M$ grows. It’s also why GPU architecture matters so much to FlashAttention performance. For typical values ($d$ of 64 to 128, $M$ around 100KB), $d^2$ is many times smaller than $M$, so FlashAttention needs many times fewer HBM accesses.²

Important Fact

A common incorrect summary is “FlashAttention reduces attention IO from $O(N^2)$ to $O(N)$.” The linear quantity is the large auxiliary memory footprint with respect to sequence length, not the general HBM-access expression above.

3.4 Memory Complexity: Linear Auxiliary State, Not Linear Compute

FlashAttention papers often say memory becomes linear instead of quadratic. That’s true, but it’s easy to misread, so it’s worth being precise about what’s linear and what isn’t:

\[\text{linear auxiliary state} \neq \text{linear compute.}\]

A naive attention implementation creates $O(N^2)$ intermediates. FlashAttention doesn’t. It keeps only:

$Q/K/V$ tiles
row-wise $m$
row-wise $\ell$
the output accumulator $a$

The inputs and output already occupy $O(Nd)$ memory on their own. A materializing dense attention implementation additionally creates $O(N^2)$ score and probability state. FlashAttention skips those full matrices and holds only tile-sized working state on chip, plus $O(N)$ row-wise statistics. So the extra memory tied to the attention operation grows linearly with sequence length rather than quadratically. But linear memory is not linear compute. The computation is still $O(N^2 d)$.

Example. For $N = 10{,}000$, there are roughly $10{,}000^2 = 100{,}000{,}000$ Q-K interactions. FlashAttention doesn’t make those disappear. It just prevents you from needing a giant intermediate holding all of them at once.

Training is where this difference bites hardest. A straightforward backward pass would otherwise want the full probability matrix $P$ kept around. FlashAttention stores compact information instead, the output plus row-wise normalization statistics, and recomputes score and probability tiles during the backward pass.

Important Fact

Linear memory does not imply linear runtime. The number of dense query-key interactions is still quadratic in $N$.

Two more things to keep in mind. First, don’t claim that total model memory is $O(N)$. Other Transformer components consume activation memory, and autoregressive serving has KV-cache memory that grows with context length. FlashAttention changes how the attention computation manages its intermediates, not the whole model’s memory profile. Second, FlashAttention and activation checkpointing can coexist, since both trade recomputation for reduced stored state, just at different scopes of the training graph.

3.5 Causal Masking and Tile Skipping

In autoregressive attention, query token $i$ can’t attend to future key token $j > i$. The attention matrix looks like

\[\begin{bmatrix} \checkmark & \times & \times & \times \\ \checkmark & \checkmark & \times & \times \\ \checkmark & \checkmark & \checkmark & \times \\ \checkmark & \checkmark & \checkmark & \checkmark \end{bmatrix}\]

with mask

\[M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}\]

Apply the mask before softmax and you get

\[\mathrm{softmax}(S') = \begin{bmatrix} \text{value} & 0 & 0 & 0 \\ \text{value} & \text{value} & 0 & 0 \\ \text{value} & \text{value} & \text{value} & 0 \\ \text{value} & \text{value} & \text{value} & \text{value} \end{bmatrix}\]

A materialized mask is another $N \times N$ object. But a tiled kernel can reason about the geometry of each tile directly:

Blocks strictly above the causal boundary are fully masked, so they can be skipped.
Blocks strictly below the boundary are fully valid.
Only blocks intersecting the diagonal need element-level causal masking.

That’s a lot of invalid tiles you never have to compute. FA2 explicitly exploits this structure,⁵ and current implementations also define alignment rules for unequal query/key lengths. The FlashAttention interface aligns the causal mask to the bottom-right corner of the attention matrix, whereas PyTorch’s is_causal=True assumes upper-left alignment. The two only differ when $N_q \neq N_k$, so check which convention your code path uses.¹⁰ $^,$ ³

Important Fact

Mask semantics are an API detail that must not be guessed. For example, current PyTorch SDPA treats a Boolean attn_mask value of True as a position that participates in attention, while other PyTorch mask APIs use different conventions.

Connection to local attention. Sliding-window attention works on the same principle. If each token only attends to 128 nearby tokens, most tiles can be skipped the same way. But this is a different move: you’ve changed the mathematical attention pattern itself. The model is now doing sparse or local attention, and FlashAttention is just the kernel executing that pattern. It’s no longer the same problem as unrestricted dense attention.

3.6 Backward Pass: Recompute Instead of Save

A naive training implementation saves $P = \mathrm{softmax}(S)$ for the backward pass. That’s an $N^2$ tensor sitting in HBM. FlashAttention doesn’t keep it. It stores compact row-wise normalization information instead, then recomputes score and probability tiles when gradients are needed.

Forward:

Store the output and row-wise normalization statistics.
Do not store $P \in \mathbb{R}^{N \times N}$.

Backward:

For each tile: recompute $QK^T$, reconstruct the local probabilities, compute gradients, and discard the tile.
Same trade as before: more arithmetic, less memory traffic.

For

\[S = QK^T / \sqrt{d}, \quad P = \mathrm{softmax}(S), \quad O = PV,\]

the backward pass relies on a row-wise identity, where $D_i$ is a per-row scalar:

\[D_i = \sum_r dO_{ir}\, O_{ir} = \sum_j P_{ij}\, dP_{ij}\]

After recomputing a tile of $P$, the local derivatives fall out as

\[dV \mathrel{+}= P^T dO, \qquad dP = dO\, V^T, \qquad dS = P \odot (dP - D_i[\text{:}, \text{None}]),\] \[dS_{ij} = P_{ij}\,(dP_{ij} - D_i),\]

and then

\[dQ \mathrel{+}= dS\, K / \sqrt{d}, \qquad dK \mathrel{+}= dS^T\, Q / \sqrt{d}.\]

Masks contribute zero probability and zero gradient. FA2 stores a row-wise log-sum-exp quantity so the probability tile can be reconstructed stably from recomputed scores.⁵

Important Fact

Backward recomputation is intentional. It spends extra matrix-multiply work to avoid reading and writing a giant probability tensor, which is a favorable trade on GPUs.

This is the memory-compute tradeoff applied to the backward pass.

3.7 Why More FLOPs Can Still Be Faster

This is one of the biggest lessons from FlashAttention. The usual assumption is that fewer FLOPs means faster. On a GPU, that’s not always true.

Different operations run at completely different throughputs. Tensor cores chew through matrix multiplication. Everything else is comparatively expensive: exponentials, reductions, synchronization, shared-memory operations, and HBM transfers.² $^,$ ⁵ Therefore, doing extra arithmetic can be the right move if it eliminates expensive memory traffic.

Consider two options:

A: compute, write huge $P$ to HBM, read $P$ back, compute
B: compute, discard, recompute later

B does more arithmetic, but it avoids pushing a giant tensor through HBM. On modern accelerators, that trade usually wins.

Every FlashAttention generation is built around this. FA2’s goal is partly to reduce non-matmul FLOPs, since those don’t have the same throughput as tensor-core GEMMs. It also reworks the work partitioning so more of the GPU is doing useful work at once.⁵ FA3 pushes further on Hopper, overlapping matrix multiplication, softmax, and data movement with asynchronous hardware features.⁶ FA4 responds to Blackwell, where tensor-core throughput grew faster than other resources. That shift makes exponentials and shared-memory traffic relatively more important, and FA4 is designed around it.⁷

Important Fact

A better performance question is not merely “How many FLOPs?” but “Which operations, on which units, with what data movement, reuse, parallelism, and synchronization?”

This is the systems lesson that makes FlashAttention matter beyond attention itself. Hardware efficiency doesn’t come from the math alone. It comes from co-designing the mathematical scheduling with the memory/execution hierarchy.

4. Architectural Compatibility

4.1 MHA, MQA, and GQA Compatibility (Shared K/V Heads)

FlashAttention isn’t tied to standard Multi-Head Attention (MHA). In general, the inputs have shapes

\[Q \in \mathbb{R}^{B \times N_q \times H_q \times d}, \quad K, V \in \mathbb{R}^{B \times N_k \times H_{kv} \times d}.\]

Let’s look into FlashAttention’s compatibility with MHA, Multi-Query Attention (MQA), and Grouped-Query Attention (GQA). Three head-sharing regimes are distinguished by $H_q$ versus $H_{kv}$:

MHA: Every query head has its own K/V head. $H_q = H_{kv}$
MQA: All query heads share a single K/V head. $H_{kv} = 1$
GQA: Several query heads share each K/V head. $1 < H_{kv} < H_q$

Here’s the distinction that matters:

MQA/GQA defines the architecture. It decides how heads share K/V projections.
RoPE defines the position transform. It decides how position gets baked into Q and K.
FlashAttention defines the execution. It decides how the resulting attention operation gets scheduled and computed efficiently.

Any combination of these three orthogonal concepts can be used together. A kernel doesn’t care how many query heads share a K/V head, or whether the Q/K inputs were RoPE-rotated. It just evaluates attention between each query head and its assigned K/V head. Head sharing changes the architecture and the KV-memory footprint. RoPE changes the input representation. FlashAttention changes how the work gets scheduled on the hardware.

Current implementations do impose one shape constraint. In the Dao-AILab kernels, the number of query heads must be divisible by the number of KV heads¹⁰:

\[H_q \bmod H_{kv} = 0\]

PyTorch SDPA exposes enable_gqa, but its docs have described GQA as an experimental feature with its own shape constraints, so production code should follow the exact version’s documentation rather than assuming universal fused-kernel support.³

4.2 Variable Lengths, Local Attention, and Dropout

Real systems aren’t always the clean case of same sequence length, dense attention, and no dropout.

Production implementations support a broader feature set:

variable-length sequences
causal attention
sliding-window attention (SWA)
dropout, in training-oriented interfaces
MQA/GQA
ALiBi-style score bias
KV-cache decoding, including optional RoPE handling

But these are implementation capabilities, not properties of the FlashAttention mathematical idea itself.¹⁰

What a given build supports depends on:

GPU
CUDA/ROCm backend (NVIDIA and AMD have separate implementations, with different support matrices)
dtype
head dimension
mask
library version
kernel generation

One practical warning on dropout. PyTorch’s scaled_dot_product_attention applies dropout according to its dropout_p argument regardless of training or eval mode, so eval code needs to pass 0.0 explicitly when dropout should be off.³

Important Fact

Do not infer feature support from the name “FlashAttention.” Check the exact library, kernel generation, device, dtype, head dimension, mask/bias, and training/inference path that will actually run.

5. Summary

Let’s walk through the whole FlashAttention story in one derivation, and hit the key points to remember.

Ordinary attention: $QK^T \in \mathbb{R}^{N \times N}$. Materializing it creates an enormous intermediate matrix.¹¹

\[O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V.\]

Naive execution: This creates huge memory traffic.

\[QK^T \to \text{store } N \times N \text{ scores} \to \text{softmax} \to \text{store } N \times N \text{ probabilities} \to \text{multiply by } V\]

FlashAttention: Tiling is used. Split $Q \to Q_i$ and $K, V \to K_j, V_j$, then calculate the score tile:

\[S_{ij} = \frac{Q_i K_j^T}{\sqrt{d}}.\]

The problem with tiling: Softmax needs the maximum and the denominator over the whole row.
The solution: online softmax
1. Maintain $(m, \ell, a)$ where:
  - $m$ = running maximum
  - $\ell$ = running sum of exponentials (the softmax denominator)
  - $a$ = running unnormalized value-weighted sum (accumulator)
2. When a new block has a larger maximum $m’$, calculate a rescaling factor $\alpha$ and rescale the old state ($\ell$ and $a$) by it:
\[\alpha = e^{m - m'}\]
Therefore, we can process the full FlashAttention pipeline without ever materializing the $N \times N$ matrix:
\[\text{tile} \to \text{score} \to \text{online softmax} \to V \text{ accumulation} \to \text{discard tile}.\]
Result: The same exact dense attention mathematics, but much better memory behavior: less HBM traffic, much smaller intermediates. Arithmetic remains $O(N^2 d)$.

FlashAttention in one table:

Property	Complexity / Formulation
Memory complexity	$O(Nd)$ auxiliary attention storage
Compute complexity	$O(N^2 d)$ dense attention compute
Attention formulation	Exact dense softmax attention

Key takeaway

Same exact dense attention mathematics, but much better memory behavior. Less HBM traffic, much smaller intermediates, arithmetic unchanged at $O(N^2 d)$.

Appendix

Reference Map

The table below shows which paper to refer to for each question.¹

Question	Best starting source
What is the original IO-aware algorithm?	FlashAttention-1 paper ²
Why does online normalization work?	Milakov & Gimelshein ⁹
Was exact subquadratic-memory attention known earlier?	Rabe & Staats ⁸
What specifically changed in FA2?	FlashAttention-2 paper ⁵
What is Hopper-specific in FA3?	FlashAttention-3 paper ⁶
What changes on Blackwell in FA4?	FlashAttention-4 paper ⁷
What does the current official package support?	Dao-AILab repository ¹⁰
How does current PyTorch select SDPA kernels?	PyTorch SDPA and `sdpa_kernel` docs ³ $^,$ ¹²
Why can exact backends differ numerically?	PyTorch reproducibility notes ⁴
How is PagedAttention different?	Kwon et al. ¹³

A note on how to use this map

References

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026FlashAttentionPt1,
  title   = "FlashAttention: Part 1",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Sep",
  url     = "https://chizkidd.github.io/2026/09/13/flashattention/"
}

Understanding FlashAttention: How IO-Aware Attention Makes Transformers Faster Without Approximating Attention. Google Drive PDF. Accessed September 2026. ↩ ↩²
Tri Dao et al. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹ ↩¹⁰ ↩¹¹ ↩¹² ↩¹³ ↩¹⁴ ↩¹⁵
PyTorch. torch.nn.functional.scaled_dot_product_attention documentation. PyTorch Documentation. Accessed September 2026. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶
PyTorch. Reproducibility documentation, including scaled-dot-product attention backend differences. PyTorch Documentation. Accessed September 2026. ↩ ↩² ↩³
Tri Dao. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. ICLR 2024. ↩ ↩² ↩³ ↩⁴ ↩⁵ ↩⁶ ↩⁷ ↩⁸ ↩⁹
Jay Shah et al. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024. ↩ ↩² ↩³ ↩⁴ ↩⁵
Ted Zadouri et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. arXiv, 2026. ↩ ↩² ↩³ ↩⁴ ↩⁵
Markus N. Rabe and Charles Staats. Self-attention Does Not Need O(n²) Memory. arXiv, 2021. ↩ ↩²
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv, 2018. ↩ ↩² ↩³ ↩⁴
Dao-AILab. flash-attention official repository and current implementation documentation. GitHub. Accessed September 2026. ↩ ↩² ↩³ ↩⁴
Ashish Vaswani et al. Attention Is All You Need. NeurIPS 2017. ↩
PyTorch. torch.nn.attention.sdpa_kernel documentation. PyTorch Documentation. Accessed September 2026. ↩
Woosuk Kwon et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. SOSP 2023. ↩

Mixture of Experts (MoE): How Transformers Scale Without Activating Everything.

Mon, 10 Aug 2026 00:00:00 +0000

Mixture of Experts (MoE) is one of the main techniques used to scale modern language models without making every token pay the full computational cost of the model.
The basic idea is surprisingly simple: instead of sending every token through one enormous feed-forward network, we split it into many smaller expert networks and only activate a few experts for each token.
These notes walk through the intuition behind MoE, starting from the role of the feed-forward network (FFN) inside a Transformer, then moving through routing, top-$k$ selection, fine-grained experts, shared experts, expert capacity, dropless MoE, load balancing, and router stability.
The focus is not just on what MoE does, but why each piece exists and what problem it is trying to solve.

I originally learned this material from Jia-Bin Huang’s visual explanation of MoE ¹. I have rewritten the ideas here as notes for myself, with the equations and implementation details that I found useful when trying to understand how modern sparse MoE models actually work. After working through it, the name makes a lot more sense. It really is a mixture of experts. The interesting part is deciding who gets called for each token, and then making sure the whole system doesn’t collapse under its own weight.

Feed-Forward Networks in Transformers
Why Mixture of Experts?
Sparse Mixture of Experts
- The Router
- Top-$k$ Routing
Fine-Grained Experts
Shared Experts
Expert Capacity & Token Overflow
Dropless MoE
The Load Balancing Problem
DeepSeek V3’s Bias Adjustment Trick
Auxiliary-Loss-Free Load Balancing
Router Stability & the Router Z-Loss
- Safe Softmax
- Router Z-Loss
Putting It All Together
The Bigger Picture

Appendix

References
Citation

Feed-Forward Networks in Transformers

A Transformer layer typically alternates between an attention mechanism and a feed-forward network (FFN). Attention allows each token to incorporate information from other tokens in the sequence. The FFN then processes each token independently. Attention basically answers this:

“Which other tokens should this token pay attention to?”

The decoder-only transformer architecture.

Figure 1: The decoder-only transformer architecture. (Source: Cameron R. Wolfe)

The FFN is different. It operates independently on each token after attention has produced its contextual representation. A simplified Transformer block looks roughly like:

\[\boxed{x \rightarrow \text{Attention} \rightarrow \text{FFN} \rightarrow x'}\]

The FFN is usually a relatively large multilayer perceptron (MLP). If the input token representation has dimension $d$, the FFN typically expands it to a much larger hidden dimension $d_h$, applies a nonlinear activation, and projects it back down:

\[x \rightarrow W_{up}x \rightarrow \sigma(\cdot) \rightarrow W_{down}a\]

For a single token embedding $x$, a simplified FFN can be written as:

\[\begin{aligned} z &= xW_{\text{up}} + b_{\text{up}} \\ a &= \sigma(z) \\ y &= aW_{\text{down}} + b_{\text{down}} \end{aligned}\]

where:

$W_{\text{up}}$ projects the token into a larger hidden dimension.
$\sigma$ is the nonlinear activation function.
$W_{\text{down}}$ projects the representation back to the original model dimension.
$b_{\text{up}}$ and $b_{\text{down}}$ are the biases which could be zero or nonzero.
and the result is added back through the residual connection.

Typically, $d_h \approx 4d$, where $d$ is the input dimension and $d_h$ is the hidden dimension of the FFN.

A Useful Intuition for FFN

One way to think about the first projection is that each hidden dimension asks the token a different learned “question.” If a row of $W_{up}$ points strongly in some semantic direction, its dot product with the token representation tells us how strongly that feature is present. The activation function then decides which features should remain active. Finally, $W_{down}$ maps those activated features back into the model dimension. This gives us a useful mental model:

up projection = detect features
activation = select features
down projection = combine features into an output

There is an important caveat here. It is tempting to say that the FFN literally stores one clean factual association per neuron, but that is too strong. The weights are distributed representations, and features can be entangled across dimensions. The “questions” and “facts” interpretation is a useful intuition, not a literal description of what every neuron is doing.

RMSNorm

Before the FFN, modern Transformers commonly use a normalization layer such as RMSNorm. For an input vector $x$, given that $\gamma$ is a learnable scaling vector, the RMSNorm is:

\[\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d}x_i^2+\epsilon}\] \[\text{RMSNorm}(x) = \gamma \odot \frac{x}{\text{RMS}(x)}\]

The important idea is that normalization keeps the magnitude of activations under control while allowing the model to learn a different scale for each dimension.

The FFN as a Knowledge Store

The FFN is more than just a generic nonlinear transformation. One useful interpretation is that the first projection asks a collection of learned questions about the token representation. Each row of $W_{\text{up}}$ can be thought of as a learned direction in representation space.

\[z_i = x \cdot W_{\text{up},i} + b_i\]

If $z_i$ is large, the input strongly matches the feature represented by that row. The activation function then suppresses irrelevant features. For a simple ReLU example:

\[a_i = \max(0,z_i)\]

The second projection maps the activated features back into the model dimension. This gives us an intuitive picture:

The FFN can be thought of as a large collection of learned feature detectors that activate different pieces of stored information depending on the input.

This interpretation is useful for understanding why increasing the FFN hidden dimension can improve model capacity.

Why make the FFN bigger?

Now suppose we want a more capable model. One obvious option is to increase $d_h$, the hidden dimension of the FFN. More hidden dimensions means more learned features and more capacity. However, if we make the FFN four times larger, eight times larger, or even more, we also increase the amount of computation performed for every token. A dense model has to use the same FFN for every token. That means the computational cost grows roughly with the total size of the network. This gives us a frustrating trade-off:

\[\boxed{ \text{more parameters} \quad\Longrightarrow\quad \text{more capacity} }\]

but also:

\[\boxed{ \text{more parameters} \quad\Longrightarrow\quad \text{more computation per token} }\]

Why Mixture of Experts?

What if we could have the extra parameters without using all of them every time? That is where MoE enters. If we simply increase $d_h$, the FFN becomes larger. That gives the model more capacity, but it also increases:

training computation,
inference computation,
parameter memory,
communication requirements.

And there is another observation: A token does not need every feature in the FFN.

For example, a token about chemistry probably does not need every feature that might be useful for programming, mathematics, or another language. So instead of making one enormous FFN that processes every token, we can divide the FFN into multiple smaller networks. Each one becomes an expert. The model can then choose which experts should process each token. This is the central idea behind a sparse Mixture of Experts (MoE).

Figure 2: Mixture of Experts (MoE) routing in a Switch Transformer encoder block. The standard dense feed-forward network (FFN) layer is replaced by a sparse Switch FFN layer (light blue) that processes sequence tokens independently. As illustrated, a router independently directs two tokens (x₁ = “More”, x₂ = “Parameters”) via solid lines to specific FFN experts across 4 FFN experts. The final layer output (dotted line) is the selected FFN's response multiplied by its router gate value. (Source: William Fedus, Barret Zoph, Noam Shazeer, 2022)

The crucial part is that we do not run every expert. The router selects only a small number of experts for each token. If we have 64 experts but activate only 8 for a token, then the model can contain a lot more total parameters than a dense model with roughly the same amount of computation per token. A simplified MoE layer can be written as:

\[\text{MoE}(x) = \sum_{i \in \text{TopK}(r(x))} p_i(x) E_i(x)\]

where:

$x$ is the token representation
$E_i$ is expert $i$
$r(x)$ is the router output
$p_i(x)$ is the routing probability
$\text{TopK}$ selects the experts that actually process the token

This is the key trick:

The model’s total parameter count can grow much faster than the number of parameters activated for each token.

That is the “sparse” in sparse MoE.

Sparse Mixture of Experts

Suppose we have $N$ experts:

\[E_1, E_2, \dots, E_N\]

Each expert is itself an FFN. Instead of evaluating all $N$ experts for every token, we select only the top $k$ experts. If $k \ll N$, then the model can contain many more parameters while only activating a small fraction of them for each token. This creates an important distinction:

Total parameters and active parameters are no longer the same thing.

A model can therefore have a very large parameter count without requiring every token to use the entire model.

Figure 3: Sparse Mixture-of-Experts (s-MoE) overview. Core mechanism, motivation, training challenge and solution. (Source: Emergent Mind)

The Router

So now we have another problem. If there are many experts, who decides which ones get used? The model needs some mechanism to decide which experts should receive each token. This is the job of the router, a small neural network. A simple router can just be a learned linear projection.

For an input token representation $x$, the router computes a score for every expert:

\[h(x) = xW_g\]

where

\[W_g \in \mathbb{R}^{d \times N}\] \[h(x) \in \mathbb{R}^{N}\]

Each value in $h(x)$ is the router’s score, or logit, for one expert. We can turn these logits into probabilities using softmax:

\[p_i(x) = \frac{\exp(h(x)_i)} {\sum_{j=1}^{N}\exp(h(x)_j)}\]

The resulting vector represents the router’s preference over the experts.Then we select the top $K$ experts.

Top-$k$ Routing

We then select only the $k$ experts with the highest routing scores. Let

\[S(x) = \text{TopK}(p(x),k)\]

Then the MoE output can be written as:

\[y = \sum_{i \in S(x)} p_i(x)E_i(x)\]

In practice, implementations may renormalize the selected routing weights, but the basic idea remains the same. The process is therefore:

\[\boxed{x \rightarrow \text{Router} \rightarrow \text{Top-}k\text{ Experts} \rightarrow \text{Weighted Combination} \rightarrow y}\]

The routing happens at the token level. Different tokens in the same sequence can therefore be sent to completely different combinations of experts. For example, if the router produces:

[0.05, 0.10, 0.03, 0.42, 0.07, 0.33]

and $K=2$, experts 4 and 6 win. The token is sent only to those experts. Their outputs are then combined using the corresponding routing weights. This means routing happens per token. Two neighboring tokens can go through completely different experts.

Fine-Grained Experts

A natural question is:

Why not just have a small number of large experts?

One answer is flexibility. Suppose we have eight large experts and select two of them. The number of possible combinations is limited.

Instead of having a few large experts, why not have many smaller experts?

We can split the expert computation into many smaller experts and activate more of them. This is the idea behind fine-grained expert segmentation. Suppose we have roughly the same total expert capacity. We could build:

8 large experts
OR 
64 smaller experts 

and activate a subset of them.

The second approach gives the router many more possible combinations to choose from. This can be useful because the model does not have to commit a token to one very large “specialist.” It can combine several smaller specialists. For example:

Token A → Expert 3 + Expert 11 + Expert 42 + Expert 57
Token B → Expert 2 + Expert 8 + Expert 19 + Expert 61

The resulting combinations give the model a much larger routing space. OLMoE, for example, uses 64 small experts and activates 8 per token ². Its released configuration has about 6.9B total parameters with about 1.3B active parameters per token. DeepSeek-V3 pushes this idea much further, using hundreds of routed experts while activating only a small subset for each token. DeepSeekMoE proposed splitting the experts into a much larger number of smaller experts and activating a correspondingly larger number of them. The goal is to allow more flexible combinations of specialized knowledge ³. This gives the router a finer-grained set of choices:

\[\boxed{ \text{many small experts} + \text{top-}k\text{ routing} \rightarrow \text{more possible combinations}}\]

So MoE is not just:

“Let’s make several copies of the FFN.”

The design of the expert pool itself matters. This is one of the important ideas behind modern DeepSeek-style MoE architectures.

Shared Experts

DeepSeekMoE also introduced the idea of shared experts. A shared expert is always active rather than being selected by the router. The intuition is that some information is useful for almost every token. Instead of forcing the routed experts to repeatedly learn this common information, we can give the model a dedicated expert for broadly useful patterns. The architecture therefore contains:

\[\boxed{\text{Shared Experts} + \text{Routed Experts}}\]

The shared component handles more general information, while the routed experts can specialize. The DeepSeekMoE paper found benefits from combining fine-grained experts with shared experts, although the value of shared experts is architecture-dependent and is not universally guaranteed ³.

Figure 4: Illustration of DeepSeekMoE. Subfigure (a) showcases an MoE layer with the conventional top-2 routing strategy. Subfigure (b) illustrates the fine-grained expert segmentation strategy. Subsequently, subfigure (c) demonstrates the integration of the shared expert isolation strategy, constituting the complete DeepSeekMoE architecture. It is noteworthy that across these three architectures, the number of expert parameters and computational costs remain constant. (Source: Damai Dai et al., 2024)

Expert Capacity & Token Overflow

There is a practical problem with routing. Suppose we have $16$ tokens and $8$ experts. If tokens were distributed perfectly evenly, each expert would receive $\frac{16}{8}=2$ tokens. So we could give every expert a capacity of 2. But routing decisions depend on the actual input. There is nothing inherently stopping the router from doing this instead:

Expert 1:  8 tokens
Expert 2:  5 tokens
Expert 3:  2 tokens
Expert 4:  1 token
Expert 5:  0 tokens
Expert 6:  0 tokens
Expert 7:  0 tokens
Expert 8:  0 tokens

OR

Expert 1:  5 tokens
Expert 2:  3 tokens
Expert 3:  2 tokens
Expert 4:  2 token
Expert 5:  1 tokens
Expert 6:  1 tokens
Expert 7:  1 tokens
Expert 8:  1 tokens

Now the first expert has received more tokens than it can process in both cases. This creates token overflow. Some experts are overloaded while others are sitting idle in the first case. And in a distributed system, experts may live on different GPUs. So routing is not merely a modeling decision, but also a systems problem. Historically, one way to deal with this is to give every expert a fixed batch size. This allows us to define the expert capacity, which is the maximum number of tokens in a batch that can be routed to each expert. Expert capacity allows more buffer to help mitigate token overflow during routing ⁴.

\[\text{expert capacity} = \text{capacity factor} \times \frac{T}{N}\]

where:

$T$ = number of total tokens per batch
$N$ = number of experts

Figure 5: Illustration of token routing dynamics. (Source: William Fedus, Barret Zoph, Noam Shazeer, 2022)

If an expert receives more tokens than its capacity, something has to happen. One option is to drop the excess tokens. The token essentially skips the MoE computation and continues through the residual pathway. That avoids exceeding the allocated compute and memory budget, but now some tokens are not getting the expert computation the model intended. Another traditional solution is to increase the expert’s capacity. But then we have the opposite problem. For example, a capacity factor of $1.5$ would increase the available capacity above the ideal balanced allocation in the second case as shown above, but larger capacity means:

more memory,
more computation,
more padding,
more communication.

Suppose an expert receives only 2 tokens but we allocate space for 8. We now have six empty slots. The hardware still has to deal with those padded tensor shapes. So we have a trade-off:

\[\boxed{ \text{low capacity} \rightarrow \text{token overflow} }\] \[\boxed{ \text{high capacity} \rightarrow \text{wasted computation/memory} }\]

This is one reason modern MoE implementations care so much about efficient routing kernels. So can we avoid both?

Dropless MoE

Instead of forcing every expert to process the same number of tokens, we can organise the expert computation around the actual routing assignments. One way to think about this is using block-sparse matrix multiplication. Rather than padding every expert to the same token count, the computation can use blocks whose sizes correspond to the number of tokens actually assigned to each expert. Conceptually:

\[\boxed{ X \rightarrow \text{Routing} \rightarrow \text{Variable-sized expert batches} \rightarrow \text{Expert computation} \rightarrow \text{Scatter back} }\]

For example, we can use sparse/block operations that operate on the actual routed tokens:

Token routing:
Expert 1 → x1, x7, x9
Expert 2 → x2
Expert 3 → x3, x4, x5, x6

Rather than padding everything to four tokens:
Expert 1 → x1, x7, x9, PAD
Expert 2 → x2, PAD, PAD, PAD
Expert 3 → x3, x4, x5, x6

This allows us to avoid dropping tokens simply because an expert received more tokens than expected, while also avoiding unnecessary padding. This style of implementation is often referred to as dropless MoE. The exact implementation depends heavily on the GPU kernels and distributed training system, but the underlying goal is simple:

Compute only the expert work that actually exists.

OLMoE, for example, uses token-choice routing with dropless routing. The important distinction is that MoE sparsity is not just about the neural network architecture. We need the underlying hardware and kernels to exploit that sparsity too. Otherwise, theoretical efficiency can disappear in implementation overhead.

The Load Balancing Problem

There is a deeper problem with sparse routing. At the beginning of training, all experts are randomly initialised. Suppose the router happens to send many of the first tokens to experts $E_1$ and $E_2$. Those experts are updated more often so they become better. The router then has even more reason to send tokens to them, and the cycle continues. This can create a self-reinforcing feedback loop:

\[\boxed{ \text{more tokens} \rightarrow \text{more updates} \rightarrow \text{better experts} \rightarrow \text{higher routing score} \rightarrow \text{more tokens} \rightarrow \text{same better experts} }\]

Eventually, some experts may receive very few tokens. These underused experts become effectively dead. This is the load balancing problem. The goal is not necessarily to make every expert equally good, but to prevent the router from collapsing onto a small subset of experts by distributing the processing of tokens as evenly as possible among the experts.

Noisy Top-$k$ Gating

One early approach is to add Gaussian noise to the router logits:

\[h(x) = xW_g\] \[h'(x)_i = h(x)_i + \epsilon \cdot \text{Softplus}\big((xW_{\text{noise}})_i\big) \quad \text{where } \epsilon \sim \mathcal{N}(0, 1)\]

The noise encourages exploration. Instead of always selecting the same experts, the router occasionally explores other experts, giving them opportunities to receive tokens and learn. The router selects a more diverse set of experts and prevents the same experts from always being chosen. We can further improve this by adding a tunable scaling factor for each expert. This is called noisy top-K gating. This idea appears in noisy top-$k$ gating approaches to MoE routing ⁵. But exploration alone is not enough. We still need an explicit mechanism to encourage balanced expert utilisation.

Importance vs Load

Given N experts indexed by $i = 1 \text{ to } N$ and a batch $B$ with $T$ tokens, there are two slightly different things we can measure.

Importance

We can measure how much routing probability an expert receives. For expert $i$, the importance $I_i$ is:

\[I_i = \sum_{x \in B} p_i(x)\]

An expert can therefore have high importance even if it is not selected very often.

Load

Instead, we can count how many tokens are actually routed/assigned to the expert $i$. For expert $i$, the load $L_i$ is:

\[L_i = \sum_{x \in B} \mathbb{1}[i\in S(x)]\] \[\text{where } S(x) = \text{TopK}(p(x),k)\]

This measures actual expert usage. These two quantities ($L_i, I_i$) are not necessarily the same. For example, the router might give eight experts reasonably balanced probabilities while repeatedly selecting only four of them through top-$k$ routing. Therefore, balancing probabilities alone does not guarantee balanced computation. Unfortunately, the number of tokens received by each expert is a discrete quantity that cannot be used directly for backpropagation. Instead, we can define a smooth estimator for the number of examples assigned to each expert in the batch. The smoothness enables gradients to flow through the estimator, making it possible to apply backpropagation.

Load Balancing Loss

A commonly used loss formulation combines the fraction of tokens routed to each expert with the average routing probability assigned to that expert. Mathematically, the load balancing loss is:

\[P_i = \frac{1}{T} I_i = \frac{1}{T} \sum_{x \in B} p_i(x), \qquad f_i = \frac{1}{T} L_i = \frac{1}{T} \sum_{x \in B} \mathbb{1}[i \in S(x)]\] \[\mathcal{L}_{\text{LB}} = N \sum_{i=1}^{N} f_i P_i\]

where:

$f_i$ is the fraction of tokens routed to expert $i$.
$P_i$ is the average router probability assigned to expert $i$.
$N$ is the number of experts.

The factor $N$ keeps the scale of the loss comparable as the number of experts changes. Essentially, it keeps the scale of the loss from shrinking simply because we increased the number of experts. Under perfectly uniform routing where $f_i = P_i = \frac{1}{N}$, the scaling factor $N$ cancels out the splitting effect of the cross-expert summation:

\[N \sum_{i=1}^{N} f_i P_i = N \sum_{i=1}^{N} \left( \frac{1}{N} \cdot \frac{1}{N} \right) = N \sum_{i=1}^{N} \frac{1}{N^2} = N \cdot \frac{N}{N^2} = 1\]

A common approach is to add an auxiliary load balancing loss to the language modeling objective. Conceptually:

\[\mathcal{L} = \mathcal{L}_{\text{CE}} + \alpha \mathcal{L}_{\text{LB}}\]

where:

$\mathcal{L}_{\text{CE}}$ is the normal next-token prediction loss (cross-entropy loss).
$\mathcal{L}_{\text{LB}}$ encourages more uniform expert usage (load-balancing loss).
$\alpha$ is a hyperparameter that controls the influence of the load balancing loss during training. It controls how strongly the balancing objective affects training.

The important idea is that the balancing loss gives the router a reason to spread tokens across experts instead of collapsing onto a few. However, there is an obvious downside.

If $\alpha$ is too large, the model may optimise for balanced routing at the expense of the actual language modeling objective. The router is forced to care too much about uniformity, and uniform routing is not necessarily what we actually want. The model should learn to specialize.
If $\alpha$ is too small, the balancing loss may have almost no effect.

So:

\[\boxed{ \alpha \uparrow \rightarrow \text{better balance, potentially worse model objective} }\] \[\boxed{ \alpha \downarrow \rightarrow \text{less interference, potentially worse balance} }\]

From experimental results, we see that a clear advantage of using load balancing loss is that it produces lower training and validation loss. Without load balancing, two out of eight experts receive all the training tokens while the remaining six experts are poorly trained. With load balancing loss, the tokens are distributed much more evenly across all experts.

Device-Level Load Balancing

In a real distributed MoE model, experts are often spread across different GPUs. This introduces another problem, a systems problem. Even if the experts are balanced globally, the devices might not be. For example:

\[\boxed{ \text{GPU}_1 \rightarrow \text{many tokens} }\] \[\boxed{ \text{GPU}_2 \rightarrow \text{few tokens} }\]

The model is still bottlenecked by the overloaded GPU. Therefore, large-scale MoE systems may also consider load at the device level. These experts may live on different devices. Therefore, in addition to expert-level balance, we can further distribute the computational load evenly across devices and alleviate bottlenecks. The computation of device-level balance loss follows a similar formulation but aggregates token counts and router probabilities at the device level. The same general principle applies:

\[\boxed{ \text{balanced experts} + \text{balanced devices} \rightarrow \text{better hardware utilization } }\]

This matters because MoE introduces communication between devices whenever tokens need to be dispatched to experts located on different GPUs.

DeepSeek V3’s Bias Adjustment Trick

DeepSeek V3 uses a simple trick to balance the load. Instead of using softmax, the router in DeepSeek V3 directly applies a sigmoid function on the predicted logits, one for each expert. Suppose the router is slightly biased towards selecting the first two experts. To overcome this imbalance, we add a small bias to each expert depending on whether the number of tokens routed to the expert is larger or smaller than the average. This simple bias adjustment increases the likelihood of tokens being routed to a broader range of experts. By dynamically adjusting the bias for overloaded or underloaded experts, we keep balanced expert load while achieving better performance.

Note: The bias adjustments were only used for determining the top-K routing. The final output still uses the original probabilities as weights for combining expert outputs.

Auxiliary-Loss-Free Load Balancing

DeepSeek-V3 introduced an interesting alternative to the traditional auxiliary balancing loss. Instead of adding another loss term that directly competes with the language modeling objective, DeepSeek-V3 uses a dynamic expert-level bias to influence which experts are selected ⁶. Conceptually:

CASE 1:
overloaded expert
      ↓
lower routing bias
      ↓
less likely to be selected 

CASE 2:
underloaded expert
      ↓
higher routing bias
      ↓
more likely to be selected

The idea is roughly:

Measure how many tokens each expert is receiving.
Compare each expert’s load with the average.
Increase the routing bias for under-utilised experts.
Decrease the routing bias for overloaded experts.
Use the adjusted scores for top-$k$ selection.
Keep the original router scores for weighting the selected expert outputs.

Mathematically:

\[h'(x)_i = h(x)_i + b_i\]

where $b_i$ is a dynamically updated expert-specific bias.

If an expert is overloaded $b_i \downarrow$
If an expert is underloaded $b_i \uparrow$

The bias is used to influence the selection of the top-$K$ experts. But the original routing probabilities are still used when combining the expert outputs. That means the mechanism can encourage balanced routing without directly adding a load-balancing penalty to the main optimization objective. This is why it is often described as loss-free load balancing. It is a nice example of a recurring theme in large-scale ML:

Sometimes the best way to fix an optimization problem is to change the algorithm around the objective rather than adding another term to the objective.

Essentially, the key detail is that this bias is used for routing selection, rather than changing the probabilities used to combine the final expert outputs. This allows DeepSeek-V3 to perform load balancing without introducing the same auxiliary-loss trade-off. The DeepSeek-V3 technical report describes this as an auxiliary-loss-free load balancing strategy ⁶. Results demonstrate that this auxiliary-loss-free load balancing strategy avoids the trade-off between load balance and model performance. There is also an important nuance here. DeepSeek-V3 does not completely eliminate the need to monitor routing balance. Sequence-level imbalance can still matter, especially when individual sequences produce skewed routing patterns. DeepSeek-V3 uses a sequence-wise balance loss to prevent severe imbalance within individual sequences; so it’s not completely loss-free, but it avoids a global auxiliary loss that competes directly with the main objective.

Router Stability & the Router Z-Loss

Load balancing is not the only issue with the router. Recall that the router begins with logits $h(x)$ and converts them to probabilities using softmax:

\[p_i(x) = \frac{e^{h(x)_i}}{\sum_{j=1}^{N} e^{h(x)_j}}\]

An interesting property of softmax is that it is shift invariant. For any constant $c$:

\[\text{Softmax}(h) = \text{Softmax}(h+c)\]

because:

\[\frac{e^{h_i+c}} {\sum_j e^{h_j+c}} = \frac{e^c e^{h_i}} {e^c\sum_j e^{h_j}} = \frac{e^{h_i}} {\sum_j e^{h_j}}\]

This creates a subtle problem. The router can increase all of its logits by the same amount without changing its output probabilities. So the probabilities remain stable while the underlying logits can grow arbitrarily large. This becomes dangerous when using reduced-precision arithmetic such as FP16. For example:

\[[2, 1, 0] \text{ and } [102, 101, 100]\]

produce exactly the same softmax distribution. The model therefore has no incentive, from the softmax probabilities alone, to keep the absolute magnitude of the logits small. This poses a numerical problem because large logits can interact badly with exponentiation especially in modern LLMs that are trained using reduced-precision arithmetic (BF16 or FP16) for memory reduction and training speed increase.

Safe Softmax

A standard numerical trick is to subtract the maximum logit before exponentiation:

\[p_i(x) = \text{Softmax}(h)_i = \frac{e^{h(x)_i - \max(h(x))}}{\sum_{j=1}^{N} e^{h(x)_j - \max(h(x))}}\]

This prevents excessively large exponentials from overflowing. However, it only treats the numerical symptom. The underlying logits can still drift to extremely large values. Using safe softmax is like wearing a seat belt while driving. It does provide important protections, but it does not mean we should press the accelerator recklessly and allow the logits to grow without constraints.

Router Z-Loss

To prevent the logits from becoming excessively large, a straightforward approach is to add a loss term that penalizes the sum of their square values. Unfortunately, this hurts model performance because it interferes with the router probability from the softmax. A more direct approach is to regularise the normalisation term. The idea is to penalize large values of the router’s log-normalizer. This ensures that the gradient from this loss does not change the router probabilities. Let’s define:

\[Z(x) = \log \sum_{i=1}^{N} e^{h_i(x)}\]

We square the log normalizer to penalize both positive and negative shifts and average them across the batch. This is known as the router Z-loss. The router Z-loss can be written as:

\[\mathcal{L}_Z = \frac{1}{T} \sum_{x \in B} Z(x)^2\]

The logarithm prevents the regularisation term from growing exponentially, while the squared term penalises large positive or negative shifts (large values of the log-normalizer). The router Z-loss keeps the logits centered and prevents them from drifting to arbitrarily large values, ensuring training stability is crucial for reliable scaling of mixture of experts to larger and more powerful models. The intuition is:

Keep the router logits numerically well-behaved without unnecessarily changing which experts the router prefers.

This is particularly important when scaling MoE training to large models reliably. Scaling the number of experts is not enough. We also have to scale the stability machinery around them. From controlled experiments in OLMoE ², without Z-loss, the training curve has many loss spikes. In contrast, adding router Z-loss leads to more stable training and better performance.

Putting It All Together

At this point, the architecture looks something like this:

                    Token representation
                           |
                         RMSNorm
                           |
                      Transformer
                        Attention
                           |
                      MoE / FFN
                           |
                    +------+------+
                    |             |
                  Router       Experts
                    |             |
             score every      Expert 1
               expert         Expert 2
                    |          ...
                 Top-K         Expert N
                    |             |
                    +------+------+
                           |
                    Weighted combine
                           |
                      Residual add
                           |
                         Output

And during training, we are optimizing more than just next-token prediction. A simplified objective is:

\[\boxed{ \mathcal{L}_{\text{CE}} + \alpha\mathcal{L}_{\text{LB}} + \beta\mathcal{L}_{\text{Z}} }\]

where:

$\mathcal{L}_{\text{CE}}$ teaches the model to predict the next token
$\mathcal{L}_{\text{LB}}$ encourages healthy expert utilization
$\mathcal{L}_{\text{Z}}$ keeps the router numerically stable

Some modern architectures replace the explicit load-balancing loss with other routing strategies, such as DeepSeek-V3’s loss-free bias adjustment. So there is not one universal MoE recipe. There is a family of design decisions around:

number of experts
expert size
number of active experts
routing function
load balancing
capacity
communication
numerical stability
shared experts
hardware kernels

That is what makes modern MoE systems much more interesting than the simple diagram suggests. The entire sparse MoE pipeline can now be summarised as:

\[x \rightarrow \boxed{\text{Router}} \rightarrow \boxed{\text{Top-}k\text{ Selection}} \rightarrow \boxed{\text{Expert Dispatch}} \rightarrow \boxed{\text{Expert FFNs}} \rightarrow \boxed{\text{Weighted Combination}} \rightarrow y\]

But making this work at scale requires solving several different problems:

Problem	Technique
Too many FFN parameters	Sparse expert activation
Limited expert specialisation	Fine-grained experts
Common information repeated across experts	Shared experts
Uneven token assignments	Capacity management
Token dropping / padding	Dropless MoE
Expert collapse	Load balancing
GPU bottlenecks	Device-level balancing
Balancing hurts the main objective	Auxiliary-loss-free routing
Router logits grow uncontrollably	Router Z-loss

The thing I find most interesting about MoE is that the main idea is actually very simple:

Build a much larger network, but only use a small part of it for each token.

The difficult part is everything that comes afterward. The router has to learn where each token should go. The experts have to specialise without some of them becoming useless. The tokens have to be distributed across GPUs without creating communication bottlenecks. And the whole routing system has to remain numerically stable while the model is being trained. That is what makes modern MoE architectures more than simply “a bunch of FFNs with a router.”

The Bigger Picture

The thing I found most interesting about MoE is that it changes how we think about scaling neural networks.

A dense model says:

“Make the whole network bigger, and every token pays for the bigger network.”

An MoE model says:

“Make the network much bigger, but only activate the parts that matter for this token.”

That gives us two different notions of model size:

\[\boxed{\text{Total Parameters}}\]

and:

\[\boxed{\text{Active Parameters per Token}}\]

Those numbers can be dramatically different. This distinction is now fundamental when comparing large MoE models. A model with 1T total parameters is not necessarily performing 1T parameters’ worth of computation for every token. Only a subset of those parameters may participate in a given forward pass. But there is an important caveat that is easy to miss:

Sparse compute does not mean sparse memory.

The expert weights still have to exist somewhere. During training and inference, the system has to store and distribute those parameters across devices. And the router has to move tokens to the devices that contain the selected experts. That introduces communication overhead, especially the all-to-all communication used in distributed MoE systems. So MoE does not magically make a trillion-parameter model equivalent to a tiny model. Instead, it changes the trade-off:

\[\boxed{ \text{more total capacity} \quad\text{vs.}\quad \text{roughly controlled active compute} }\]

while introducing additional memory, communication, and systems complexity. That trade-off is the real reason MoE is powerful. The easiest way I’ve found to remember MoE is shown below.

Key Takeaways

The FFN gives the model capacity.
The experts split that capacity into specialists.
The router decides which specialists a token needs.
Load balancing prevents the specialists from becoming unused or overloaded.
Router Z-loss keeps the routing mechanism numerically stable.
And underneath all of this, the hardware has to make the sparse computation actually efficient.

So the architecture is really two problems intertwined:

\[\boxed{ \text{Machine Learning} + \text{Systems Engineering} }\]

The machine learning problem is:

How do we learn useful routing and specialization?

The systems problem is:

How do we execute that routing efficiently across accelerators?

That combination is what makes Mixture of Experts such a powerful scaling strategy. And honestly, after working through it, the name makes a lot more sense. It really is a mixture of experts. The interesting part is deciding who gets called for each token.

Appendix

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026mixtureOfExperts,
  title   = "Mixture of Experts (MoE): How Transformers Scale Without Activating Everything",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Aug",
  url     = "https://chizkidd.github.io/2026/08/10/mixture-of-experts/"
}

References

Jia-Bin Huang. Mixture of Experts (MoE), Visually Explained. YouTube. Accessed August 2026. ↩
Muennighoff et al. OLMoE: Open Mixture-of-Experts Language Models. arXiv, 2024. ↩ ↩²
Damai Dai et al. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. ACL 2024. ↩ ↩²
William Fedus et al. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv, 2021. ↩
Noam Shazeer et al. Sparsely-Gated Mixture-of-Experts. arXiv, 2017. ↩
DeepSeek-AI et al. DeepSeek-V3 Technical Report. arXiv, 2024. ↩ ↩²

How Attention Became Efficient & Scalable: KV Caching, MQA, GQA, MLA, and Sparse Attention.

Wed, 05 Aug 2026 00:00:00 +0000

Attention mechanisms have evolved considerably to make transformer inference faster and more memory-efficient.
These notes trace that evolution: from vanilla self-attention, through KV caching, to memory-saving variants like MQA, GQA, and MLA, and finally to sparse attention methods like SWA and Deepseek Sparse Attention (DSA).

Self-Attention
Masked (Causal) Self-Attention
Multi-Head Attention
Key-Value Caching (KV Caching)
- Deepseek R1/V3 memory example
Multi-Query Attention (MQA)
Grouped Query Attention (GQA)
Multihead Latent Attention (MLA)
- MLA at Inference Time
Sliding Window Attention (SWA)
Deepseek Sparse Attention (DSA)
- Quantization & Rotation in DSA
- DSA Training

Appendix

References
Citation

Self-Attention

See detailed notes on transformer and attention here.

\[Q = XW_Q, \quad K = XW_K, \quad V = XW_V\] \[\text{Attention}(Q,K,V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V = O\]

Masked (Causal) Self-Attention

\[\text{Masked Attention}(Q,K,V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V\] \[\begin{aligned} \text{where} \\ M &\equiv \text{lookahead mask} \end{aligned}\]

Multi-Head Attention

\[\text{head}_i = \text{Attention}(Q,K,V)\] \[\text{MHA} = \text{multi-head attention}(Q,K,V) = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_H)W_O = Z\] \[\text{OR}\] \[\begin{aligned} O_i &= \text{Attention}(Q_i, K_i, V_i) \\ O &= \text{MultiheadAttention}(Q,K,V) \\ O &= \text{Concat}(O_1, O_2, \dots) \\ \Delta X &= OW_O \end{aligned}\]

Key-Value Caching (KV Caching)

Due to masking, the transformer model can compute attention for all the tokens at the same time. This is called the prefilling stage.
During inference, the key and value vectors from the previous tokens DO NOT change at all across multiple attention heads & all attention layers. It is therefore wasteful to recompute these vectors every time for each new token.
A simple solution is to store the key & value vectors from previous tokens in memory, and reuse them when needed. This is known as KV caching.
KV caching helps avoid unnecessary computation & therefore speeds up the decoding process.
In KV caching, we DO NOT need to cache any of the previous query vectors because of the masked out entries, which have no effect on the attention output.
KV caching is very effective but memory intensive.

Deepseek R1/V3 memory example

\[\text{Required memory} = (2)(\text{KV dim})(\text{Precision})(\#\text{Heads})(\#\text{Layers})(\text{Sequence length})\] \[= (2)(128)(2\ \text{bytes/element})(128\ \text{heads})(61\ \text{layers})(32{,}768\ \text{tokens})\] \[= 131\ \text{GB} \quad (\text{A lot of memory!!!})\]

Multi-Query Attention (MQA)

How do we reduce the memory requirements?
In the MHA mechanism, we cannot adjust the KV dim, precision, or the number of layers, but maybe we can look at the number of heads.
One simple idea is to reduce the number of heads for key and value matrices from $h$ to $1$. Therefore, there is only a single copy of key & value vectors for each token in a layer. This single copy of KV vectors is then shared across all attention heads. This is known as multi-query attention.

\[\begin{aligned} W_Q &\in \mathbb{R}^{d \times (d_k \times h)} \\ W_K &\in \mathbb{R}^{d \times (d_k \times 1)} \\ W_V &\in \mathbb{R}^{d \times (d_v \times 1)} \end{aligned}\]

	KV cache per token
MHA	4 MB
MQA	31 KB

128x reduction in memory usage

An issue with MQA, however, is that it sacrifices the ability to capture complex relationships between tokens. As a result, MQA’s performance degrades considerably compared to the original MHA.
The drastic reduction of number of heads from $128$ to $1$ is likely the reason for the performance degradation of MQA. We can instead reduce the number of heads to a smaller value denoted as $n_g$.

Grouped Query Attention (GQA)

Instead of the total reduction of number of heads to 1 for the key and value matrices as in MQA, GQA reduces the number of heads for the key and value matrices to a smaller value, $n_g$.
GQA maintains the same overall pattern of attention as MHA, but collapses the number of key-value heads by sharing them across multiple query heads.

Multi-head attention (MHA) vs Grouped query attention (GQA)

Figure 1: Multi-Head Attention vs Group Query Attention. (Source: Sebastian Raschka)

MHA: each of head₁, head₂, head₃, head₄ has its own Q, K, V.
GQA: head₁ and head₂ share one K/V pair; head₃ and head₄ share another K/V pair.
MQA: head₁, head₂, head₃, head₄ all share a single K/V pair.

GQA strikes a balance between the memory efficiency of MQA and the expressive power of MHA.
GQA is a popular choice in modern LLMs, including Llama 3 8B (Meta), Qwen 3 4B (Alibaba), Gemma 3 27B (Google), Mistral Small 3.1 24B (Mistral), SmolLM3 3B (Hugging Face), etc.

	KV cache per token
MHA	4 MB
MQA	31 KB
GQA ($n_g = 16$)	500 KB

8x reduction in memory usage from MHA to GQA ($n_g = 16$)

Multihead Latent Attention (MLA)

See detailed notes on MLA here.
MLA reduces memory usage & slightly improves model performance compared to others using MHA.

Figure 2: Multihead Latent Attention.

MLA compresses $Q, K, V$ representations into a low-rank space to reduce the size of the KV cache using a down-projection matrix which produces latent representations:
\[C^Q = XW^Q_\downarrow \qquad C^{KV} = XW^{KV}_\downarrow\]
These compressed latent vectors are then mapped back to $W^Q_\uparrow$, $W^K_\uparrow$, $W^V_\uparrow$ for each attention head using the corresponding up-projection matrices.
Issue: This low-rank KV compression fails with rotary position embedding (RoPE).
Fix: A decoupled RoPE technique which introduces extra multi-head queries $W^{QR}$ and a shared key $W^{KR}$ to encode positional information.
For MLA, the Query, Key & Values are regrouped for each head:
- The Query is constructed by concatenating the compressed query $Q^C$ with the rotated query $Q^R$.
- The Key is constructed similarly by concatenating the compressed key $K^C$ with the rotated key $K^R$.

Multihead Latent Attention (MLA)

MLA compresses $Q, K, V$ representations into a low-rank latent space via down-projection matrices $W^Q_\downarrow$ and $W^{KV}_\downarrow$, then maps back up via up-projection matrices. A decoupled RoPE technique adds rotary queries $W^{QR}$ and a shared rotary key $W^{KR}$.

	KV cache per token
MHA	4 MB
MQA	31 KB
GQA ($n_g = 16$)	500 KB
MLA	70 KB

57x reduction in memory usage from MHA to MLA

MLA at Inference Time

\[\Delta X = OW_O = [O_1, O_2, \dots, O_h]W_O\] \[= \left[\dots, \text{Softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_k}} + M\right)V_i, \dots\right]W_O\] \[\begin{aligned} Q_i &= XW_{\downarrow}^Q W_{\uparrow,i}^Q \\ K_i &= C^{KV}W_{\uparrow,i}^K = XW_{\downarrow}^{KV}W_{\uparrow,i}^K \\ V_i &= C^{KV}W_{\uparrow,i}^V \end{aligned}\] \[\begin{aligned} \text{where} \\ C^{KV} &\equiv \text{cached latent} \end{aligned}\] \[\Delta X = \left[\dots, \underbrace{\text{Softmax}\left(\frac{XW_{\downarrow}^QW_{\uparrow,i}^QW_{\uparrow,i}^{K^T}C^{KV^T}}{\sqrt{d_k}} + M\right)}_{A_i}V_i, \dots\right]W_O\] \[\Delta X = [A_1V_1, A_2V_2, \dots, A_hV_h]W_O\] \[\Delta X = \left[A_1C^{KV}W_{\uparrow,1}^V, \dots, A_hC^{KV}W_{\uparrow,h}^V\right]W_O\] \[= [A_1C^{KV}, A_2C^{KV}, A_3C^{KV}, \dots, A_hC^{KV}] \begin{bmatrix} W_{\uparrow,1}^V & 0 & 0 & \cdots & 0 \\ 0 & W_{\uparrow,2}^V & 0 & \cdots & 0 \\ 0 & 0 & W_{\uparrow,3}^V & \cdots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \cdots & W_{\uparrow,h}^V \end{bmatrix}W_O\]

Note: the up-projection matrix for the value can be incorporated into the output matrix.

Due to the associative property, we can merge $W_{\downarrow}^QW_{\uparrow,i}^QW_{\uparrow,i}^{K^T}$ into a single matrix, which means we do not need to explicitly compute the key vectors during inference.
Also, $W_{\uparrow}^V$ can be absorbed into $W^O$. By incorporating the up-projection matrix into the query and output matrices, we can compute attention efficiently without incurring extra computation.
MLA has a fundamental limitation: it is incompatible with Rotary Positional Embedding (RoPE):

\[XW_{\downarrow}^QW_{\uparrow,i}^Q\underbrace{R_mR_n^T}_{\text{hinders merge}}W_{\uparrow,i}^{K^T}C^{KV^T} \quad (\text{RoPE})\]

We can no longer absorb the up-projection matrix $W_{\uparrow}^K$ during inference. Therefore, during inference, we must recompute the keys for all preceding tokens, which decreases inference efficiency.
To fix this, we introduce additional multihead queries & a shared key in order to effectively encode position information when using RoPE.
We concatenate the original query/key vectors and the rotated ones as input to the MHA computation. This approach is known as decoupled RoPE.
MLA is used in recent LLMs such as Deepseek V3 and R1, Kimi K2, GLM-5, Ling 2.5, Mistral Large 3, Sarvam 105B, etc.
However, as we generate more tokens, we must calculate attention between the current token and all preceding tokens. This results in slower token-per-second throughput. To fix this, Deepseek proposed a lightning indexer.

Sliding Window Attention (SWA)

Instead of attending to the entire prefix, each token only attends to a fixed window of recent tokens around its position.
SWA reduces the memory & compute cost of long-context inference (slower tok/s throughput) by limiting how many previous tokens each position can attend to.
Because attention is restricted to a local token neighborhood, this mechanism is often referred to as local attention.

Regular (causal) self-attention mask vs Sliding window attention

Figure 3: Causal Self Attention vs Sliding Window Attention (Window = 3). (Source: Sebastian Raschka)

Regular (causal) self-attention mask: each token attends to all preceding tokens.
Sliding Window Attention mask (window = 3): each token attends only to itself and the previous two tokens.

Deepseek Sparse Attention (DSA)

Deepseek proposed a lightning indexer to address the slower token-per-second throughput of MLA.
The key concept of the lightning indexer is to quickly assess the relevance of each token and select only the most relevant ones for attention computation.
The selected tokens are not determined by a fixed-width local window like in SWA, but by an indexer-plus-selector setup (top-k selector), where a lightning indexer computes index relevance scores, and a token selector keeps only a smaller set of high-scoring past positions ¹.

\[I^{t,s} = \sum_{j=1}^{n_h} w_j^t \, \text{ReLU}(q_j^t \cdot k_j^s) \equiv \text{Index score}\]

Apply partial RoPE to $W^Q$ and $W^K$ before rotation & quantization, then lightning indexing.

Quantization & Rotation in DSA

How is the attention computation of the indexer between the query vector and all the key vectors made efficient?
The key step is to first quantize the query & key vectors into 8-bit (BF16/FP32 → FP8) representation.
Quantization offers a coarse approximation, greatly speeding up the calculation of index scores, and is fine for use since we only want to find the most relevant tokens rather than compute the exact attention scores.
However, naive quantization of the query & key vectors may result in inaccuracy. To mitigate this issue, we need to apply a rotation.
Floating point formats with higher precision (float16 or above) can accurately represent the full dynamic range of values of $q^t$ and $k^s$ ($q^t = 0.85, -0.000000017$; $k^s = -0.99, 0.000001$). But this is problematic when we quantize these vectors with lower precision.

\[S_{orig} = q^t \cdot k^s \quad (\text{in FP32/FP16/BF16})\] \[S_{direct} = \text{Quantize}(q^t) \cdot \text{Quantize}(k^s)\]

Multiplying the query & key vectors with a random orthogonal matrix $R$ reduces the range of values, but this requires another matrix multiplication: $q^tR, k^sR$.
A simpler & even more effective approach is to use the Hadamard transform $H$, which yields deterministically computed output vectors that exhibit uniform mixing across all coordinates.
The Hadamard transform effectively spreads out large spikes across all coordinates, allowing the quantizer to retain more information.
The Hadamard transform can be implemented efficiently with highly optimized GPU kernels.
For implementation, a Fast Walsh-Hadamard Transform (FWHT) is used instead of a dense Hadamard matrix, $H_D$ ².

\[H_D = \begin{bmatrix} 1 & 1 & 1 & 1 \\ 1 & -1 & 1 & -1 \\ 1 & 1 & -1 & -1 \\ 1 & -1 & -1 & 1 \end{bmatrix}\frac{1}{\sqrt{4}} \quad \text{for a 4-head DSA}\]

Example

\[q^t = [0.999, -0.001, 0.001, -0.001]\] \[k^s = [-0.001, 0.897, -0.001, 0.001]\] \[\begin{aligned} S_{orig} &= q^t \cdot k^s \\ S_{direct} &= \text{Quantize}(q^t) \cdot \text{Quantize}(k^s) \\ S_{rotated-R} &= \text{Quantize}(q^tR) \cdot \text{Quantize}(k^sR) \\ S_{rotated-H} &= \text{Quantize}(q^tH) \cdot \text{Quantize}(k^sH) \end{aligned}\] \[\text{MAE}(S_{orig}, S_{direct}) = 0.0112, \quad \text{std} = 0.0085\] \[\text{MAE}(S_{orig}, S_{rotated-R}) = 0.0074, \quad \text{std} = 0.0061\] \[\text{MAE}(S_{orig}, S_{rotated-H}) = 0.0025, \quad \text{std} = 0.0020\]

DSA’s use of a lightning indexer to efficiently identify the most relevant tokens allows for 2 to 3x faster processing of long sequences, while reducing memory consumption by ~30 to 40%. This efficiency boost maintains the same performance level as previous models.

DSA Training

I) Dense warm-up stage: First, we freeze the main MLA layer & only update the parameters of the lightning indexer.

\[L^I = \sum_t D_{KL}\left(p^{t,:}, \text{Softmax}(I^{t,:})\right)\] \[\begin{aligned} \text{where} \\ p^{t,:} &\equiv \text{target distribution} \\ \text{Softmax}(I^{t,:}) &\equiv \text{prediction distribution} \end{aligned}\]

The target distribution is obtained by summing the mean attention scores across all heads for each query token, followed by applying L1 normalization along the sequence dimension.
This dense warm-up stage aligns the indexer’s outputs with the main attention distribution.

II) Sparse stage: After warming up the indexer, we add fine-grained token selection and train all the model parameters so that the model can learn the sparse attention patterns of DSA.

During this stage, we still encourage the indexer to match the main attention distribution, but only for the selected tokens.
We also detach the indexer inputs from the computational graph, allowing us to optimize it separately.
The indexer is trained only on its own loss, while the main model is optimized solely on the language modeling loss.

\[L^I = \sum_t D_{KL}\left(p^{t,S^t}, \text{Softmax}(I^{t,S^t})\right)\]

Appendix

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026attentionEfficientScalable,
  title   = "How Attention Became Efficient & Scalable",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Aug",
  url     = "https://chizkidd.github.io/2026/08/05/attention-efficient-scalable/"
}

References

Sebastian Raschka. A Visual Guide to Attention Variants in Modern LLMs. magazine.sebastianraschka.com. 2026. ↩
How Attention Got So Efficient (YouTube Video). YouTube. 2025. ↩

Policy Gradient Methods: REINFORCE, Actor-Critic, and the Policy Gradient Theorem (S&B Ch. 13)

Thu, 07 May 2026 00:00:00 +0000

Almost all the algorithms/methods covered so far have been action-value methods (except gradient-bandit algorithms, Section 2.8).
Action-value methods learn the values of actions and then derive the policy thereafter to select actions based on their estimated action values.
Here, we explicitly learn a parametrized policy that can select actions without consulting a value function.
A value function is not required for action selection, but still could be used to learn the policy parameters $\boldsymbol{\theta}$.
The parametrized policy $\pi(a \vert s, \boldsymbol{\theta})$ is now the probability that action $a$ is taken at time $t$ given that the environment is in state $s$ at time $t$ with parameter $\boldsymbol{\theta} \in \mathbb{R}^{d’}$:

\[\pi(a \vert s, \boldsymbol{\theta}) = \Pr\{A_t = a \vert S_t = s,\ \boldsymbol{\theta}_t = \boldsymbol{\theta}\}\]

We will consider methods for learning $\boldsymbol{\theta}$ based on the gradient of some scalar performance measure $J(\boldsymbol{\theta})$ w.r.t. $\boldsymbol{\theta}$. The goal is to maximize performance, hence the use of gradient ascent:

\[\boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t + \alpha \widehat{\nabla J(\boldsymbol{\theta}_t)}\] \[\begin{aligned} \text{where} \\ \widehat{\nabla J(\boldsymbol{\theta}_t)} &\in \mathbb{R}^{d'} \equiv \text{a stochastic estimate whose expectation approximates} \\ &\phantom{{}\equiv{}} \text{the gradient of the performance measure } J \text{ w.r.t. } \boldsymbol{\theta}_t \end{aligned}\]

This general methodology applies to policy gradient methods.
Methods that learn approximations to both policy & value functions are often called Actor-Critic methods
- Actor refers to the learned policy.
- Critic refers to the learned value function (state-value function).

13.1 Policy Approximation & Its Advantages
13.2 The Policy Gradient Theorem
13.3 REINFORCE: Monte Carlo Policy Gradient
13.4 REINFORCE with Baseline
13.5 Actor-Critic Methods
13.6 Policy Gradient for Continuing Problems
13.7 Policy Parametrization for Continuous Actions
13.8 Summary

Appendix

Citation

13.1 Policy Approximation & Its Advantages

In policy gradient methods, the policy can be parametrized in any way, as long as $\pi(a \vert s, \boldsymbol{\theta})$ is differentiable w.r.t. $\boldsymbol{\theta}$ (essentially the partial derivatives’ column vector exists and is finite):

\[\nabla \pi(a \vert s, \boldsymbol{\theta}) \text{ exists and is finite for all } s \in S,\ a \in A(s),\ \boldsymbol{\theta} \in \mathbb{R}^{d'}\]

If the action-space is discrete and not too large, then a natural & common kind of parametrization is to form parametrized numerical preferences $h(s, a, \boldsymbol{\theta}) \in \mathbb{R}$ for each $(s, a)$ pair.
The actions with the highest preferences in each state are given the highest probabilities of being selected, for example, according to an exponential soft-max distribution:

\[\boxed{\pi(a \vert s, \boldsymbol{\theta}) \doteq \frac{e^{h(s,a,\boldsymbol{\theta})}}{\sum_b e^{h(s,b,\boldsymbol{\theta})}}}\]

This kind of policy parametrization is called softmax in action preferences.
The action preferences can be parametrized arbitrarily, they can be computed by a deep artificial neural network (ANN) or could simply be linear in features:

\[h(s, a, \boldsymbol{\theta}) = \boldsymbol{\theta}^T \mathbf{x}(s, a)\]

Advantages of softmax in action preferences policy parametrization:

The approx. policy can approach a deterministic policy, unlike $\varepsilon$-greedy action selection.
It enables the selection of actions with arbitrary probabilities, which is useful for cases that require near-optimal stochastic policy (e.g. significant function approximation).
- E.g. useful in environments with imperfect information (e.g., card games) where it is optimal to act stochastically, such as when bluffing in Poker; it is important to do so randomly to unnerve/confuse an opponent.
The policy from policy parametrization may be a simpler function to approximate than the action-value function.
Often the most important reason for using a policy-based learning method is that it is a good way to inject prior knowledge about the desired form of the policy into the RL system.

13.2 The Policy Gradient Theorem

Policy-gradient methods have stronger convergence guarantees compared to action-value methods due to smooth continuity in change of action probabilities as a function of the learned parameter (gradient ascent).
Let’s consider the episodic performance measure, which is the value of the start state of the episode:

\[J(\boldsymbol{\theta}) \doteq v_{\pi_{\boldsymbol{\theta}}}(s_0)\] \[\begin{aligned} \text{where} \\ v_{\pi_{\boldsymbol{\theta}}} &\equiv \text{the true value function for } \pi_{\boldsymbol{\theta}}\text{, the policy determined by } \boldsymbol{\theta} \\ s_0 &\equiv \text{some non-random state} \end{aligned}\]

How can we estimate the performance gradient w.r.t. the policy parameter when the gradient depends on the unknown effect of policy changes on the state distribution?
- The theoretical answer is the policy gradient theorem.
The policy gradient theorem, which provides an analytic expression for the performance gradient w.r.t. the policy parameter, for the episodic case establishes that:

\[\boxed{\nabla J(\boldsymbol{\theta}) \propto \sum_s \mu(s) \sum_a q_\pi(s,a)\, \nabla\pi(a \vert s, \boldsymbol{\theta})}\] \[\begin{aligned} \text{where} \quad &\\ \nabla J(\boldsymbol{\theta}) &\equiv \text{column vectors of partial derivatives w.r.t. the components of } \boldsymbol{\theta} \\ \mu &\equiv \text{the on-policy distribution under the policy } \pi \\ &\phantom{{}\equiv{}} \text{(in the episodic case, proportionality constant is the average length of an episode;} \\ &\phantom{{}\equiv{}} \text{in the continuing case it is 1)} \end{aligned}\]

Proof: Policy Gradient Theorem (Episodic Case)

Let’s prove the policy gradient theorem from first principles using elementary calculus.

Note: To keep the notation simple, we leave it implicit in all cases that $\pi = f(\boldsymbol{\theta})$ and all gradients $\nabla[…]$ are also implicit w.r.t. $\boldsymbol{\theta}$.

\[\begin{align*} \nabla v_\pi(s) &= \nabla \!\left[\sum_a \pi(a \vert s)\, q_\pi(s,a)\right], \quad \text{for all } s \in \mathcal{S} \\ &= \sum_a \!\Biggl[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s)\, \nabla q_\pi(s,a)\Biggr] \\ &= \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s)\, \nabla \sum_{s',r} p(s',r \vert s,a)\!\left(r + v_\pi(s')\right)\right] \\ &= \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s) \sum_{s'} p(s' \vert s,a)\, \nabla v_\pi(s')\right] \\ &= \sum_a \Biggl[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s) \sum_{s'} p(s' \vert s,a) \sum_{a'} \Biggl[\nabla\pi(a' \vert s')\, q_\pi(s',a') \\ &\qquad\qquad + \pi(a' \vert s') \sum_{s''} p(s'' \vert s',a')\, \nabla v_\pi(s'')\Biggr]\Biggr] \end{align*}\]

Unrolling this recursion:

\[\boxed{\nabla v_\pi(s) = \sum_{x \in \mathcal{S}} \sum_{k=0}^{\infty} \Pr(s \to x, k, \pi) \sum_a \nabla\pi(a \vert x)\, q_\pi(x,a)}\] \[\begin{aligned} \Pr(s \to x, k, \pi) &\equiv \text{probability of transitioning from state } s \text{ to state } x \text{ in } k \text{ steps under policy } \pi \end{aligned}\]

Then:

\[\begin{align*} \nabla J(\boldsymbol{\theta}) &= \nabla v_\pi(s_0) \\ &= \sum_s \!\left(\sum_{k=0}^{\infty} \Pr(s_0 \to s, k, \pi)\right) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \\ &= \sum_s \eta(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \\ &= \sum_{s'} \eta(s') \sum_s \frac{\eta(s)}{\sum_{s'} \eta(s')} \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \\ &= \sum_{s'} \eta(s') \sum_s \mu(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \end{align*}\] \[\boxed{\nabla J(\boldsymbol{\theta}) \propto \sum_s \mu(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a)}\]

13.3 REINFORCE: Monte Carlo Policy Gradient

REINFORCE is our first policy gradient algorithm.
The goal/strategy is to find a way to get samples such that the expectation of the sample gradient is proportional to the actual performance gradient as a function of the parameters.
The sample gradients need to only be proportional to the performance gradient because any proportionality constant can be absorbed into the step size $\alpha$.
The policy gradient theorem gives an exact expression proportional to the gradient; so all that is needed is a way of sampling whose expectation equals or approximates this expression.
Recall the RHS of the policy gradient theorem is a sum over states weighted by how often the states occur under the target policy $\pi$:

\[\begin{align*} \nabla J(\boldsymbol{\theta}) &\propto \sum_s \mu(s) \sum_a q_\pi(s,a)\, \nabla\pi(a \vert s, \boldsymbol{\theta}) \\ &= \mathbb{E}_\pi \!\left[\sum_a q_\pi(S_t, a)\, \nabla\pi(a \vert S_t, \boldsymbol{\theta})\right] \end{align*}\]

We get REINFORCE by replacing the sum over the random variable’s possible values by an expectation under $\pi$, and then sampling the expectation:

\[\begin{align*} \nabla J(\boldsymbol{\theta}) &\propto \mathbb{E}_\pi \!\left[\sum_a \pi(a \vert S_t, \boldsymbol{\theta})\, q_\pi(S_t, a)\, \frac{\nabla\pi(a \vert S_t, \boldsymbol{\theta})}{\pi(a \vert S_t, \boldsymbol{\theta})}\right] \\ &= \mathbb{E}_\pi \!\left[q_\pi(S_t, A_t)\, \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta})}{\pi(A_t \vert S_t, \boldsymbol{\theta})}\right] \quad \text{(replacing } a \text{ by the sample } A_t \sim \pi\text{)} \\ &= \mathbb{E}_\pi \!\left[G_t\, \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta})}{\pi(A_t \vert S_t, \boldsymbol{\theta})}\right] \quad \text{(because } \mathbb{E}_\pi[G_t \vert S_t, A_t] = q_\pi(S_t, A_t)\text{)} \end{align*}\]

The last expression is the required expression; a quantity that can be sampled on each time step whose expectation is proportional to the gradient. This leads to the REINFORCE update:

\[\boxed{\boldsymbol{\theta}_{t+1} \doteq \boldsymbol{\theta}_t + \alpha G_t \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}{\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}}\]

This update is intuitively coherent:
- The gradient term represents the direction in parameter space that most increases the probability of repeating/taking again the same action $A_t$ in the future.
- The update moves the parameter the most in the directions that favour actions that yield the highest return.
- The update ensures a balancing act (lower impact) of frequently selected actions with lower returns.
REINFORCE has good convergence properties but may be of high variance and slow learning as a MC method.

REINFORCE: Monte Carlo Policy Gradient Control (Episodic) for $\pi_*$

\[\boxed{ \begin{aligned} &\textbf{Input: } \text{a differentiable policy parametrization } \pi(a \vert s, \boldsymbol{\theta}) \\ &\textbf{Algorithm parameter: } \text{step size } \alpha > 0 \\ &\textbf{Initialize } \text{policy parameter } \boldsymbol{\theta} \in \mathbb{R}^{d'} \text{ (e.g., to } \mathbf{0}\text{)} \\ &\textbf{Loop forever } \text{(for each episode):} \\ &\quad \text{Generate an episode } S_0, A_0, R_1, \ldots, S_{T-1}, A_{T-1}, R_T \text{ following } \pi(\cdot \vert \cdot, \boldsymbol{\theta}) \\ &\quad \textbf{Loop for each step of the episode } t = 0, 1, 2, \ldots, T-1\text{:} \\ &\qquad G \leftarrow \sum_{k=t+1}^{T} \gamma^{k-t-1} R_k \\ &\qquad \boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha\gamma^t G\, \nabla \ln \pi(A_t \vert S_t, \boldsymbol{\theta}) \end{aligned} }\]

13.4 REINFORCE with Baseline

The policy gradient theorem can be generalized to include a comparison of the action value to an arbitrary baseline $b(s)$:

\[\nabla J(\boldsymbol{\theta}) \propto \sum_s \mu(s) \sum_a \!\left[q_\pi(s,a) - b(s)\right] \nabla\pi(a \vert s, \boldsymbol{\theta})\]

The baseline can be any function that is not dependent on the action $a$:

\[\sum_a b(s)\, \nabla\pi(a \vert s, \boldsymbol{\theta}) = b(s)\, \nabla \sum_a \pi(a \vert s, \boldsymbol{\theta}) = b(s)\, \nabla 1 = 0\]

Now the REINFORCE update with baseline is:

\[\boxed{\boldsymbol{\theta}_{t+1} \doteq \boldsymbol{\theta}_t + \alpha \!\left[G_t - b(S_t)\right] \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}{\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}}\]

The Baseline REINFORCE update reduces the variance despite leaving the expected update value unchanged, which speeds up learning.
One natural choice for the baseline is an estimate of the state value $\hat{v}(S_t, \mathbf{w})$.
Since REINFORCE is a Monte Carlo method for learning the policy parameter $\boldsymbol{\theta}$, it’s natural to also use a Monte Carlo method to learn the state-value weights $\mathbf{w}$.

REINFORCE with Baseline (Episodic), for estimating $\pi_{\boldsymbol{\theta}} \approx \pi_*$

\[\boxed{ \begin{aligned} &\textbf{Input: } \text{a differentiable policy parametrization } \pi(a \vert s, \boldsymbol{\theta}) \\ &\textbf{Input: } \text{a differentiable state-value function parametrization } \hat{v}(s, \mathbf{w}) \\ &\textbf{Algorithm parameters: } \text{step sizes } \alpha^{\boldsymbol{\theta}} > 0,\ \alpha^{\mathbf{w}} > 0 \\ &\textbf{Initialize } \text{policy parameter } \boldsymbol{\theta} \in \mathbb{R}^{d'} \text{ and state-value weights } \mathbf{w} \in \mathbb{R}^d \text{ (e.g., to } \mathbf{0}\text{)} \\ &\textbf{Loop forever } \text{(for each episode):} \\ &\quad \text{Generate an episode } S_0, A_0, R_1, \ldots, S_{T-1}, A_{T-1}, R_T \text{ following } \pi(\cdot \vert \cdot, \boldsymbol{\theta}) \\ &\quad \textbf{Loop for each step of the episode } t = 0, 1, \ldots, T-1\text{:} \\ &\qquad G \leftarrow \sum_{k=t+1}^{T} \gamma^{k-t-1} R_k \\ &\qquad \delta \leftarrow G - \hat{v}(S_t, \mathbf{w}) \\ &\qquad \mathbf{w} \leftarrow \mathbf{w} + \alpha^{\mathbf{w}}\, \delta\, \nabla \hat{v}(S_t, \mathbf{w}) \\ &\qquad \boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha^{\boldsymbol{\theta}} \gamma^t\, \delta\, \nabla \ln \pi(A_t \vert S_t, \boldsymbol{\theta}) \end{aligned} }\]

This algorithm has 2 step sizes $\alpha^{\boldsymbol{\theta}}$ and $\alpha^{\mathbf{w}}$.
Choosing $\alpha^{\mathbf{w}}$ is relatively easy; in the linear case a good rule of thumb is (see Section 9.6):

\[\alpha^{\mathbf{w}} = \frac{0.1}{\mathbb{E}\!\left[\|\nabla \hat{v}(S_t, \mathbf{w})\|_\mu^2\right]}\]

Choosing $\alpha^{\boldsymbol{\theta}}$ is much less clear since its best value depends on the range of variation of the rewards and on the policy parametrization.

13.5 Actor-Critic Methods

REINFORCE with baseline cannot be used to evaluate actions because its state-value function only estimates the value of the 1st state of each state transition (1st state to 2nd state), which serves as a baseline of what to expect for the subsequent return to be, prior to the actual transition’s action.
Actor-critic methods, however, apply the state-value function to the 2nd state of the transition thereby estimating its value and thus evaluating the action.
The policy is the actor that maps states to actions, while the state-value function used to assess actions in this way is the critic.
The estimated value of the 2nd state, when discounted & added to the reward, yields the one-step return, $G_{t:t+1}$.

13.5.1 One-Step Actor-Critic Methods

One-step actor-critic methods replace the REINFORCE full return with the one-step return and use a learned state-value function as the baseline as follows:

\[\begin{align*} \boldsymbol{\theta}_{t+1} &\doteq \boldsymbol{\theta}_t + \alpha \!\left[G_{t:t+1} - \hat{v}(S_t, \mathbf{w})\right] \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}{\pi(A_t \vert S_t, \boldsymbol{\theta}_t)} \\ &= \boldsymbol{\theta}_t + \alpha \!\left[R_{t+1} + \gamma \hat{v}(S_{t+1}, \mathbf{w}) - \hat{v}(S_t, \mathbf{w})\right] \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}{\pi(A_t \vert S_t, \boldsymbol{\theta}_t)} \\ &= \boldsymbol{\theta}_t + \alpha\, \delta_t\, \frac{\nabla\pi(A_t \vert S_t, \boldsymbol{\theta}_t)}{\pi(A_t \vert S_t, \boldsymbol{\theta}_t)} \end{align*}\]

The semi-gradient TD(0) could serve as the natural state-value function learning method.
PROS: simple, fully online & incremental.
It is analogous to TD(0), Sarsa(0) & Q-learning.

One-Step Actor-Critic (Episodic), for estimating $\pi_{\boldsymbol{\theta}} \approx \pi_*$

\[\boxed{ \begin{aligned} &\textbf{Inputs: } \text{a differentiable policy } \pi(a \vert s, \boldsymbol{\theta}) \text{ and state-value function } \hat{v}(s, \mathbf{w}) \text{ parametrization} \\ &\textbf{Parameters: } \text{step sizes } \alpha^{\boldsymbol{\theta}} > 0,\ \alpha^{\mathbf{w}} > 0 \\ &\textbf{Initialize } \text{policy parameter } \boldsymbol{\theta} \in \mathbb{R}^{d'} \text{ and state-value weights } \mathbf{w} \in \mathbb{R}^d \text{ (e.g., to } \mathbf{0}\text{)} \\ &\textbf{Loop forever } \text{(for each episode):} \\ &\quad \text{Initialize } S \text{ (1st state of episode)} \\ &\quad I \leftarrow 1 \\ &\quad \textbf{Loop while } S \text{ is not terminal (for each time step):} \\ &\qquad A \sim \pi(\cdot \vert S, \boldsymbol{\theta}) \\ &\qquad \text{Take action } A\text{, observe } S', R \\ &\qquad \delta \leftarrow R + \gamma \hat{v}(S', \mathbf{w}) - \hat{v}(S, \mathbf{w}) \quad \text{(if } S' \text{ is terminal, then } \hat{v}(S', \mathbf{w}) \doteq 0\text{)} \\ &\qquad \mathbf{w} \leftarrow \mathbf{w} + \alpha^{\mathbf{w}}\, \delta\, \nabla \hat{v}(S, \mathbf{w}) \\ &\qquad \boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha^{\boldsymbol{\theta}} I\, \delta\, \nabla \ln \pi(A \vert S, \boldsymbol{\theta}) \\ &\qquad I \leftarrow \gamma I \\ &\qquad S \leftarrow S' \end{aligned} }\]

13.5.2 Actor-Critic with Eligibility Traces

We generalize to the forward view of $n$-step methods and then to a $\lambda$-return algorithm.
We replace the one-step return $G_{t:t+1}$ by $G_{t:t+n}$ or $G_t^\lambda$ respectively for either $n$-step or $\lambda$-return.
The backward view of the $\lambda$-return algorithm uses eligibility traces for the actor and the critic.

Actor-Critic with Eligibility Traces (Episodic), for estimating $\pi_{\boldsymbol{\theta}} \approx \pi_*$

\[\boxed{ \begin{aligned} &\textbf{Inputs: } \text{a differentiable policy and state-value function parametrization: } \pi(a \vert s, \boldsymbol{\theta}),\ \hat{v}(s, \mathbf{w}) \\ &\textbf{Parameters: } \text{trace-decay rates } \lambda^{\boldsymbol{\theta}} \in [0,1],\ \lambda^{\mathbf{w}} \in [0,1]\text{; step sizes } \alpha^{\boldsymbol{\theta}} > 0,\ \alpha^{\mathbf{w}} > 0 \\ &\textbf{Initialize } \text{policy parameter } \boldsymbol{\theta} \in \mathbb{R}^{d'} \text{ and state-value weights } \mathbf{w} \in \mathbb{R}^d \text{ (e.g., to } \mathbf{0}\text{)} \\ &\textbf{Loop forever } \text{(for each episode):} \\ &\quad \text{Initialize } S \text{ (1st state of episode)} \\ &\quad \mathbf{z}^{\boldsymbol{\theta}} \leftarrow \mathbf{0} \quad \text{($d'$-component eligibility trace vector)} \\ &\quad \mathbf{z}^{\mathbf{w}} \leftarrow \mathbf{0} \quad \text{($d$-component eligibility trace vector)} \\ &\quad \textbf{Loop while } S \text{ is not terminal (for each time step):} \\ &\qquad A \sim \pi(\cdot \vert S, \boldsymbol{\theta}) \\ &\qquad \text{Take action } A\text{, observe } S', R \quad \text{(if } S' \text{ is terminal, then } \hat{v}(S', \mathbf{w}) \doteq 0\text{)} \\ &\qquad \delta \leftarrow R + \gamma \hat{v}(S', \mathbf{w}) - \hat{v}(S, \mathbf{w}) \\ &\qquad \mathbf{z}^{\mathbf{w}} \leftarrow \gamma\lambda^{\mathbf{w}} \mathbf{z}^{\mathbf{w}} + \nabla \hat{v}(S, \mathbf{w}) \\ &\qquad \mathbf{z}^{\boldsymbol{\theta}} \leftarrow \gamma\lambda^{\boldsymbol{\theta}} \mathbf{z}^{\boldsymbol{\theta}} + I\, \nabla \ln \pi(A \vert S, \boldsymbol{\theta}) \\ &\qquad \mathbf{w} \leftarrow \mathbf{w} + \alpha^{\mathbf{w}}\, \delta\, \mathbf{z}^{\mathbf{w}} \\ &\qquad \boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha^{\boldsymbol{\theta}}\, \delta\, \mathbf{z}^{\boldsymbol{\theta}} \\ &\qquad I \leftarrow \gamma I \\ &\qquad S \leftarrow S' \end{aligned} }\]

13.6 Policy Gradient for Continuing Problems

From Section 10.3 on continuing problems, lack of episode boundaries requires a new performance measure definition in terms of the average rate of reward per time step:

\[\begin{align*} J(\boldsymbol{\theta}) \doteq r(\pi) &\doteq \lim_{h \to \infty} \frac{1}{h} \sum_{t=1}^{h} \mathbb{E}\!\left[R_t \vert S_0, A_{0:t-1} \sim \pi\right] \\ &= \lim_{t \to \infty} \mathbb{E}\!\left[R_t \vert S_0, A_{0:t-1} \sim \pi\right] \\ &= \sum_s \mu(s) \sum_a \pi(a \vert s) \sum_{s',r} p(s',r \vert s,a)\, r \end{align*}\] \[\begin{aligned} \text{where} \\ \mu(s) &\doteq \lim_{t \to \infty} \Pr\{S_t = s \vert A_{0:t} \sim \pi\} \equiv \text{steady state distribution under } \pi\text{,} \\ &\phantom{{}\doteq{}} \text{assumed to exist and be independent of } S_0 \textbf{ [ergodicity assumption]} \end{aligned}\] \[\sum_s \mu(s) \sum_a \pi(a \vert s, \boldsymbol{\theta})\, p(s' \vert s,a) = \mu(s') \quad \text{for all } s' \in S \quad \text{(ergodicity)}\]

In the continuing case, we define values $v_{\pi}(s) \doteq \mathbb{E}_{\pi}[G_t \vert S_t = s]$ and $q_{\pi}(s,a) \doteq \mathbb{E}_{\pi}[G_t \vert S_t = s, A_t = a]$ w.r.t. the differential return:

\[G_t \doteq R_{t+1} - r(\pi) + R_{t+2} - r(\pi) + R_{t+3} - r(\pi) + \ldots\]

Proof: Policy Gradient Theorem (Continuing Case)

Leave the notation implicit in all cases that $\pi = f(\boldsymbol{\theta})$ and that the gradients $\nabla[…]$ are w.r.t. $\boldsymbol{\theta}$. In the continuing case $J(\boldsymbol{\theta}) = r(\pi)$, and $v_\pi$ & $q_\pi$ denote values w.r.t. the differential return.

\[\begin{align*} \nabla v_\pi(s) &= \nabla \!\left[\sum_a \pi(a \vert s)\, q_\pi(s,a)\right], \quad \text{for all } s \in \mathcal{S} \\ &= \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s)\, \nabla q_\pi(s,a)\right] \\ &= \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s)\, \nabla \sum_{s',r} p(s',r \vert s,a)\!\left(r - r(\boldsymbol{\theta}) + v_\pi(s')\right)\right] \\ &= \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s)\!\left[-\nabla r(\boldsymbol{\theta}) + \sum_{s'} p(s' \vert s,a)\, \nabla v_\pi(s')\right]\right] \end{align*}\] \[\nabla r(\boldsymbol{\theta}) = \sum_a \!\left[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s) \sum_{s'} p(s' \vert s,a)\, \nabla v_\pi(s')\right] - \nabla v_\pi(s)\]

Since $\nabla J(\boldsymbol{\theta})$ does not depend on $s$, we sum over all $s \in \mathcal{S}$ weighted by $\mu(s)$ (because $\sum_s \mu(s) = 1$):

\[\begin{align*} \nabla J(\boldsymbol{\theta}) &= \sum_s \mu(s) \Bigl(\sum_a \Bigl[\nabla\pi(a \vert s)\, q_\pi(s,a) + \pi(a \vert s) \sum_{s'} p(s' \vert s,a)\, \nabla v_\pi(s')\Bigr] - \nabla v_\pi(s)\Bigr) \\ &= \sum_s \mu(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \\ &\quad + \sum_{s'} \underbrace{\sum_s \mu(s) \sum_a \pi(a \vert s)\, p(s' \vert s,a)}_{\mu(s')} \nabla v_\pi(s') - \sum_s \mu(s)\, \nabla v_\pi(s) \\ &= \sum_s \mu(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) + \sum_{s'} \mu(s')\, \nabla v_\pi(s') - \sum_s \mu(s)\, \nabla v_\pi(s) \\ &= \sum_s \mu(s) \sum_a \nabla\pi(a \vert s)\, q_\pi(s,a) \qquad \text{Q.E.D.} \end{align*}\]

Actor-Critic with Eligibility Traces (Continuing), for estimating $\pi_{\boldsymbol{\theta}} \approx \pi_*$

\[\boxed{ \begin{aligned} &\textbf{Inputs: } \pi(a \vert s, \boldsymbol{\theta}),\ \hat{v}(s, \mathbf{w}) \\ &\textbf{Parameters: } \lambda^{\mathbf{w}} \in [0,1],\ \lambda^{\boldsymbol{\theta}} \in [0,1],\ \alpha^{\mathbf{w}} > 0,\ \alpha^{\boldsymbol{\theta}} > 0,\ \alpha^{\bar{R}} > 0 \\ &\textbf{Initialize } \bar{R} \in \mathbb{R} \text{ (e.g., to 0)},\ \mathbf{w} \in \mathbb{R}^d\ \&\ \boldsymbol{\theta} \in \mathbb{R}^{d'} \text{ (e.g., to } \mathbf{0}\text{)},\ S \in \mathcal{S} \text{ (e.g., to } s_0\text{)} \\ &\mathbf{z}^{\mathbf{w}} \leftarrow \mathbf{0};\ \mathbf{z}^{\boldsymbol{\theta}} \leftarrow \mathbf{0} \\ &\textbf{Loop forever } \text{(for each time step):} \\ &\quad A \sim \pi(\cdot \vert S, \boldsymbol{\theta}) \\ &\quad \text{Take action } A\text{, observe } S', R \\ &\quad \delta \leftarrow R - \bar{R} + \hat{v}(S', \mathbf{w}) - \hat{v}(S, \mathbf{w}) \\ &\quad \bar{R} \leftarrow \bar{R} + \alpha^{\bar{R}}\, \delta \\ &\quad \mathbf{z}^{\mathbf{w}} \leftarrow \lambda^{\mathbf{w}} \mathbf{z}^{\mathbf{w}} + \nabla \hat{v}(S, \mathbf{w}) \\ &\quad \mathbf{z}^{\boldsymbol{\theta}} \leftarrow \lambda^{\boldsymbol{\theta}} \mathbf{z}^{\boldsymbol{\theta}} + \nabla \ln \pi(A \vert S, \boldsymbol{\theta}) \\ &\quad \mathbf{w} \leftarrow \mathbf{w} + \alpha^{\mathbf{w}}\, \delta\, \mathbf{z}^{\mathbf{w}} \\ &\quad \boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha^{\boldsymbol{\theta}}\, \delta\, \mathbf{z}^{\boldsymbol{\theta}} \\ &\quad S \leftarrow S' \end{aligned} }\]

13.7 Policy Parametrization for Continuous Actions

Policy gradient methods enable us to handle large (and even continuous) action spaces by learning the statistics of the probability distribution instead of computing learned probabilities for each of the many actions.
The probability density function for the normal distribution is conventionally written as:

\[p(x) \doteq \frac{1}{\sigma\sqrt{2\pi}} \exp\!\left(-\frac{(x-\mu)^2}{2\sigma^2}\right)\] \[\begin{aligned} \text{where} \\ \mu &\equiv \text{mean of the normal distribution} \\ \sigma &\equiv \text{standard deviation of the normal distribution} \end{aligned}\]

Probability Density Function (PDF)

The probability density functions for several different means and standard deviations are shown above.

$p(x)$ is the density of the probability at $x$, and not the probability. It can be greater than 1; it is the total area under the graph that must sum to 1.
To get the probability of $x$ falling within a range, take the integral under $p(x)$ for that specific range of $x$ values.
The policy parametrization, defining the policy as the normal probability density over a real-valued scalar action $a$ with mean & standard deviation given by state-dependent, parametric function approximators, is as follows:

\[\pi(a \vert s, \boldsymbol{\theta}) \doteq \frac{1}{\sigma(s, \boldsymbol{\theta})\sqrt{2\pi}} \exp\!\left(-\frac{(a - \mu(s, \boldsymbol{\theta}))^2}{2\sigma(s, \boldsymbol{\theta})^2}\right)\] \[\begin{aligned} \text{where} \\ \mu &: S \times \mathbb{R}^{d'} \to \mathbb{R} \equiv \text{mean function approximator} \\ \sigma &: S \times \mathbb{R}^{d'} \to \mathbb{R}^{+} \equiv \text{standard deviation function approximator} \end{aligned}\]

The approximators need a representation form, so we split the policy’s parameter vector into 2 parts, $\boldsymbol{\theta} = [\boldsymbol{\theta}_{\mu}, \boldsymbol{\theta}_{\sigma}]^T$, one for mean approximation and the other for standard deviation approximation.
The mean can be approximated as a linear function while the standard deviation can be approximated as the exponential of a linear function:

\[\begin{aligned} \mu(s, \boldsymbol{\theta}) &\doteq \boldsymbol{\theta}_\mu^T \mathbf{x}_\mu(s) \\ \sigma(s, \boldsymbol{\theta}) &\doteq \exp\!\left(\boldsymbol{\theta}_\sigma^T \mathbf{x}_\sigma(s)\right) \end{aligned}\]

13.8 Summary

Going from action-value methods to parametrized policy methods that take actions without consulting action-value estimates.
More specifically, policy gradient methods update the policy parameter on each step in the direction of an estimate of the performance gradient w.r.t. the policy parameter.
Advantages of parametrized policy methods over $\varepsilon$-greedy & action-value methods:
- They can learn specific probabilities for taking actions.
- They can learn appropriate levels of exploration & approach deterministic policies asymptotically.
- They can naturally handle continuous action spaces.
- Important theoretical advantage over action-value methods in the form of the policy gradient theorem.
REINFORCE uses the policy gradient theorem with Monte Carlo returns.
Addition of a state-value function as a baseline with REINFORCE reduces its variance without introducing bias and speeds up learning.
If the state-value function is used to criticize/assess the policy’s action selections, then the value function is called a critic and the policy is called an actor. Overall this is referred to as the actor-critic method.
The critic introduces bias into the actor’s gradient estimates, but is still often desirable for the same reason that bootstrapping TD methods are superior to Monte Carlo methods (significant variance reduction).
Overall, policy-gradient methods provide a significantly different set of strengths & weaknesses than action-value methods.

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026RLsuttonBartoCh13notes,
  title   = "Sutton & Barto, Ch. 13: Policy Gradient Methods (Personal Notes)",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "May",
  url     = "https://chizkidd.github.io/2026/05/07/rl-sutton-barto-notes-ch013/"
}

Transformer Architecture Explained: Self-Attention, Encoders, and Decoders

Fri, 17 Apr 2026 00:00:00 +0000

Transformers are a sequence-to-sequence model: given an input sequence, produce an output sequence.
Architecture: an Encoder processes the input; a Decoder generates the output autoregressively.

\[\text{(En) "I am sorry"} \xrightarrow{\text{Encoder}} \xrightarrow{\text{Decoder}} \texttt{<start>}\ \text{Je suis désolé}\ \texttt{<end>}\]

Autoregressive: the decoder generates one token at a time, conditioning on all previously generated tokens.

Input Text Sequence Representation
Encoders
From MLP to Attention
Self-Attention & Multi-Head Attention
Decoders
Masked Attention
Encoder-Decoder Cross Attention

Appendix

Citation

Transformer architecture

Transformer architecture (Source: Attention is all you need)

Input Text Sequence Representation

Tokenization

Two approaches to representing input text:

1. One-hot encoding

No semantic similarity or meaning of words encoded.

2. Token embedding

Encodes semantic similarity between words.
Embedding matrix is learned (Lookup Table).
Each token embedding is stored as a column vector.

\[\begin{aligned} \text{where} \\ W_E &\equiv \text{embedding matrix, } d \times \text{\#tokens} \\ d &\equiv \text{embedding dimension} \end{aligned}\]

Why We Need Context

Many words have different meanings in different contexts:
- “I bought an apple & an orange”
- “I bought an apple watch”
We need to rely on context to resolve the ambiguity.

Encoders

The encoder pipeline:

\[\text{Input} \to \text{token} + \text{POS embed} \to \text{Norm} \to \text{MHA(self)} \to \text{Add} \to \text{Norm} \to \text{FFN} \to \text{Add} \to \text{Output}\]

Input tokens are embedded using $W_E$ and combined with positional encodings to produce the input matrix $X$.
The architecture stacks: Multi-Head Self-Attention + residual, then FeedForward Network + residual.

From MLP to Attention

MLP only

MLP stands for Multilayer Perceptron
No contextual information; each token is processed independently.

Concatenation of nearby token embeddings before MLP

Need a sufficiently large window to cover the entire input sequence.
Cannot handle variable sequence lengths.
Requires many model parameters.

Attention

Use token similarity to determine the relevance of each token to every other token by performing a dot product.
Allows the model to dynamically weight which parts of the input are relevant for each position.

Self-Attention & Multi-Head Attention

Self-Attention

\[Q = XW_Q, \quad K = XW_K, \quad V = XW_V\] \[\text{head}_i = \text{Softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right)V = \text{Attention}(Q, K, V)\]

Multi-Head Attention (MHA)

\[\text{MHA}(X) = \text{multi-head}(Q, K, V) = \text{Concat}(h_1, h_2, \ldots, h_H)\, W_O = Z\] \[\begin{aligned} \text{where} \\ h_i &\equiv i\text{-th attention head} \\ W_O &\equiv \text{output projection matrix} \\ Z &\equiv \text{final output} \end{aligned}\]

Decoders

The decoder pipeline:

\[\text{Masked MHA} \to \text{Cross-Attn} \to \text{FFN} \to \text{Linear} \to \text{Softmax} \to \text{Output}\]

The decoder takes as input the previously generated tokens $(z_i)$ along with their positional encodings $(p_i)$, and at each step attends both to itself (masked) and to the encoder output (cross attention).

Masked Attention

In the decoder, we use masked (causal) self-attention to prevent the decoder from attending to future tokens:

\[\text{Masked Attn}(Q, K, V) = \text{Softmax}\!\left(\frac{QK^T}{\sqrt{d_k}} + M\right)V\] \[\begin{aligned} \text{where} \\ M &\equiv \text{lookahead mask} \end{aligned}\] \[M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix}\]

Adding $-\infty$ to future positions drives their softmax weights to zero, ensuring position $i$ can only attend to positions $\leq i$.

Encoder-Decoder Cross Attention

The decoder queries the encoder output $E$ (encoder output) to incorporate source context:

\[\begin{aligned} Y' &= Y + \text{Masked MHA}(\text{Norm}(Y)) \\[6pt] Q &= Y' W_Q^{\text{dec}} \\ K^{\text{enc}} &= E W_K^{\text{enc}} \\ V^{\text{enc}} &= E W_V^{\text{enc}} \end{aligned}\] \[\begin{aligned} \text{where} \\ E &\equiv \text{encoder output} \end{aligned}\] \[\text{Cross Attn}(Y', E) = \text{Softmax}\!\left(\frac{Q K_{\text{enc}}^T}{\sqrt{d_k}}\right)V^{\text{enc}}\]

Then the rest of the decoder:

\[\begin{aligned} Y'' &= Y' + \text{Cross Attn}(\text{Norm}(Y'),\ E) \\ Y^O &= Y'' + \text{FFN}(\text{Norm}(Y'')) \\ D &= Y^O \\ \text{logits} &= DW^{\text{out}} + b \\ P(y_t) &= \text{Softmax}(\text{logits}_t) \end{aligned}\] \[\begin{aligned} \text{where} \\ D &\equiv \text{decoder output} \\ W^{\text{out}} &\equiv \text{output projection to vocabulary} \\ P(y_t) &\equiv \text{probability distribution over next token} \end{aligned}\]

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026transformers,
  title   = "Transformers",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Apr",
  url     = "https://chizkidd.github.io/2026/04/17/transformers/"
}

SAM 2 Explained: Meta's Promptable Visual Segmentation Model

Fri, 17 Apr 2026 00:00:00 +0000

Meta’s unified model for promptable image and video segmentation.
- A foundation model for solving promptable visual segmentation in images & videos.
Built a data engine to collect the largest video segmentation dataset to date.
Model: Simple transformer architecture with streaming memory for real-time video processing.
Trained on a wide range of tasks: video segmentation and image segmentation.
The paper can be found here.

1. Introduction
2. Related Work
3. Task: Promptable Visual Segmentation (PVS)
4. Model
5. Data
6. Zero-Shot Experiments
7. Comparison to SOTA in Semi-Supervised VOS
8. Conclusion
9. Discussion

Appendix

Citation
References

1. Introduction

Why video and not image?

Image is only a static snapshot of the real world; lacks motion information (temporal).
Video captures temporal information.
Many vital applications (robotics, AR/VR, autonomous vehicles) require temporal localization beyond image-level segmentation.
A universal visual segmentation system should be applicable to both images & videos.
Video segmentation aims to determine the spatio-temporal extent of entities, which presents unique challenges beyond those in images.
Significant changes in appearance encountered by entities & lower quality nature of videos than images present challenges for video segmentation.
SAM successfully solves image segmentation, but existing video segmentation models & datasets fall short in providing a comparable capability to “segment anything in videos.”
SAM 2: A unified model for video & image segmentation.
Promptable Visual Segmentation (PVS): Task that generalizes image segmentation to the video domain.
A data engine that generates training data via an in-the-loop model with annotators and produces the Segment Anything Video (SA-V) dataset.

Video Object Segmentation (VOS)
- Video augmentation datasets
Interactive Video Object Segmentation (iVOS)
Image Segmentation task, model and dataset
- Research Paper: Segment Anything (SA)

Segment Anything (Adapted from the Paper)

We aim to build a foundation model for segmentation by introducing three interconnected components: a promptable segmentation task, a segmentation model (SAM) that powers data annotation and enables zero-shot transfer to a range of tasks via prompt engineering, and a data engine for collecting SA-1B, our dataset of over 1 billion masks.

3. Task: Promptable Visual Segmentation (PVS)

\[\text{PVS} \longrightarrow \text{SAM 2} \longrightarrow \text{SA-V dataset}\]

PVS task allows providing prompts to the model on any frame of a video.
The interactive segmentation with SAM2 involves the steps below:
- SAM 2 is prompted on a single frame and responds instantly with a valid segmentation mask of the target object on this frame.
- SAM 2 then propagates the target object’s segment to multiple frames to form a masklet.
- Multiple initial prompts are received and propagated by the model to obtain the masklet of the object across the entire video, which leads to localization of the segmentation mask of the target on every single video frame.
- Additional prompts on any frame can be added to SAM 2 for segmentation mask refinement.
SAM 2 is applied as a data collection tool to the PVS task for building the SA-V dataset.
Model evaluation is done via simulation of interactive video segmentation scenarios across multiple frames in the conventional first-frame, limited, semi-supervised VOS setting, and for image segmentation on the SA benchmarks.

4. Model

SAM 2 is a generalization of SAM to the video (& image) domain. Essentially, it employs taking point, box & mask prompts on individual frames to define the spatial extent of the object to be segmented spatio-temporally.

Figure 1: The SAM 2 architecture. For a given frame, the segmentation prediction is conditioned on the current prompt and/or on previously observed memories. Video frames are processed in a streaming fashion by the image encoder, cross-attended to memories of the target object from previous frames stored in the memory bank, and decoded via the mask decoder (optionally prompted by the prompt encoder) to predict the segmentation mask for that frame. Finally, a memory encoder transforms the prediction and image encoder embeddings for use in future frames.

Components:

Image encoder: For real-time processing of arbitrarily long videos via a streaming, hierarchical approach.
Memory attention: Condition the current frame features on the past frames features and predictions or any new prompts.
Prompt encoder & mask decoder: Encode and send the input prompts to the mask decoder to predict the segmentation mask for the current frame.
Memory encoder: Generates a memory by downsampling, element-wise summation and light-weight convolution of the output mask which fuses the information.
Memory bank: Stores spatial memory information (from prompts) about past predictions for the target object in the video, and high-level semantic information of the object to segment based on each frame’s mask decoder output tokens.

Training:

The model is trained jointly on image and video data via interactive prompting simulation of the model, sequential sampling of 8 frames and random selection of 2 frames to prompt, with the goal of sequential and interactive prediction of the gound-truth masklet.

5. Data

The data engine was built, by an interactive in-the-loop model setup with human annotators, to collect a large and diverse video segmentation dataset to develop the “segemnt anything” capability in video. The data engine underwent 3 phases, each based on the level of model assistance provided to annotators.

5.1 Data Engine

Phase 1: SAM per frame
Phase 2: SAM + SAM 2
Phase 3: SAM 2
Quality verification including a separate set of QA annotators for the masklets (satisfactory or unsatisfactory)
Auto masklet generation enables the “anything capability” of the model by ensuring annotation diversity.
Analysis 1: Comparison in each data engine phase through a controlled experiment of the average annotation time per frame, the average percentage of manually edited frames per masklet, and the average number of clicks per clicked frame. For QA, the defined metric is called Phase 1 Mask Alignment Score, the percentage of masks whose IoU compared to the corresponding masks in Phase 1 (highest quality manual annotations) exceeds 0.75.
Analysis 2: Performance comparison of SAM 2 trained on the available data at the end of each phase keeping the number of iterations fixed, therefore measuring solely the impact of the additional data. Evaluation is done on the SA-V val dataset and 9 zero-shot benchmarks using the standard $J\&F$ accuracy metric.

Analysis Results

Phase 3 is 8.4× faster than Phase 1, has the lowest edited frame percentage and clicks per frame, and results in better alignment.
Consistent segmentation accuracy improvement from iteratively adding data from each data engine phase (1, 2 & 3) for both SA-V val set and 9 zero-shot benchmarks is observed.

5.2 SA-V Dataset

Videos (50.9K total, 54% indoor + 46% outdoor scenes, average duration of 14 secs)
Masklets (190.9K manual + 451.7K automatic, 642.6K total)
SA-V training, validation (293 masklets & 155 videos) & test splits (278 masklets & 150 videos)
Internal dataset (62.9K videos & 69.6K masklets annotated in Phase 2 & 3 for training; 96 videos & 189 masklets annotated using Phase 1 for testing)

6. Zero-Shot Experiments

Compare SAM 2 with previous work on zero-shot video & image tasks using the $J\&F$ accuracy metric.

6.1 Promptable Video Segmentation (PVS)

First evaluate PVS, which involves simulating an interactive setting that is akin to the user experience. Both evaluation settings below are done on 9 densely annotated zero-shot video datasets using $N_{click} = 3$ clicks per frame and are compared to 2 strong baselines based on 2 SOTA VOS models (XMem++ & Cutie)
- Offline evaluation: multiple passes made through a video for frame selection via the largest model error.
- Online evaluation: single forward pass for video frames’ annotation.

6.2 Semi-Supervised Video Object Segmentation

Evaluate the semi-supervised VOS setting with click, box or mask prompts only on the 1st frame of the video.
For click prompts, try either 1, 3, or 5 clicks on the 1st video frame (mIoU).
Comparison is done via $J\&F$ accuracy between SAM + XMem++, SAM + Cutie and SAM 2.

6.3 Image Segmentation

Evaluate SAM 2 on the Segment Anything task across 37 zero-shot datasets using 1-click and 5-click mIoUs.

7. Comparison to SOTA in Semi-Supervised VOS

Evaluate 2 versions of SAM 2 with varying image encoder sizes with different speed-vs-accuracy tradeoffs.
Comparison with existing SOTA models/methods via accuracy ($J\&F, G$) using standard protocols.
Evaluate existing work on the SA-V val & test sets which measure performance for open-world segments of “any” object class via $J\&F$ accuracy metric.

Comparison Results

SAM 2 performs well in accuracy for video segmentation based on first-frame ground-truth mask prompts.
SAM 2 performs significantly better on SA-V val/test.

8. Conclusion

Present a natural extension of Segment Anything into the video domain, based on 3 key aspects:

Extending the promptable segmentation task to video.
Equipping the SAM architecture to use memory when applied to video.
The diverse SA-V dataset for training & benchmarking video segmentation.

The authors believed SAM 2 marked a significant advancement in visual perception, positioning their contributions as milestones that will propel further research & applications.

9. Discussion

When I first read about SAM 2’s memory bank component, my immediate thought was: this feels familiar. Not because it copies prior work, but because it sits at an intersection I’ve been circling for a while: how do neural systems remember what matters, without remembering everything?

SAM 2 maintains object identity across video frames by storing two kinds of information in its memory bank:

Spatial feature maps from up to N recent frames and M prompted frames (stored in FIFO queues)
Object pointers: lightweight semantic vectors derived from the mask decoder’s output tokens

Memory attention then cross-attends over both when predicting the current frame. This design elegantly sidesteps the sequential bottleneck that plagued RNNs. Instead of compressing history into a single fixed-length hidden state, SAM 2 gives the model direct, position-invariant access to past representations. The “distance” between frame 1 and frame 200 is functionally the same as between frame 199 and 200. In that sense, attention acts as a temporal superhighway.

Distance vs. Capacity

Key distinction: RNNs encode history into a single vector (making long-range dependencies hard). SAM 2’s attention mechanism bypasses that bottleneck entirely, so distance becomes irrelevant. But attention solves the access problem, not the retention problem.

And that’s where the FIFO design matters. The paper explicitly states that the memory bank evicts the oldest frame once the queue is full, regardless of semantic importance.¹ This validates a subtle but critical observation: SAM 2’s forgetting mechanism is a fixed heuristic, not a learned one. The model doesn’t decide what to remember based on future tracking utility; it drops frames based on arrival time.

This creates a tangible trade-off between memory availability and diagnostic utility. Consider an object that:

Undergoes dramatic lighting change
Is occluded for dozens of frames
Reappears with significant appearance deformation

The frames that best capture its initial identity might be evicted long before it reappears. Meanwhile, newer frames with degraded or ambiguous representations linger in the queue simply because they arrived later. The “object pointers” partially mitigate this by storing lightweight semantic summaries, but they’re still bound by the same FIFO eviction policy. If the pointer for frame t is overwritten, the model loses not just the spatial map but also its high-level anchor.

flowchart LR A[FIFO Queue
time-based eviction] --> B[Eviction Policy
fixed heuristic] B --> C[Memory Attention
position-invariant access] C --> D{Tracking Outcome} D -->|Object reappears
with useful memory| E[✓ Success] D -->|Key frame evicted
or pointer overwritten| F[✗ Failure] style A fill:#e3f2fd style B fill:#fff3e0 style C fill:#e8f5e9 style E fill:#c8e6c9,stroke:#2e7d32 style F fill:#ffcdd2,stroke:#c62828

Figure 2: The memory bank’s fixed eviction policy (FIFO) interacts with attention’s position-invariant access. When evicted frames contain critical identity information, tracking fails; even if attention could theoretically retrieve them.

The paper’s handling of temporal position encoding reinforces this pragmatic trade-off. Temporal embeddings are injected into the N recent frames to capture short-term motion, but deliberately omitted from prompted frames due to sparse training signals and inference generalization concerns. This is a sound engineering decision, but it reveals a boundary: SAM 2 optimizes for stable, short-to-medium horizon tracking, not open-ended temporal reasoning.

Where This Fits in the Broader Literature

SAM 2’s memory bank isn’t operating in a vacuum. It shares conceptual DNA with several prior lines of work:

Neural Turing Machines introduced differentiable external memory with read/write heads, allowing networks to learn what to store and where to retrieve from ². SAM 2’s memory attention is a specialized, non-differentiable cousin: it retrieves, but doesn’t learn the eviction policy.
RETRO demonstrated that retrieval-augmented transformers can scale knowledge without scaling parameters, by querying a frozen corpus at inference ³. SAM 2 does something analogous for video: query a frozen buffer of past frames. The open question is whether that buffer should be learned, not fixed.
TimeSformer showed that spatiotemporal attention alone can handle video understanding without recurrent components ⁴. SAM 2 extends this by adding explicit memory, and also inherits TimeSformer’s assumption that all frames are equally worth attending to.

What SAM 2 adds is a practical instantiation of these ideas for promptable segmentation. It’s not trying to solve general memory-augmented reasoning; it’s solving “keep this object tracked, please.” That focus is its strength; and its limitation.

Counter Argument to the Paper's Memory Design

About memory design, I’d say: “FIFO is computationally cheap and training-stable, which makes sense for a production model. But from a research standpoint, it hardcodes a failure mode: important frames get evicted by time, not relevance. A differentiable memory controller or retrieval-augmented eviction policy could close that gap.”

Why This Matters

There’s a tendency in ML to treat architectural choices as purely engineering decisions. But memory management isn’t just about compute budgets, it’s also about what the model values. When SAM 2 drops a frame because it’s old, not because it’s uninformative, it’s making a silent claim: recency matters more than relevance.

That claim works well for many tracking scenarios. But it breaks down when objects reappear after long occlusions, or when appearance changes dramatically. In those cases, the model isn’t failing because attention is weak; it’s failing because the right information was never kept around to attend to.

This isn’t unique to SAM 2. It’s a fundamental tension in any system that must balance finite memory against infinite context. But because SAM 2 is a foundation model positioned for broad adoption, its design choices will influence how thousands of downstream applications handle temporal reasoning. Getting the memory story right matters.

So where does this leave us? SAM 2 is undoubtedly a milestone in promptable video segmentation. But its memory bank inadvertently frames a deeper research problem: attention removes the barrier of temporal distance, but leaves the bottleneck of memory management wide open.

Open Questions I’d Want to Explore

Learnable eviction: Could we replace FIFO with a lightweight, content-aware, learnable memory eviction mechanistic network that predicts which frames to retain based on tracking confidence, appearance stability, or semantic salience? Would SAM 2 maintain robust identity over arbitrarily long horizons? What’s the compute trade-off from an engineering system to a long-context reasoning model?
Pointer robustness: The object pointers are a clever compression trick, but they’re still overwritten by FIFO. Could we decouple pointer retention from spatial memory eviction?
Cross-video retrieval: RETRO retrieves from a corpus of documents; could SAM 2 retrieve from a corpus of past videos to bootstrap tracking of familiar objects?
Failure diagnostics: Can we design a probe that predicts when SAM 2 is likely to lose an object, based on memory bank state? That would be valuable for safety-critical applications.

I don’t have answers to these yet. But they feel like the right questions to ask if we want video models that don’t just see, but remember. SAM 2 shows we’ve mastered access to the past. The next step is mastering retention of what matters.

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026sam2,
  title   = "SAM 2: Segment Anything in Images & Videos",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Apr",
  url     = "https://chizkidd.github.io/2026/04/17/sam2/"
}

References

Ravi N, Gabeur V, Hu Y-T, et al. SAM 2: Segment Anything in Images and Videos. arXiv. 2024. ↩
Graves A, Wayne G, Danihelka I. Neural Turing Machines. arXiv. 2014. ↩
Borgeaud S, Mensch A, Hoffmann J, et al. Improving Language Models by Retrieving from Trillions of Tokens. arXiv. 2022. ↩
Bertasius G, Wang H, Torresani L. Is Space-Time Attention All You Need for Video Understanding?. arXiv. 2021. ↩

Muon Optimizer Explained: Newton-Schulz Orthogonalization Beyond Adam

Sat, 04 Apr 2026 00:00:00 +0000

Muon stands for MomentUm Orthogonalized by Newton-Schulz and was invented by Keller Jordan.
The key idea: Instead of applying Adam-style per-element adaptive updates to model parameters, Muon orthogonalizes the momentum matrix before using it as the update direction.

Adam Optimizer
Matrix Orthogonalization
Newton-Schulz 5 Iteration
Muon
QK-Clip
Multihead Latent Attention (MLA)
MuonClip

Appendix

Citation
References

Moonshot AI, creator of Kimi, pioneered the improvements to Muon (MuonClip) which I dive into in detail in this post. But they have an article on X with more academic details on why they chose Muon titled “Why We Chose Muon: Our Chain of Thought”¹ by Jianlin Su, the first author of RoPE (Rotary Position Embedding).

https://t.co/dxZnLxvPae
— Kimi.ai (@Kimi_Moonshot) March 7, 2025

This blog post is directly from my personal handwritten notes on studying Muon & MuonClip. I posted those notes on X.

My notes on the Muon optimizer by @kellerjordan0 and MuonClip by @Kimi_Moonshot which integrates Muon with weight decay, RMS matching, & QK-Clip.

"Using MuonClip, we successfully pre-trained Kimi K2 on 15.5 trillion tokens without a single loss spike." - @Kimi_Moonshot pic.twitter.com/KW84xOi51E
— Chiz (@latentchiz) April 5, 2026

Adam Optimizer

Let’s start by looking at Adam optimizer, the most common optimizer for training neural networks, which I cover in this article guide on common deep learning optimizers. Adam stands for Adaptive Moment Estimation.

It combines momentum and adaptive learning rates so the model not only remembers the direction it has been moving in, but also adjusts how big each step should be for every parameter. More specifcally, it combines momentum (first moment) and RMSProp (second moment) with bias corrections to handle noisy gradients and early training instability.² The forward pass looks like:

Gradient Descent

Gradient descent updates the model parameters by stepping in the direction that reduces the loss, using the gradient as a guide for how to adjust each parameter. The momentum and velocity terms refine this process by smoothing past gradients and scaling updates adaptively, which helps stabilize training and converge faster, especially in noisy or complex loss landscapes.

\[\begin{aligned} \text{simply} \quad & \Theta_i \leftarrow \Theta_{i-1} - \alpha \frac{\partial L}{\partial \Theta_i} \\[6pt] \text{momentum} \quad & M_i \leftarrow \beta_1 M_{i-1} + (1 - \beta_1) \frac{\partial L}{\partial \Theta_i} \\[6pt] \text{velocity} \quad & V_i \leftarrow \beta_2 V_{i-1} + (1 - \beta_2) \left(\frac{\partial L}{\partial \Theta_i}\right)^2 \\[6pt] & \Theta_i \leftarrow \Theta_{i-1} - \alpha \left(\frac{M_i}{\sqrt{V_i} + \varepsilon}\right) \end{aligned}\] \[\begin{aligned} \text{where} \\ M_i &\equiv \text{momentum (1st moment)} \\ V_i &\equiv \text{velocity: gradient squared (2nd moment)} \\ \beta_1, \beta_2 &\equiv \text{decay hyperparameters} \\ \varepsilon &\equiv \text{small constant for numerical stability} \end{aligned}\]

Cons of Adam

Memory intensive.
Challenging hyper-parameter tuning.
Independent update of each value in a single, long, parameter vector without considering any internal structure (vector-based optimizer behavior) of the model parameters.³

Question: Can we explicitly account for the underlying matrix structure of the model parameters?

The Linear Layer & Matrix Momentum

For a linear layer with 3 inputs $(x_1, x_2, x_3)$ and 4 outputs $(z_1, z_2, z_3, z_4)$:

\[z_i = \Theta_{i1} x_1 + \Theta_{i2} x_2 + \Theta_{i3} x_3, \quad \forall\, i = 1, 2, 3, 4\]
In matrix form:
\[\begin{bmatrix} z_1 \\ z_2 \\ z_3 \\ z_4 \end{bmatrix} = \begin{bmatrix} \Theta_{11} & \Theta_{12} & \Theta_{13} \\ \Theta_{21} & \Theta_{22} & \Theta_{23} \\ \Theta_{31} & \Theta_{32} & \Theta_{33} \\ \Theta_{41} & \Theta_{42} & \Theta_{43} \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \\ x_3 \end{bmatrix}\]
Since $\Theta$ is now a matrix, we need a matrix momentum $\hat{M}$ for $\hat{\Theta}$:
\[\hat{M} = \begin{bmatrix} m_{11} & m_{12} & m_{13} \\ m_{21} & m_{22} & m_{23} \\ m_{31} & m_{32} & m_{33} \\ m_{41} & m_{42} & m_{43} \end{bmatrix}\]
So the matrix momentum update is:
\[\hat{M}_i \leftarrow \beta \hat{M}_i + \frac{\partial L}{\partial \hat{\Theta}_i}\]

The Problem with Vector-Based Optimizers

With vector-based optimizers like Adam, the momentum for a linear layer (a 2D matrix) tends to become almost low rank in practice.
Essentially, only a small number of dominant directions really drive the update, while the many remaining other directions contribute very little.

Question: How can we tackle this update direction imbalance and what makes a good optimizer?

From fundamental first principles, a good optimizer possesses two characteristics: stability and speed. The goal of each update of a good optimizer is to minimize model variance and maximize loss reduction contribution, which correspond to stability and speed respectively.¹

Matrix Orthogonalization

We can fix the imbalance of update directions via orthogonalization.
Orthogonalize the momentum matrix. This is where Muon comes in.
It amplifies the effect of rare directions – the directions that typically receive small or infrequent updates.
Even though these rare directions seem minor, they are often essential for effective learning and can help capture more nuance patterns in the data.

Orthogonalization via SVD

We want the orthogonal matrix $O$ closest to $M$:
\[\text{Ortho}(M) = \arg\min_{O} \left\{\| O - M \|_F\right\} \quad \text{subject to } OO^T = I \text{ or } O^TO = I\]
Using SVD (Singular Value Decomposition): $M = USV^T$, where:
\[\begin{aligned} UU^T &= U^TU = I \\ VV^T &= V^TV = I \end{aligned} \quad \bigg\} \text{ orthonormal matrices}\]
Since $OO^T = O^TO = I$:
\[\therefore\quad O = USV^T \quad \text{where } S = \begin{bmatrix} 1 & & 0 \\ & \ddots & \\ 0 & & 1 \end{bmatrix} \equiv \text{unit diagonal matrix}\]
Issue: SVD on a matrix is computationally expensive.
Fix: Use an odd polynomial matrix:
\[p(X) = aX + b(XX^T)X\]

Odd Polynomial Matrix

\[\begin{align*} p(M) &= a(M) + b(MM^T)M \\ &= \left(aI + b(MM^T)\right)M \\ &= \left(aI + b(USV^T VS U^T)\right)USV^T \\ &= \left(aI + b(US^2 U^T)\right)USV^T \\ &= aUSV^T + bUS^2 U^T USV^T \\ p(M) &= aUSV^T + bUS^3V^T \\ \end{align*}\]

This applies to any odd polynomial, so in general:
\[p(M) = U\!\left(aS + bS^3 + cS^5 + \ldots + (\text{const})\, S^{2n+1}\right)V^T, \quad \forall\, n \geq 0,\ n \to \infty\]
Sticking to the 5th order polynomial:
\[\boxed{p(M) = U\!\left(aS + bS^3 + cS^5\right)V^T}\]
We need to determine the coefficients $(a, b, c)$. The goal is to get $S$ to a unit diagonal matrix – i.e., get the diagonal values as close to 1 as possible.
Example: Plot $y = p(x)$ against $x$. If $(a, b, c) = (1.5,\ -0.5,\ 0)$:
\[y = 1.5x - 0.5x^3\]
Applying this repeatedly via Newton-Schulz iteration:

\[\begin{aligned} y_1 &= p(x) \\ y_2 &= p(p(x)) \\ y_3 &= p(p(p(x))) \\ y_4 &= p(p(p(p(x)))) \\ y_5 &= p(p(p(p(p(x))))) \\ \end{aligned}\]

Each $y_k$ represents one more composition: $y_1 \to y_5$ are multiple iterations aimed at converging the singular values toward 1.

Newton-Schulz-5 Iteration

After 5 iterations, almost all input values end up very close to 1.
We can change $(a, b, c)$ to see the effect on convergence of $y$ to 1.
$(a, b, c) = (2,\ -1.5,\ 0.5)$ speeds up the convergence to 1.
Empirically, we don’t need the singular values to converge to exactly 1.
Let’s set an upper & lower bound, e.g. $(0.7,\ 1.3)$, which is basically $[1 - \varepsilon, 1 + \varepsilon]$.
Tuned coefficients:
\[\boxed{(a,\ b,\ c) = (3.4445,\ -4.775,\ 2.0315)}\]
Each Newton-Schulz 5 iteration involves only matrix multiplication:
\[X \leftarrow aX + b(XX^T)X + c(XX^T)^2 X\]
With GPUs, no need to use SVD since GPUs can efficiently compute matrix multiplication.

def newtonschulz5(G, steps=5, eps=1e-7):
    assert G.ndim == 2
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.bfloat16()
    X /= (X.norm() + eps)
    if G.size(0) > G.size(1):
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    if G.size(0) > G.size(1):
        X = X.T
    return X

# Retrieved from https://kellerjordan.github.io/posts/muon/

Muon

Muon is designed specifically for 2D weight matrices in neural network hidden layers.⁴ Unlike traditional optimizers that treat each parameter independently, Muon leverages the geometric structure of weight matrices by orthogonalizing gradients using the Newton-Schulz iteration. Muon is specifically designed for linear neural network layers, which aligns with ongoing research that argues that different layer types require different optimizers due to their varying geometry.⁵

The optimizer formulates weight updates as a constrained optimization problem in the RMS-to-RMS operator norm space:

\[\text{Ortho}(M) = \arg\min_{O} \left\{\| O - M \|_F\right\} \quad \text{subject to } OO^T = I \text{ or } O^TO = I\]

Where $M$ is the gradient matrix. The solution involves projecting the gradient onto the set of orthogonal matrices, which aims to standardize all singular values to 1 while preserving gradient directions.

Pseudo-Algorithm for Muon

for t = 1, 2, ..., do:
    Compute gradient      G_t  ←  ∇L_t(θ_{t-1})
    Compute momentum      M_t  ←  βM_{t-1} + G_t
    Normalize             M'_t ←  M_t / ||M_t||_F
    Orthogonalization     O_t  ←  NewtonSchutz5(M'_t)
    Update parameter      θ_t  ←  θ_{t-1} - α O_t

With weight decay $(+)$:
\[\Theta_t \leftarrow \Theta_{t-1} - \alpha\!\left(O_t + \lambda\Theta_{t-1}\right)\]
With adjusted learning rate $(+)$:
\[\Theta_t \leftarrow \Theta_{t-1} - \alpha\!\left[\left(0.2\sqrt{\max(n,m)}\right) O_t + \lambda\Theta_{t-1}\right]\] \[\begin{aligned} \text{where} \\ \lambda &\equiv \text{weight decay coefficient} \\ n, m &\equiv \text{dimensions of the 2D parameter matrix} \end{aligned}\]

Muon + Weight Decay + RMS Alignment

Weight decay is used to address the diminished performance gains of Muon over AdamW when scaling up to train a larger model.
The learning rate also gets adjusted by taking into account the size of the 2D matrix. This is the underlying principle behind the RMS (Root Mean Squared) Alignment.
- The scaling factor used to scale the Muon update for each matrix to ensure per-matrix parameter update RMS alignment of around 1 of matrices of different shapes is $\sqrt(\max(A, B))$ for a full-rank weight matrix of shape $[A,\ B]$.⁶
- The $0.2$ factor is used to match Muon’s update RMS to that of AdamW. From empirical observations, AdamW’s update RMS is usually around $0.2$ to $0.4$.⁶, ¹
These 2 improvements (weight decay & adjusting the per-parameter update scale) help to stabilize the training of large models.

The Exploding Attention Logit Crisis

Issue: Attention logits can grow larger & larger as training continues, which may cause the training process to become unstable.
Consider a sequence of 4 tokens & assume self-attention for simplicity. Each token is mapped to an embedding vector of dimension $d$. Let the embedding matrix be $X$.

Self-Attention

$O = \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

In self attention above, $Q = XW^Q$, $K = XW^K$, and $V = XW^V$.
The attention logits $S$:
\[\begin{align*} S &= QK^T \\ &= (XW^Q)(XW^K)^T \\ &= X\!\left(W^Q W^{K^T}\right)X^T \end{align*}\]
where $X$ and $X^T$ denote the embedding vectors, which are typically normalized to have unit norms.
To prevent the attention logits from becoming excessively large, we must control the scale of $W^Q$ and $W^K$:
\[S = X\underbrace{\left(W^Q W^{K^T}\right)}_{\text{scale control}}X^T\]

QK-Clip

A common strategy is to apply a scaling vector to these matrices.
During training, monitor the maximum value of the attention logits, $S_{\max}$. If it exceeds a certain threshold $\tau$, calculate a scaling ratio $\gamma$:

\[\text{if } S_{\max} > \tau: \quad \gamma = \frac{\tau}{S_{\max}} \implies \frac{\tau}{S_{\max}} < 1\]

Directly constrains attention logits, ensuring they stay within a safe range by rescaling the query & key projection weights.

Idea: Scale the relevant model parameters by $\gamma$ when the attention logits surpass the threshold. Scale both $W^Q$, $W^K$ by $\sqrt{\gamma}$:

\[\begin{align*} S &= X\!\left(\gamma W^Q W^{K^T}\right)X^T \\ &= X\!\left(\sqrt{\gamma}\, W^Q\ \sqrt{\gamma}\, W^{K^T}\right)X^T \end{align*}\]

Revised pseudo-algorithm (QK-Clip):

θ_t ← MuonOptimizer(θ_{t-1}, G_t)
if S_max > τ:
    W^Q ← √γ W^Q
    W^K ← √γ W^K

QK-Clip for Multi-Head Attention

In practice, self-attention consists of multiple heads ($n_{\text{heads}} = h$).
When the maximum attention logits exceed the threshold, instead of rescaling all heads in the same way (which doesn’t make sense), we introduce an individual scaling factor for each head to control their logits separately.

Multi-Head Attention

$O = \text{MultiheadAttention}(Q, K, V) = \text{Concat}(head_0, head_1, .., head_h) W_o$

Attention logits per head:
\[S = X\!\left(\sqrt{\gamma_h}\, W_h^Q\ \sqrt{\gamma_h}\, W_h^{K^T}\right)X^T\] \[\text{if } S_{\max}^h > \tau \quad \text{then} \quad \gamma_h = \frac{\tau}{S_{\max}}\]

Algorithm for Muon + QK-Clip for MHA:

θ_t ← MuonOptimizer(θ_{t-1}, G_t)
if S^h_max > τ:
    W^Q_h ← √γ_h W^Q_h
    W^K_h ← √γ_h W^K_h

Multihead Latent Attention (MLA)

For multihead Latent Attention (MLA) by DeepSeek, things get tricky.
MLA compresses $Q, K, V$ representations into a low-rank space to reduce the size of the KV cache using a down-projection matrix which produces latent representations:
\[C^Q = XW^Q_\downarrow \qquad C^{KV} = XW^{KV}_\downarrow\]
These compressed latent vectors are then mapped back to $W^Q_\uparrow$, $W^K_\uparrow$, $W^V_\uparrow$ for each attention head using the corresponding up-projection matrices.
Issue: This low-rank KV compression fails with rotary position embedding (RoPE).
Fix: A decoupled RoPE technique which introduces extra multi-head queries $W^{QR}$ and a shared key $W^{KR}$ to encode positional information.
For MLA, the Query, Key & Values are regrouped for each head:
- The Query is constructed by concatenating the compressed query $Q^C$ with the rotated query $Q^R$.
- The Key is constructed similarly by concatenating the compressed key $K^C$ with the rotated key $K^R$.

Multihead Latent Attention (MLA)

MuonClip

In MLA, it is important to carefully decide how to rescale these 4 matrices: $W^Q_\uparrow,\ W^{QR},\ W^{KR},\ W^K_\uparrow$.

For the up-projection matrices $W^Q_\uparrow$, $W^K_\uparrow$: rescale these parameters for each head individually.
The RoPE components $W^{QR}$, $W^{KR}$ rescaling is more tricky:
- Each head has its own rotary query $W^{QR}$.
- But all heads share a single rotary key matrix $W^{KR}$.

Issue: Applying the same per-head scaling for both RoPE components leads to the shared $W^{KR}$ being rescaled multiple times, which is undesirable.

Fix: Rescale only the head-specific rotary query $W^{QR}$ by their respective $\gamma^h$, while leaving the shared rotary key matrix $W^{KR}$ unchanged. This technique is called MuonClip. MuonClip improves upon Muon with the QK-Clip technique to handle training instability while benefiting from Muon’s advanced token efficiency.⁷

MuonClip Algorithm:

θ_t ← MuonOptimizer(θ_{t-1}, G_t)
if S^h_max > τ:
    γ_h = τ / S^h_max
    W^Q_{h,↑}  ← √γ_h  W^Q_{h,↑}
    W^K_{h,↑}  ← √γ_h  W^K_{h,↑}
    W^{QR}_h   ← γ_h   W^{QR}_h
    W^{KR}     ← W^{KR}   (unchanged)

In math:

\[\begin{aligned} \gamma_h &= \frac{\tau}{S_{\max}^h} \\[6pt] W_{h,\uparrow}^Q &\leftarrow \sqrt{\gamma_h}\, W_{h,\uparrow}^Q \\ W_{h,\uparrow}^K &\leftarrow \sqrt{\gamma_h}\, W_{h,\uparrow}^K \\ W_h^{QR} &\leftarrow \gamma_h\, W_h^{QR} \\ W^{KR} &\leftarrow W^{KR} \quad \text{(unchanged)} \end{aligned}\]

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026muonmuonclip,
  title   = "Muon & MuonClip Optimizers",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Apr",
  url     = "https://chizkidd.github.io/2026/04/04/muon-muonclip/"
}

References

Kimi (Moonshot AI). Why We Chose Muon: Our Chain of Thought. X (Twitter). 2025. ↩ ↩² ↩³
Chizoba Obasi. A Complete Guide to Neural Network Optimizers. chizkidd.github.io. 2026. ↩
Jia-Bin Huang This Simple Optimizer Is Revolutionizing How We Train AI [Muon] (YouTube Video). YouTube. ↩
Keller Jordan. Muon: An optimizer for hidden layers in neural networks. kellerjordan.github.io. 2024. ↩
Jeremy Bernstein. Deriving Muon. jeremybernste.in. 2025. ↩
Muon is Scalable for LLM Training. arXiv. 2025. ↩ ↩²
Kimi K2: Open Agentic Intelligence. arXiv. 2025. ↩

Inkcast: Turn Any EPUB or PDF into an Audiobook in Your Browser

Mon, 16 Mar 2026 00:00:00 +0000

Earlier this year, I decided to force myself to read more. Not a New Year’s resolution, because those never last. The reason is that growing up as a child and young teenager, reading often felt like punishment. My mum required my siblings and me to read a certain number of pages from a designated book every day throughout elementary school. Missing a day meant mandatory punishment. In boarding secondary school, this eventually led to a stubborn, subconscious resistance to non-essential reading. Over the six years I spent there, I probably read only five to ten non-academic fiction books (though Artemis Fowl was a delight). So it is not hard to see where my indifference to reading came from.

During the COVID-19 pandemic, however, I fell deep into podcasts. As an avid sports fan and TV show buff, I listened to everything: sports recaps, tech podcasts, expert interviews, the works (meeting Walter White at Anfield would be the ultimate dream). So when I decided to read more this year, audiobooks felt like the natural bridge. I already had a few EPUBs in the Apple Books app on my reading list and wondered: Can I listen to these EPUBs using Apple Dictation’s two-finger swipe-down feature? Unfortunately, it only works for the current page. It is quite janky, not very user-friendly, and frankly does not work well for my use case.

Recently, I worked on evaluating how a Facebook state-of-the-art (SOTA) automatic speech recognition (ASR) model handles Igbo tones, trying to see whether it actually “listens” properly. So I have been dabbling with audio quite a bit this year. You could say I have been thinking about listening a lot. In the past, I also experimented with WaveNet (a generative model for raw audio) and its fundamental building block, the dilated causal convolution.

With these experiences in mind, I wondered: Can I build an iPhone Shortcut that lets me listen to EPUBs properly? That question eventually led to Inkcast. The goal was not to build a Speechify competitor. I simply wanted to solve a personal problem. My aim was to create a low-effort, frictionless tool for personal use, so I made a GitHub repository and started building. Within a few days, I had a working website that could take EPUBs and PDFs and let users listen to the content organized by chapters in a sidebar with one-tap navigation. It included basic controls such as play/pause, rewind (15 seconds), forward (30 seconds), playback speed control (0.75× to 2×), and voice selection.

It worked well on desktop, so I used the URL to create a Shortcut on my iPhone. On mobile it also worked, but there was one problem: the reader voices sounded robotic and monotone, which is not ideal for long-form listening. The irony was that only weeks earlier I had been evaluating how machines handle speech. So I went back to the drawing board to figure out how to get more natural-sounding reader voices on Inkcast.

While researching audiobook-quality text-to-speech, I came across several APIs (OpenAI, ElevenLabs, Google Cloud). None were free, and for a personal project I wanted something that required no subscriptions or API keys. Most resources suggested that human-quality narration requires a dedicated TTS service. Eventually I discovered that the Web Speech API can access premium voices already installed on the device. These voices are free, require no API keys, and remain available offline. They are not state of the art, but they are surprisingly good. Many people do not realize that higher-quality Siri voices can be downloaded. Since then I have also added optional HD voices that run on-device via Kokoro-82M (a one-time download of about 90 to 330 MB, then they work offline) for the highest quality. The voice quality improved, but the project also started evolving in another direction.

I have always wanted to work through Paul Graham’s essays properly. There are 229 of them, and they read almost like long-form podcasts. But they live on webpages, which raised another question: Why limit the input to EPUBs and PDFs? So I added URL support. I pasted Paul Graham’s archive page into Inkcast, and it automatically pulled all 200+ essays into the sidebar. That was the moment I realized the idea actually worked.

The app itself lives in a single HTML file. There are no accounts and no installations, and your EPUBs and PDFs never leave your device. The one exception is web articles: when you paste a URL, a small proxy fetches the page for you, since browsers do not let a page fetch other sites directly. Everything else runs locally, so it ended up functioning as a privacy-preserving tool by default. In a way, I started the year studying whether machines listen well. Along the way, I realized that humans do not have many good free tools for listening either. Speechify costs $139 per year and Audible requires a subscription, so I built something that worked for me.

Feel free to try it, and if you find it useful or it saves you a Speechify subscription, you can buy me a coffee.

Eligibility Traces Explained: TD(λ), Sarsa(λ), and the λ-Return (S&B Ch. 12)

Fri, 13 Mar 2026 00:00:00 +0000

Eligibility traces are one of the basic mechanisms of RL that unify and generalize TD and Monte Carlo (MC) methods.
TD methods augmented with eligibility traces produce a family of methods spanning a range from MC methods at one end ($\lambda = 1$) to one-step TD (TD(0)) methods at the other end ($\lambda = 0$).
With eligibility traces, MC methods can be implemented online and on continuing problems.
$n$-step methods also unify TD and MC methods but are not as elegant algorithmically as eligibility traces (ET).
The eligibility traces algorithm entails:
- First, we have a short-term memory vector, the eligibility trace $\mathbf{z}_t \in \mathbb{R}^d$, that parallels the long-term weight vector $\mathbf{w}_t \in \mathbb{R}^d$.
- Then when a component of $\mathbf{w}_t$ participates in producing an estimated value, the corresponding component of $\mathbf{z}_t$ is bumped up and then begins to fade away.
- Learning occurs in that component of $\mathbf{w}_t$ if a non-zero TD error occurs before the trace falls back to zero (fades away).
- The trace-decay parameter $\lambda \in [0,1)$ determines the rate at which the trace falls.
Advantages of ET over $n$-step methods:
- Requires only a single trace vector $\mathbf{z}_t$ rather than storing the last $n$ feature vectors.
- Learning occurs continually and uniformly in time rather than being delayed and playing “catch up” at episode end. This leads to immediate behavior effect from learning rather than being delayed.
- ET is a backward view algorithm, unlike $n$-step methods that are forward view algorithms, which is less complex to implement.
Forward view algorithms are based on looking forward from the updated state, and the updated state depends on all the future rewards.
Backward view algorithms use the current TD error, looking backward to recently visited states, to achieve nearly the same updates as forward view.
We start with ideas for state values and prediction, then extend them to action values and control, then do on-policy, then extend to off-policy learning. The field of focus is linear function approximation (covers tabular and state aggregation cases).

12.1 The $\lambda$-return
12.2 TD($\lambda$)
12.3 $n$-step Truncated $\lambda$-return Methods
12.4 Redoing Updates: Online $\lambda$-return Algorithm
12.5 True Online TD($\lambda$)
12.6 Dutch Traces in Monte Carlo Learning
12.7 Sarsa($\lambda$)
12.8 Variable $\lambda$ and $\gamma$
12.9 Off-Policy Traces with Control Variates
12.10 Watkins’s Q($\lambda$) to Tree-Backup($\lambda$)
12.11 Stable Off-Policy Methods with Traces
12.12 Implementation Issues
12.13 Conclusions

Appendix

Citation

12.1 The $\lambda$-return

Recall in Chapter 7 we defined an $n$-step return as the sum of the first $n$ rewards plus the estimated value of the state reached in $n$ steps, each appropriately discounted:

\[G_{t:t+n} \doteq R_{t+1} + \gamma R_{t+2} + \ldots + \gamma^{n-1} R_{t+n} + \gamma^n \hat{v}(S_{t+n}, \mathbf{w}_{t+n-1}), \quad\quad 0 \leq t \leq T - n\]

A valid update can be done not just towards any $n$-step return, but also towards any average of $n$-step returns.
- E.g. average the 2-step and 4-step return: $\frac{1}{2} G_{t:t+2} + \frac{1}{2} G_{t:t+4}$

Backup Diagram for Compound Update

The compound update mixing half of a two-step return and half of a four-step return.

Any set of $n$-step returns can be averaged, even an infinite set, as long as the weights on the component returns are positive and sum to $1$.
What if instead of using one $n$-step return, we use a weighted average of all $n$-step returns? This leads to averaging which produces a substantial new range of algorithms.E.g.,
1. Averaging one-step and infinite-step returns to interrelate TD and MC methods.
2. Averaging experience-based updates with Dynamic Programming (DP) updates to obtain a single combination of experience-based and model-based methods.
An update that averages simpler component updates is called a compound update or the $\lambda$-return.
The TD($\lambda$) algorithm is one way of averaging $n$-step updates, each weighted proportionally by $\lambda^{n-1}$ (where $\lambda \in [0,1]$) and normalized by a factor of $(1-\lambda)$ to ensure that the weights sum to $1$.

Backup Diagram for TD($\lambda$)

If $\lambda = 0$, then the overall update reduces to its first component, the TD(0) update, whereas if $\lambda = 1$, then the overall update reduces to its last component, the MC update.

Essentially, $\lambda$-return, $G_t^\lambda$, combines all $n$-step returns $G_{t:t+n}$ in a weighted average manner, $(1-\lambda)\lambda^{n-1}$, and is defined in its state-based form by:

\[\boxed{G_t^\lambda \doteq (1-\lambda) \sum_{n=1}^{\infty} \lambda^{n-1} G_{t:t+n}}\]

TD($\lambda$) Weighting

The TD($\lambda$) weighting function diagram illustrates the weighting on the sequence of $n$-step returns in the $\lambda$-return:
- $1$-step return gets the largest weight, $1-\lambda$
- $2$-step return gets the next (2nd) largest weight, $(1-\lambda)\lambda$
- $3$-step return gets the 3rd largest weight, $(1-\lambda)\lambda^2$
- $n$-step return gets the $n$-th largest (smallest) weight, $(1-\lambda)\lambda^{n-1}$
- The weight fades by $\lambda$ with each additional step.
- After a terminal state has been reached, all subsequent $n$-step returns are equal to the conventional return $G_t$.
- So essentially, we can decompose $G_t^\lambda$ based on the TD($\lambda$) weighting function diagram into the main sum and post-termination terms:
\[\begin{array}{l} G_t^\lambda = (1-\lambda) \sum\nolimits_{n=1}^{T-t-1} \lambda^{n-1} G_{t:t+n} + \lambda^{T-t-1} G_t \\ \hspace{3em} \underbrace{\hspace{11em}}_{\text{pre-termination}} \kern{0.5em}\underbrace{\hspace{4em}}_{\text{post-termination}} \end{array}\]
- So now we can see the impact of $\lambda$ more clearly:
\[\begin{aligned} \text{if } \lambda = 1: \quad & G_t^\lambda = G_t \hspace{18em} \text{(MC)} \\[6pt] \text{if } \lambda = 0: \quad & G_t^\lambda = \left\{ \begin{array}{ll} \sum_{n=1}^{T-t-1} G_{t:t+n} & \text{for } n=1 \\ 0 & \text{for } n > 1 \end{array} \right\} = G_{t:t+1} \quad \text{(TD(0))} \end{aligned}\]

TD($\lambda$) Weighting

Weighting given in the $\lambda$-return to each of the $n$-step returns.

Our first learning algorithm based on the $\lambda$-return is the off-line $\lambda$-return algorithm, which waits until the end of an episode to make updates. Its semi-gradient, $\lambda$-return, target update for $t = 0, 1, 2, \ldots, T-1$ is:

\[\boxed{\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha \!\left[G_t^\lambda - \hat{v}(S_t, \mathbf{w}_t)\right] \nabla \hat{v}(S_t, \mathbf{w}_t)}\]

The $\lambda$-return allows us to smoothly move between MC and TD(0) methods, comparable to $n$-step returning.

Forward View

This approach is called the theoretical, forward view of a learning algorithm:
- Update value function towards the $\lambda$-return.
- Look forward in time to all the future rewards to compute $G_t^\lambda$.
- Like MC, can only be computed from complete return.

Forward View

We decide how to update each state by looking forward to future rewards and states.

12.2 TD($\lambda$)

TD($\lambda$) was the first algorithm that showed a formal relationship between a forward view and backward view using eligibility traces.
TD($\lambda$) improves over the off-line $\lambda$-return algorithm in 3 ways:
- Updates the weight vector on every step of an episode and not just the end.
- Equal distributions in time of its computations rather than at an episode’s end.
- Can be applied to continuing problems and not just episodic problems.
Let’s focus on the semi-gradient version of TD($\lambda$) with function approximation:
- The eligibility trace $\mathbf{z}_t$ has the same number of components as $\mathbf{w}_t$.
- $\mathbf{z}$ is initialized to $\mathbf{0}$, incremented on each time step by the value gradient, and then fades away by $\gamma\lambda$:
\[\begin{align*} \mathbf{z}_{-1} &\doteq \mathbf{0} \\ \mathbf{z}_t &\doteq \gamma\lambda \mathbf{z}_{t-1} + \nabla \hat{v}(S_t, \mathbf{w}_t), \quad 0 \leq t \leq T \end{align*}\] \[\text{where } \lambda \equiv \text{trace decay parameter and } \gamma \equiv \text{discount rate}\]
The eligibility trace keeps track of which $\mathbf{w}_t$ components have contributed, positively or negatively, to recent state valuations.
This is the recency heuristic used for credit assignment, where more credit is assigned to the most recent states. Recent is defined in terms of $\gamma\lambda$.
The TD error for state-value prediction is:
\[\delta_t \doteq R_{t+1} + \gamma \hat{v}(S_{t+1}, \mathbf{w}_t) - \hat{v}(S_t, \mathbf{w}_t)\]
and the weight vector update in TD($\lambda$) is proportional to the scalar TD error and the vector eligibility trace:

\[\boxed{\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha\, \delta_t\, \mathbf{z}_t}\]

Backward View

Forward view provides theory but backward view provides mechanism (practical) where we update online, every step, from incomplete sequences.
Keep an eligibility trace for every state $s$.
Update value $V(s)$ for every state $s$ in proportion to TD-error $\delta_t$ and eligibility trace $\mathbf{z}_t$:

\[\begin{aligned} \delta_t &= R_{t+1} + \gamma V(S_{t+1}) - V(S_t) \\ V(s) &\leftarrow V(s) + \alpha\, \delta_t\, \mathbf{z}_t \end{aligned}\]

Backward View of TD($\lambda$)

In the backward or mechanistic view of TD($\lambda$), each update depends on the current TD error combined with the current eligibility traces of past events.

Let’s look at the effect of $\lambda$ to understand the backward view of TD($\lambda$):

\[\begin{align*} \text{if } \lambda = 0: \quad & \mathbf{z}_t = \nabla \hat{v}(S_t, \mathbf{w}_t) \\ & \mathbf{w}_{t+1} = \mathbf{w}_t + \alpha\, \delta_t \nabla \hat{v}(S_t, \mathbf{w}_t) \quad \longrightarrow \quad \text{TD(0)} \\[6pt] \text{if } 0 < \lambda < 1: \quad & \text{earlier states are given less credit for the TD error} \\[6pt] \text{if } \lambda = 1: \quad & \mathbf{z}_t = \gamma \mathbf{z}_{t-1} + \nabla \hat{v}(S_t, \mathbf{w}_t) \quad \longrightarrow \quad \text{credit for earlier states falls by } \gamma \text{ per step} \\[6pt] \text{if } \lambda = 1,\ \gamma = 1: \quad & \mathbf{z}_t = \mathbf{z}_{t-1} + \nabla \hat{v}(S_t, \mathbf{w}_t) \quad \longrightarrow \quad \text{MC-like behavior (no time decay for ET)} \\[6pt] \text{if } \lambda = 1: \quad & \text{we get TD(1)} \end{align*}\]

In summary, $\lambda = 1$ yields TD(1).
TD(1) implements MC algorithms in a more general way and for a wider range of applicability:
- Not limited to episodic tasks, can be applied to discounted continuing tasks.
- Can be performed incrementally and online.
- Learns immediately and alters behavior during an episode if something good or bad happens, for control methods.
Linear TD($\lambda$) converges in the on-policy case if the step-size parameter $\alpha$ is reduced over time according to stochastic approximation theory conditions.
The convergence of linear TD($\lambda$) is not to the minimum-error weight vector but to a nearby weight vector that depends on $\lambda$.
The bound on solution quality generalized for any $\lambda$, for the continuing, discounted case is:

\[\overline{\text{VE}}(\mathbf{w}_\infty) \leq \frac{1 - \gamma\lambda}{1 - \gamma} \min_\mathbf{w} \overline{\text{VE}}(\mathbf{w})\]

That is, the asymptotic error is no more than $\dfrac{1-\gamma\lambda}{1-\gamma}$ times the smallest possible error for TD($\lambda$):

\[\begin{align*} \text{as } \lambda \to 1: \quad & \overline{\text{VE}}(\mathbf{w}_\infty) \to \min_\mathbf{w} \overline{\text{VE}}(\mathbf{w}) \\[6pt] \text{as } \lambda \to 0: \quad & \overline{\text{VE}}(\mathbf{w}_\infty) \to \frac{1}{1-\gamma} \min_\mathbf{w} \overline{\text{VE}}(\mathbf{w}) = \overline{\text{VE}}(\mathbf{w}_\text{TD}) \quad \text{(TD(0))} \end{align*}\]

However, $\lambda = 1$ is often the poorest choice.

12.3 $n$-step Truncated $\lambda$-return Methods

The off-line $\lambda$-return is of limited use because the $\lambda$-return is not known until the episode ends.
The off-line $\lambda$-return approximation is to truncate the sequence after a fixed number of steps.
Hence, a natural approximation is to truncate the sequence where $\lambda$-return cannot be calculated for an arbitrarily large $n$. This handles the continuing case.
The truncated $\lambda$-return for time $t$, given data only up to some later horizon $h$, is:

\[G_{t:h}^\lambda \doteq (1-\lambda) \sum_{n=1}^{h-t-1} \lambda^{n-1} G_{t:t+n} + \lambda^{h-t-1} G_{t:h}, \quad 0 \leq t \leq h \leq T\] \[\begin{aligned} \text{where } h &\equiv \text{horizon (plays same role as time of termination } T\text{)} \end{aligned}\]

Here the residual weighting is given to the longest available $n$-step return $G_{t:h}$.
The truncated $\lambda$-return gives rise to a family of $n$-step $\lambda$-return algorithms, known in the state-value case as Truncated TD($\lambda$) or TTD($\lambda$).
TTD($\lambda$) is defined for $0 \leq t < T$ by:

\[\boxed{\mathbf{w}_{t+n} \doteq \mathbf{w}_{t+n-1} + \alpha \!\left[G_{t:t+n}^\lambda - \hat{v}(S_t, \mathbf{w}_{t+n-1})\right] \nabla \hat{v}(S_t, \mathbf{w}_{t+n-1})}\]

Efficient implementation of TTD($\lambda$) relies on the $k$-step $\lambda$-return:

\[\boxed{G_{t:t+k}^\lambda = \hat{v}(S_t, \mathbf{w}_{t-1}) + \sum_{i=t}^{t+k-1} (\delta\lambda)^{i-t} \delta_i'}\] \[\begin{aligned} \text{where } \delta_i' &\equiv R_{t+1} + \gamma \hat{v}(S_{t+1}, \mathbf{w}_{t-1}) - \hat{v}(S_t, \mathbf{w}_{t-1}) \end{aligned}\]

Backup Diagram for Truncated TD($\lambda$)

The truncated $\lambda$-return gives rise to a family of $n$-step $\lambda$-return algorithms called TTD($\lambda$).

12.4 Redoing Updates: Online $\lambda$-return Algorithm

How do we choose the truncation parameter $n$ in TTD($\lambda$)?
It involves a tradeoff:
- $n$ should be large so that TTD($\lambda$) closely approximates off-line $\lambda$-return, but
- $n$ should also be small so that the updates can be made sooner and can influence behavior sooner.
In principle, we can achieve both cases via the online $\lambda$-return algorithm, but at the cost of computational complexity.
Essentially at each time step, we go back and redo all the updates since the beginning of the episode as we gather new increment of data:
- The new updates are better than the old ones because now they account for the time step’s new data.
- Basically this conceptual algorithm involves multiple passes over the episode, one at each horizon, each generating a different sequence of weight vectors.
Let’s distinguish between the weight vectors computed at the different horizons by seeing the first 3 sequences:

\[\begin{align*} h=1: \quad & \mathbf{w}_1^1 \doteq \mathbf{w}_0^1 + \alpha \!\left[G_{0:1}^\lambda - \hat{v}(S_0, \mathbf{w}_0^1)\right] \nabla \hat{v}(S_0, \mathbf{w}_0^1) \\[6pt] h=2: \quad & \mathbf{w}_0^2 \doteq \mathbf{w}_0^2 + \alpha \!\left[G_{0:2}^\lambda - \hat{v}(S_0, \mathbf{w}_0^2)\right] \nabla \hat{v}(S_0, \mathbf{w}_0^2) \\ & \mathbf{w}_2^2 \doteq \mathbf{w}_1^2 + \alpha \!\left[G_{1:2}^\lambda - \hat{v}(S_1, \mathbf{w}_1^2)\right] \nabla \hat{v}(S_1, \mathbf{w}_1^2) \\[6pt] h=3: \quad & \mathbf{w}_1^3 \doteq \mathbf{w}_0^3 + \alpha \!\left[G_{0:3}^\lambda - \hat{v}(S_0, \mathbf{w}_0^3)\right] \nabla \hat{v}(S_0, \mathbf{w}_0^3) \\ & \mathbf{w}_2^3 \doteq \mathbf{w}_1^3 + \alpha \!\left[G_{1:3}^\lambda - \hat{v}(S_1, \mathbf{w}_1^3)\right] \nabla \hat{v}(S_1, \mathbf{w}_1^3) \\ & \mathbf{w}_3^3 \doteq \mathbf{w}_2^3 + \alpha \!\left[G_{2:3}^\lambda - \hat{v}(S_2, \mathbf{w}_2^3)\right] \nabla \hat{v}(S_2, \mathbf{w}_2^3) \end{align*}\] \[\begin{aligned} \text{where} \\ \mathbf{w}_t^h &\equiv \text{weights used to generate the value at time } t \text{ in the sequence up to horizon } h \\ \mathbf{w}_0^h &\equiv \text{1st weight vector in each sequence that is inherited from the previous episode} \\ \mathbf{w}_n^h &\equiv \text{last weight vector in each sequence or the ultimate weight-vector sequence} \end{aligned}\]

The general form of the online $\lambda$-return update for $0 \leq t < h \leq T$ is:

\[\boxed{\mathbf{w}_{t+1}^h \doteq \mathbf{w}_t^h + \alpha \!\left[G_{t:h}^\lambda - \hat{v}(S_t, \mathbf{w}_t^h)\right] \nabla \hat{v}(S_t, \mathbf{w}_t^h)}\] \[\mathbf{w}_t \doteq \mathbf{w}_t^t\]

12.5 True Online TD($\lambda$)

The original ideal online $\lambda$-return algorithm shown in Section 12.4 is very complex so we use online TD($\lambda$) to approximate it.
We’ll use eligibility trace to invert the forward view, online $\lambda$-return to an efficient backward view algorithm. This is called the True Online TD($\lambda$).
It uses a simple strategy trick with the weight matrix, where we only need to keep the last weight vector from all the updates at the last time step (the diagonals of online $\lambda$-return weight matrix).

\[\begin{aligned} \begin{bmatrix} \mathbf{w}_0^0 & & & & \\ \mathbf{w}_0^1 & \mathbf{w}_1^1 & & & \\ \mathbf{w}_0^2 & \mathbf{w}_1^2 & \mathbf{w}_2^2 & & \\ \mathbf{w}_0^3 & \mathbf{w}_1^3 & \mathbf{w}_2^3 & \mathbf{w}_3^3 & \\ \vdots & \vdots & \vdots & \vdots & \ddots \\ \mathbf{w}_0^T & \mathbf{w}_1^T & \mathbf{w}_2^T & \mathbf{w}_3^T & \cdots & \mathbf{w}_T^T \end{bmatrix} &\longrightarrow \begin{bmatrix} \mathbf{w}_0^0 \\ & \mathbf{w}_1^1 \\ & & \mathbf{w}_2^2 \\ & & & \mathbf{w}_3^3 \\ & & & & \ddots \\ & & & & & \mathbf{w}_T^T \end{bmatrix} \\ \end{aligned}\] \[\text{Online } \lambda\text{-return} \hspace{8em} \text{True Online TD}(\lambda)\]

For the linear case in which $\hat{v}(s, \mathbf{w}) = \mathbf{w}^T \mathbf{x}(s)$, the true online TD($\lambda$) algorithm is:

\[\boxed{\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha\, \delta_t\, \mathbf{z}_t + \alpha \!\left(\mathbf{w}_t^T \mathbf{x}_t - \mathbf{w}_{t-1}^T \mathbf{x}_t\right)\!\left(\mathbf{z}_t - \mathbf{x}_t\right)}\] \[\begin{aligned} \text{where} \\ \mathbf{w}_t &\doteq \mathbf{w}_t^t \\ \mathbf{x}_t &\doteq \mathbf{x}(S_t) \\ \mathbf{z}_t &\doteq \gamma\lambda \mathbf{z}_{t-1} + (1 - \alpha\gamma\lambda\, \mathbf{z}_{t-1}^T \mathbf{x}_t)\, \mathbf{x}_t \end{aligned}\]

The per-step computational complexity of true online TD($\lambda$) is the same as TD($\lambda$), $O(d)$.
$\mathbf{z}_t$ used in true online TD($\lambda$) is called a dutch trace, unlike that of TD($\lambda$) which is called an accumulating trace.
Earlier work used a 3rd kind of trace called the replacing trace, defined only for the tabular case or for binary feature vectors (tile coding). It is defined:

\[\tilde{z}_{i,t} \doteq \left\{ \begin{array}{ll} 1 & \text{if } x_{i,t} = 1 \\ \gamma\lambda\, z_{i,t-1} & \text{otherwise} \end{array} \right\}\]

Nowadays, dutch traces usually perform better than replacing traces and have a clearer theoretical basis.
Accumulating traces remain of interest for nonlinear function approximations where dutch traces are unavailable.

12.6 Dutch Traces in Monte Carlo Learning

Eligibility traces have nothing to do with TD learning despite their close historical association.
Eligibility traces arise even in Monte Carlo learning.
Using dutch traces, we can invert the forward view MC algorithm to an equivalent, yet computationally cheaper backward view algorithm.
This is the only equivalence of forward and backward view that is explicitly demonstrated in this book.
The linear, gradient MC prediction algorithm makes the following sequence of updates, one for each time step of the episode:

\[\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha \!\left[G - \mathbf{w}_t^T \mathbf{x}_t\right] \mathbf{x}_t, \quad 0 \leq t < T\]

For simplicity, assume that the return $G$ is a single reward at the end of the episode (hence no subscript by time) and that there is no discounting.
This is known as the Least Mean Square (LMS) rule.
We introduce an additional vector memory, the eligibility trace, that keeps in it a summary of all the feature vectors seen so far. The overall update will be the same as the MC updates’ sequence shown above and is:
\[\begin{align*} \mathbf{w}_T &= \mathbf{w}_{T-1} + \alpha \!\left(G - \mathbf{w}_{T-1}^T \mathbf{x}_{T-1}\right) \mathbf{x}_{T-1} \\ &= \mathbf{w}_{T-1} + \alpha \mathbf{x}_{T-1}\!\left(-\mathbf{x}_{T-1}^T \mathbf{w}_{T-1}\right) + \alpha G \mathbf{x}_{T-1} \\ &= \!\left(\mathbf{I} - \alpha \mathbf{x}_{T-1} \mathbf{x}_{T-1}^T\right) \mathbf{w}_{T-1} + \alpha G \mathbf{x}_{T-1} \\ &= \mathbf{F}_{T-1}\, \mathbf{w}_{T-1} + \alpha G \mathbf{x}_{T-1} \end{align*}\] \[\begin{aligned} \text{where} \\ \mathbf{F}_t &\doteq \mathbf{I} - \alpha \mathbf{x}_t \mathbf{x}_t^T \equiv \text{a forgetting or fading matrix} \end{aligned}\] \[\therefore\quad \mathbf{w}_{T-1} = \mathbf{F}_{T-2}\, \mathbf{w}_{T-2} + \alpha G \mathbf{x}_{T-2}\]
Now recursing:
\[\begin{align*} \mathbf{w}_T &= \mathbf{F}_{T-1}\, \mathbf{w}_{T-1} + \alpha G \mathbf{x}_{T-1} \\ &= \mathbf{F}_{T-1}\!\left(\mathbf{F}_{T-2}\, \mathbf{w}_{T-2} + \alpha G \mathbf{x}_{T-2}\right) + \alpha G \mathbf{x}_{T-1} \\ &= \mathbf{F}_{T-1} \mathbf{F}_{T-2}\, \mathbf{w}_{T-2} + \alpha G\!\left(\mathbf{F}_{T-1} \mathbf{x}_{T-2} + \mathbf{x}_{T-1}\right) \\ &= \mathbf{F}_{T-1} \mathbf{F}_{T-2}\!\left(\mathbf{F}_{T-3}\, \mathbf{w}_{T-3} + \alpha G\, \mathbf{x}_{T-3}\right) + \alpha G\!\left(\mathbf{F}_{T-1}\, \mathbf{x}_{T-2} + \mathbf{x}_{T-1}\right) \\ &= \mathbf{F}_{T-1} \mathbf{F}_{T-2} \mathbf{F}_{T-3}\, \mathbf{w}_{T-3} + \alpha G\!\left(\mathbf{F}_{T-1} \mathbf{F}_{T-2}\, \mathbf{x}_{T-3} + \mathbf{F}_{T-1}\, \mathbf{x}_{T-2} + \mathbf{x}_{T-1}\right) \\ &\quad \vdots \\ &= \underbrace{\mathbf{F}_{T-1} \mathbf{F}_{T-2} \cdots \mathbf{F}_0\, \mathbf{w}_0}_{\mathbf{a}_{T-1}} + \alpha G \underbrace{\sum\nolimits_{k=0}^{T-1} \mathbf{F}_{T-1} \mathbf{F}_{T-2} \cdots \mathbf{F}_{k+1}\, \mathbf{x}_k}_{\mathbf{z}_{T-1}} \\ &= \mathbf{a}_{T-1} + \alpha G\, \mathbf{z}_{T-1} \end{align*}\] \[\begin{aligned} \text{where} \\ \mathbf{a}_{T-1}\ \&\ \mathbf{z}_{T-1} &\equiv \text{values at time } T-1 \text{ of 2 auxiliary memory vectors that can be updated} \\ &\phantom{{}\equiv{}} \text{incrementally w/o knowledge of } G \text{ and with } O(d) \text{ complexity per time step} \\ \mathbf{z}_t &\equiv \text{dutch-style eligibility trace, initialized to } \mathbf{z}_0 = \mathbf{x}_0 \end{aligned}\]
The $\mathbf{z}_t$ vector is in fact a dutch-style eligibility trace, initialized to $\mathbf{z}_0 = \mathbf{x}_0$, that can be updated according to:
\[\begin{align*} \mathbf{z}_t &= \sum\nolimits_{k=0}^{t} \mathbf{F}_t \mathbf{F}_{t-1} \cdots \mathbf{F}_{k+1}\, \mathbf{x}_k, \quad 1 \leq t < T \\ &= \sum\nolimits_{k=0}^{t-1} \mathbf{F}_t \mathbf{F}_{t-1} \cdots \mathbf{F}_{k+1}\, \mathbf{x}_k + \mathbf{x}_t \\ &= \mathbf{F}_t \sum\nolimits_{k=0}^{t-1} \mathbf{F}_{t-1} \mathbf{F}_{t-2} \cdots \mathbf{F}_{k+1}\, \mathbf{x}_k + \mathbf{x}_t \\ &= \mathbf{F}_t\, \mathbf{z}_{t-1} + \mathbf{x}_t \\ &= \!\left(\mathbf{I} - \alpha \mathbf{x}_t \mathbf{x}_t^T\right) \mathbf{z}_{t-1} + \mathbf{x}_t \\ &= \mathbf{z}_{t-1} - \alpha\!\left(\mathbf{z}_{t-1}^T \mathbf{x}_t\right) \mathbf{x}_t + \mathbf{x}_t \\ &\boxed{= \mathbf{z}_{t-1} + \!\left(1 - \alpha\, \mathbf{z}_{t-1}^T \mathbf{x}_t\right) \mathbf{x}_t} \end{align*}\]
which is the dutch trace for $\gamma\lambda = 1$.
The $\mathbf{a}_t$ auxiliary vector is initialized to $\mathbf{a}_0 = \mathbf{w}_0$ and then updated according to:
\[\begin{align*} \mathbf{a}_t &= \mathbf{F}_t \mathbf{F}_{t-1} \cdots \mathbf{F}_0\, \mathbf{w}_0, \quad 1 \leq t < T \\ &= \mathbf{F}_t\, \mathbf{a}_{t-1} \\ &= \mathbf{a}_{t-1} - \alpha \mathbf{x}_t \mathbf{x}_t^T \mathbf{a}_{t-1} \\ &\boxed{= \mathbf{a}_{t-1}\!\left(\mathbf{I} - \alpha \mathbf{x}_t \mathbf{x}_t^T\right)} \end{align*}\]

Takeaways

The auxiliary vectors, $\mathbf{a}_t$ and $\mathbf{z}_t$, are updated on each time step $t < T$ and then, at time $T$ when $G$ is observed, they are used to compute:

\[\boxed{\mathbf{w}_T = \mathbf{a}_{T-1} + \alpha G\, \mathbf{z}_{T-1}}\]

The time and memory complexity per step is $O(d)$.
This is surprising and intriguing since ET is working in a non-TD setting (ET arise where long-term predictions are needed to be learned efficiently).

12.7 Sarsa($\lambda$)

Now let’s extend eligibility traces to action-value methods.
First, let’s recall the action-value form of the $n$-step return:
\[G_{t:t+n} \doteq R_{t+1} + \ldots + \gamma^{n-1} R_{t+n} + \gamma^n \hat{q}(S_{t+n}, A_{t+n}, \mathbf{w}_{t+n-1}), \quad t+n < T\] \[\text{with} \quad G_{t:t+n} = G_t \quad \text{ if } t+n \geq T.\]
With this, for $t = 0, \ldots, T-1,$ let’s form the action-value form of the off-line $\lambda$-return algorithm which uses $\hat{q}$ rather than $\hat{v}$:

\[\boxed{\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha \!\left[G_t^\lambda - \hat{q}(S_t, A_t, \mathbf{w}_t)\right] \nabla \hat{q}(S_t, A_t, \mathbf{w}_t)}\] \[\begin{aligned} \text{where} \quad G_t^\lambda &\doteq G_{t:\infty}^\lambda \end{aligned}\]

For the forward view shown in the figure below, which is similar to TD($\lambda$), the updates are:
- 1st update: one full-step lookahead
- 2nd update: two-step lookahead
- Final update: complete return.

Backup Diagram for Sarsa($\lambda$)

The first update looks ahead one full step, to the next state–action pair, the second looks ahead two steps, to the second state–action pair, and so on. A final update is based on the complete return.

The weighting of each $n$-step update in the $\lambda$-return is same as in TD($\lambda$) and $\lambda$-return.
The forward view TD for action values, Sarsa($\lambda$), has the same update rule as TD($\lambda$):

\[\boxed{\mathbf{w}_{t+1} \doteq \mathbf{w}_t + \alpha\, \delta_t\, \mathbf{z}_t}\]

The action-value form of the TD error is used:

\[\delta_t \doteq R_{t+1} + \gamma \hat{q}(S_{t+1}, A_{t+1}, \mathbf{w}_t) - \hat{q}(S_t, A_t, \mathbf{w}_t)\]

The action-value form of the eligibility trace:

\[\begin{align*} \mathbf{z}_{-1} &\doteq \mathbf{0} \\ \mathbf{z}_t &\doteq \gamma\lambda \mathbf{z}_{t-1} + \nabla \hat{q}(S_t, A_t, \mathbf{w}_t), \quad 0 \leq t \leq T \end{align*}\]

There exists an action-value version of our ideal TD method, the online $\lambda$-return algorithm and its efficient implementation as true online TD($\lambda$):
- Section 12.4: Everything there holds here except for using the action-value form of the $n$-step return, $G_{t:t+n}$.
- Sections 12.5 & 12.6: Everything holds here except for using state-action feature vectors $\mathbf{x}_t = \mathbf{x}(S_t, A_t)$ instead of state feature vectors $\mathbf{x}_t = \mathbf{x}(S_t)$.
- The resulting efficient backward algorithm obtained from using the eligibility trace to invert the action-value form of the forward view, online $\lambda$-return is called the True Online Sarsa($\lambda$).
There is also a truncated version of Sarsa($\lambda$), called Forward Sarsa($\lambda$), which appears to be a model-free, control method for use in conjunction with multi-layer ANNs.

12.8 Variable $\lambda$ and $\gamma$

To get the most general forms of the final TD algorithms, it is vital to generalize the degree of bootstrapping and discounting beyond constant parameters to functions dependent on the state and action:

\[\begin{align*} \lambda_t &\doteq \lambda(S_t, A_t), \quad & \lambda &: S \times A \to [0,1] \\ \gamma_t &\doteq \gamma(S_t), \quad & \gamma &: S \to [0,1] \end{align*}\]

$\lambda_t$ is the termination function and is significant because it changes the return $G_t$, which is now more generally defined as:

\[\begin{align*} G_t &\doteq R_{t+1} + \gamma_{t+1} G_{t+1} \\ &= R_{t+1} + \gamma_{t+1} R_{t+2} + \gamma_{t+1} \gamma_{t+2} R_{t+3} + \gamma_{t+1} \gamma_{t+2} \gamma_{t+3} R_{t+4} + \ldots \\ &= \sum_{k=t}^{\infty} \left(\prod_{i=t+1}^{k} \gamma_i\right) R_{k+1} \end{align*}\] \[\begin{aligned} \text{where } \prod_{k=t}^{\infty} \gamma_k &= 0 \text{ with probability 1 for all } t, \text{ to assure the sums are finite} \end{aligned}\]

This general return $G_t$ definition enables episodic settings to become a single stream of experience, without special terminal state, start distributions or termination times
- A terminal state just becomes a state with $\gamma(s) = 0$ that transitions to the start distribution.
Generalization to variable bootstrapping yields a new state-based $\lambda$-return:

\[\boxed{G_t^{\lambda s} \doteq R_{t+1} + \gamma_{t+1}\!\left[(1 - \lambda_{t+1})\, \hat{v}(S_{t+1}, \mathbf{w}_t) + \lambda_{t+1}\, G_{t+1}^{\lambda s}\right]}\]

Action-based $\lambda$-return is either the Sarsa form:

\[\boxed{G_t^{\lambda a} \doteq R_{t+1} + \gamma_{t+1}\!\left[(1 - \lambda_{t+1})\, \hat{q}(S_{t+1}, A_{t+1}, \mathbf{w}_t) + \lambda_{t+1}\, G_{t+1}^{\lambda a}\right]}\]

or the Expected Sarsa form:

\[\boxed{G_t^{\lambda a} \doteq R_{t+1} + \gamma_{t+1}\!\left[(1 - \lambda_{t+1})\, \bar{V}_t(S_{t+1}) + \lambda_{t+1}\, G_{t+1}^{\bar{\lambda}a}\right]}\] \[\begin{aligned} \text{where } \bar{V}_t(s) \doteq \sum_a \Pi(a \vert s)\, \hat{q}(s, a, \mathbf{w}_t) \\ \end{aligned}\]

Superscripts notation for $i$ in $G_t^{\lambda i}$

\[\begin{aligned} \text{"s"} &: \text{bootstraps from state values} \\ \text{"a"} &: \text{bootstraps from action values} \end{aligned}\]

12.9 Off-Policy Traces with Control Variates

To generalize to off-policy, we need to incorporate importance sampling using eligibility traces.
Let’s focus on the bootstrapping generalization of per-decision importance sampling with control variates (Section 7.4).
The new state-based $\lambda$-return in Section 12.8 generalizes, after the off-policy, control variate, $n$-step return (ending at horizon $h$) model to:

\[\boxed{G_t^{\lambda s} \doteq \rho_t \!\left(R_{t+1} + \gamma_{t+1}\!\left[(1 - \lambda_{t+1})\, \hat{v}(S_{t+1}, \mathbf{w}_t) + \lambda_{t+1}\, G_{t+1}^{\lambda s}\right]\right) + (1 - \rho_t)\, \hat{v}(S_t, \mathbf{w}_t)}\] \[\begin{aligned} \text{where } \rho_t &= \frac{\pi(A_t \vert S_t)}{b(A_t \vert S_t)} \end{aligned}\]

The final $\lambda$-return can be approximated in terms of sums of the state-based TD error $\delta_t^s$, with the approximation becoming exact if the approximate value function does not change:

\[\begin{align*} G_t^{\lambda s} &\approx \hat{v}(S_t, \mathbf{w}_t) + \rho_t \sum_{k=t}^{\infty} \delta_k^s \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \\ \delta_t^s &\doteq R_{t+1} + \gamma_{t+1}\, \hat{v}(S_{t+1}, \mathbf{w}_t) - \hat{v}(S_t, \mathbf{w}_t) \end{align*}\]

The forward view update of the approximate $\lambda$-return is:

\[\begin{align*} \mathbf{w}_{t+1} &= \mathbf{w}_t + \alpha \!\left[G_t^{\lambda s} - \hat{v}(S_t, \mathbf{w}_t)\right] \nabla \hat{v}(S_t, \mathbf{w}_t) \\ &\boxed{\approx \mathbf{w}_t + \alpha \rho_t \!\left(\sum_{k=t}^{\infty} \delta_k^s \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i\right) \nabla \hat{v}(S_t, \mathbf{w}_t)} \end{align*}\]

We’re interested in the equivalence (approximately) between the forward-view update summed over time and a backward-view update summed over time. The equivalence is approximate because we ignore changes in the value function.
The sum of the forward-view update over time is:

\[\begin{align*} \sum_{t=0}^{\infty} \!\left(\mathbf{w}_{t+1} - \mathbf{w}_t\right) &\approx \sum_{t=0}^{\infty} \sum_{k=t}^{\infty} \alpha \rho_t\, \delta_k^s \nabla \hat{v}(S_t, \mathbf{w}_t) \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \\ &= \sum_{k=0}^{\infty} \sum_{t=0}^{k} \alpha \rho_t \nabla \hat{v}(S_t, \mathbf{w}_t)\, \delta_k^s \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \\ &\quad \left(\text{using the summation rule: } \sum_{t=x}^{y} \sum_{k=t}^{y} = \sum_{k=x}^{y} \sum_{t=x}^{k}\right) \\ &= \sum_{k=0}^{\infty} \alpha\, \delta_k^s \sum_{t=0}^{k} \rho_t \nabla \hat{v}(S_t, \mathbf{w}_t) \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \end{align*}\]

If the entire expression from the 2nd sum on could be written and updated incrementally as an eligibility trace, then the sum of the forward-view update over time would be in the form of the sum of a backward-view TD update.
- That is, if this expression was the trace at time $k$, then we could update it from its value at time $k-1$ by:

\[\begin{align*} \mathbf{z}_k &= \sum_{t=0}^{k} \rho_t \nabla \hat{v}(S_t, \mathbf{w}_t) \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \\ &= \sum_{t=0}^{k-1} \rho_t \nabla \hat{v}(S_t, \mathbf{w}_t) \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i + \rho_k \nabla \hat{v}(S_k, \mathbf{w}_k) \\ &= \gamma_k \lambda_k \rho_k \underbrace{\sum_{t=0}^{k-1} \rho_t \nabla \hat{v}(S_t, \mathbf{w}_t) \prod_{i=t+1}^{k-1} \gamma_i \lambda_i \rho_i}_{\mathbf{z}_{k-1}} + \rho_k \nabla \hat{v}(S_k, \mathbf{w}_k) \end{align*}\] \[\boxed{\mathbf{z}_k = \rho_k \!\left[\gamma_k \lambda_k\, \mathbf{z}_{k-1} + \nabla \hat{v}(S_k, \mathbf{w}_k)\right]}\]

If we change the index from $k$ to $t$ of the $\mathbf{z}_k$ equation above, we get the general accumulating trace update for state values:

\[\boxed{\mathbf{z}_t \doteq \rho_t \!\left[\gamma_t \lambda_t\, \mathbf{z}_{t-1} + \nabla \hat{v}(S_t, \mathbf{w}_t)\right]}\]

This eligibility trace combined with the usual semi-gradient TD($\lambda$) parameter-update rule (Section 12.2) forms a general TD($\lambda$) algorithm that can be applied to either on-policy or off-policy data:
- In on-policy, the algorithm is exactly TD($\lambda$) because $\rho_t = 1$ always and the ET above becomes the usual accumulating trace for variable $\lambda$ and $\gamma$:
\[\mathbf{z}_t \doteq \gamma_t \lambda_t\, \mathbf{z}_{t-1} + \nabla \hat{v}(S_t, \mathbf{w}_t)\]
- In off-policy, the algorithm stays as it is, although not guaranteed to be stable as a semi-gradient method.
- For off-policy, we’ll consider extensions that guarantee stability in the next few sections.
Let’s derive the off-policy ET for action-value methods and corresponding general Sarsa($\lambda$) algorithms.
- Starting with either recursive general action-based $\lambda$-return of Sarsa or Expected Sarsa, $G_t^{\lambda a}$, in Section 12.8 (Expected Sarsa works out to be simpler), we can extend the Expected Sarsa $G_t^{\lambda a}$ to the off-policy case after the off-policy model of action-based, off-policy, control variate, $n$-step return:

\[\boxed{\begin{align*} G_t^{\lambda a} &\doteq R_{t+1} + \gamma_{t+1}\!\left(\!\left[1 - \lambda_{t+1}\right] \bar{V}_t(S_{t+1}) + \lambda_{t+1}\!\left[\rho_{t+1} G_{t+1}^{\lambda a} + \bar{V}_t(S_{t+1}) - \rho_{t+1}\, \hat{q}(S_{t+1}, A_{t+1}, \mathbf{w}_t)\right]\right) \\ &= R_{t+1} + \gamma_{t+1}\!\left(\bar{V}_t(S_{t+1}) + \lambda_{t+1} \rho_{t+1} \!\left[G_{t+1}^{\lambda a} - \hat{q}(S_{t+1}, A_{t+1}, \mathbf{w}_t)\right]\right) \end{align*}}\] \[\begin{aligned} \text{where } \bar{V}_t(S_{t+1}) &= \sum_a \Pi(a \vert S_{t+1})\, \hat{q}(S_{t+1}, a, \mathbf{w}_t) \end{aligned}\]

The $\lambda$-return, approximately as the sum of TD errors, is:

\[\begin{align*} G_t^{\lambda a} &\approx \hat{q}(S_t, A_t, \mathbf{w}_t) + \sum_{k=t}^{\infty} \delta_k^a \prod_{i=t+1}^{k} \gamma_i \lambda_i \rho_i \\ \delta_t^a &= R_{t+1} + \gamma_{t+1} \bar{V}(S_{t+1}) - \hat{q}(S_t, A_t, \mathbf{w}_t) \end{align*}\]

Using steps analogous to those for the state case earlier in this section, write a forward-view update based on action-based $\lambda$-return $G_t^{\lambda a}$ above, then transform the sum of the updates using the summation rule and finally derive the eligibility trace for action values:

\[\boxed{\mathbf{z}_t \doteq \gamma_t \lambda_t \rho_t\, \mathbf{z}_{t-1} + \nabla \hat{q}(S_t, A_t, \mathbf{w}_t)}\]

This ET combined with the action-based, expected TD error $\delta_t^a$ and the usual semi-gradient TD($\lambda$) parameter-update rule (Section 12.2) forms an elegant, efficient Expected Sarsa($\lambda$) algorithm that can be applied to either on-policy or off-policy data:
- On-policy case: The algorithm becomes the Sarsa($\lambda$) algorithm given constant $\lambda$ and $\gamma$, and the usual state-action TD error:
\[\begin{aligned} &\quad \rho_t = 1, \quad \nabla\lambda_t = \nabla\gamma_t = 0 \\ &\boxed{\mathbf{z}_t \doteq \gamma\lambda\, \mathbf{z}_{t-1} + \nabla \hat{q}(S_t, A_t, \mathbf{w}_t)} \end{aligned}\]
At $\lambda = 1$, these algorithms become closely related to corresponding Monte Carlo algorithms.
No episode-by-episode equivalence of updates exist, only of their expectations, even under the most favorable conditions.
- Methods have been proposed recently [Sutton, Mahmood, Precup & van Hasselt, 2014] that do achieve an exact equivalence.
- These methods require an additional vector of “provisional weights” that keep track of executed updates but may need to be retracted/emphasized depending on future actions taken.
- The state and state-action versions of these methods are called PTD($\lambda$) and PQ($\lambda$) respectively, where the ‘P’ stands for Provisional.
If $\lambda < 1$, then all these off-policy algorithms involve bootstrapping and the deadly triad applies, meaning that they can be guaranteed stable only for the tabular case, state aggregation and other limited forms of function approximation.
Recall the challenge of off-policy learning has 2 parts. Off-policy eligibility traces deal effectively with the 1st part, correcting for the expected value of the targets, but not with the 2nd part that has to do with distribution of updates (matching off-policy to on-policy).
Algorithmic strategies for handling the 2nd part of the off-policy learning challenge with eligibility traces are summarized in Section 12.11.

12.10 Watkins’s Q($\lambda$) to Tree-Backup($\lambda$)

Watkins’s Q($\lambda$)

Watkins’s Q($\lambda$) is the original method for extending Q-learning to eligibility traces.
It involves decaying the ET in the usual way as long as a greedy action was taken, then cuts the traces to 0 after the 1st non-greedy action.

Backup Diagram for Watkins's Q($\lambda$)

The series of component updates ends either with the end of the episode or with the first nongreedy action, whichever comes first.

Tree-Backup($\lambda$)

Let’s look at the eligibility trace version of Tree Backup, which is called Tree-Backup($\lambda$) or TB($\lambda$).
TB($\lambda$) is the true successor to Q-learning because it has no importance sampling.
TB($\lambda$) concept is straightforward:
- The tree-backup updates of each length (Section 7.5) are weighted dependent on the bootstrapping parameter $\lambda$.
- Using the recursive form of the action-based $\lambda$-return for Expected Sarsa and then expanding the bootstrapping target case after the model of tree-backup $n$-step return (Section 7.5):
  \[\boxed{\begin{align*} G_t^{\lambda a} &\doteq R_{t+1} + \gamma_{t+1}\!\left(\!\left[1 - \lambda_{t+1}\right] \bar{V}_t(S_{t+1}) + \lambda_{t+1}\!\left[\sum_{a \neq A_{t+1}} \pi(a \vert S_{t+1})\, \hat{q}(S_{t+1}, a, \mathbf{w}_t) + \pi(A_{t+1} \vert S_{t+1}) G_{t+1}^{\lambda a}\right]\right) \\ &= R_{t+1} + \gamma_{t+1}\!\left(\bar{V}_t(S_{t+1}) + \lambda_{t+1} \pi(A_{t+1} \vert S_{t+1}) \!\left[G_{t+1}^{\lambda a} - \hat{q}(S_{t+1}, A_{t+1}, \mathbf{w}_t)\right]\right) \end{align*}}\]
$G_t^{\lambda a}$ can be approximated (ignoring changes in approx. value function) as a sum of TD errors:
\[\begin{align*} G_t^{\lambda a} &\approx \hat{q}(S_t, A_t, \mathbf{w}_t) + \sum_{k=t}^{\infty} \delta_k^a \prod_{i=t+1}^{k} \gamma_i \lambda_i \pi(A_i \vert S_i) \\ \delta_t^a &= R_{t+1} + \gamma_{t+1} \bar{V}_t(S_{t+1}) - \hat{q}(S_t, A_t, \mathbf{w}_t) \end{align*}\]
As always, using same steps as in the previous section, we get a special eligibility trace update involving the target-policy probabilities of the selected actions:
\[\boxed{\mathbf{z}_t \doteq \gamma_t \lambda_t \pi(A_t \vert S_t)\, \mathbf{z}_{t-1} + \nabla \hat{q}(S_t, A_t, \mathbf{w}_t)}\]

Backup Diagram for Tree Backup($\lambda$)

The tree-backup updates of each length are weighted in the usual way dependent on the bootstrapping parameter $\lambda$

The ET above combined with the usual semi-gradient TD($\lambda$) parameter-update rule defines the TB($\lambda$) algorithm.
Like all semi-gradient algorithms, TB($\lambda$) is not guaranteed to be stable when used with off-policy data and a powerful function approximator (the deadly triad).

12.11 Stable Off-Policy Methods with Traces

Let’s look at 4 of the most important methods that achieve stable off-policy methods/training using eligibility traces.
All 4 are based on either the Gradient-TD or Emphatic TD and linear function approximation.

GTD($\lambda$)

Analogous to TDC, and aims to learn a parameter $\mathbf{w}_{t}$ such that $\hat{v}(s, \mathbf{w}) \doteq \mathbf{w}_{t}^T \mathbf{x}(s) \approx v_{\pi}(s)$ even from data that is due to following another policy $b$. Its update is:
\[\begin{aligned} \mathbf{w}_{t+1} &\doteq \mathbf{w}_t + \alpha\, \delta_t^s\, \mathbf{z}_t - \alpha \gamma_{t+1}(1 - \lambda_{t+1})\!\left(\mathbf{z}_t^T \mathbf{v}_t\right) \mathbf{x}_{t+1} \\ \mathbf{v}_{t+1} &\doteq \mathbf{v}_t + \beta\, \delta_t^s\, \mathbf{z}_t - \beta\!\left(\mathbf{v}_t^T \mathbf{x}_t\right) \mathbf{x}_t \end{aligned}\] \[\begin{aligned} \text{where} \\ \mathbf{v} &\in \mathbb{R}^d \equiv \text{a vector of the same dimension as } \mathbf{w}, \text{ initialized to } \mathbf{v}_0 = \mathbf{0} \\ \beta &> 0 \equiv \text{a 2nd step-size parameter} \end{aligned}\]

GQ($\lambda$)

Gradient-TD algorithm for action values with eligibility traces.
GQ($\lambda$) aims to learn $\mathbf{w}_{t}$ such that $\hat{q}(s, a, \mathbf{w}_{t}) \doteq \mathbf{w}_{t}^T \mathbf{x}(s,a) \approx q_{\pi}(s,a)$ from off-policy data.
If the target policy is $\varepsilon$-greedy, or otherwise biased towards the greedy policy for $\hat{q}$, then GQ($\lambda$) can be used as a control algorithm.
GQ($\lambda$) update is:
\[\begin{aligned} \mathbf{w}_{t+1} &\doteq \mathbf{w}_t + \alpha\, \delta_t^a\, \mathbf{z}_t - \alpha \gamma_{t+1}(1 - \lambda_{t+1})\!\left(\mathbf{z}_t^T \mathbf{v}_t\right) \bar{\mathbf{x}}_{t+1} \\ \bar{\mathbf{x}}_t &\doteq \sum_a \pi(a \vert S_t)\, \mathbf{x}(S_t, a) \\ \delta_t^a &\doteq R_{t+1} + \gamma_{t+1}\, \mathbf{w}_t^T \bar{\mathbf{x}}_{t+1} - \mathbf{w}_t^T \mathbf{x}_t \\ \mathbf{z}_t &\doteq \gamma_t \lambda_t \rho_t\, \mathbf{z}_{t-1} + \nabla \hat{q}(S_t, A_t, \mathbf{w}_t) \end{aligned}\] \[\begin{aligned} \text{where} \\ \bar{\mathbf{x}}_t &\equiv \text{average feature vector for } S_t \text{ under the target policy} \\ \delta_t^a &\equiv \text{expectation form of the TD error} \end{aligned}\]

HTD($\lambda$)

Hybrid TD($\lambda$) state-value algorithm combines aspects of GTD($\lambda$) and TD($\lambda$).
HTD($\lambda$) is a strict generalization of TD($\lambda$) to the off-policy setting, meaning it reduces exactly to TD($\lambda$) when the behavior and target policies coincide; a property GTD($\lambda$) does not share:
\[b(A_t \vert S_t) = \pi(A_t \vert S_t), \quad \rho_t = 1 \implies \text{HTD}(\lambda) = \text{TD}(\lambda)\]
HTD($\lambda$) is defined by:
\[\begin{aligned} \mathbf{w}_{t+1} &\doteq \mathbf{w}_t + \alpha\, \delta_t^s\, \mathbf{z}_t + \alpha\!\left[\!\left(\mathbf{z}_t - \mathbf{z}_t^b\right)^T \mathbf{v}_t\right]\!\left(\mathbf{x}_t - \gamma_{t+1} \mathbf{x}_{t+1}\right) \\ \mathbf{v}_{t+1} &\doteq \mathbf{v}_t + \beta\, \delta_t^s\, \mathbf{z}_t - \beta\!\left(\mathbf{z}_t^T \mathbf{v}_t\right)\!\left(\mathbf{x}_t - \gamma_{t+1} \mathbf{x}_{t+1}\right), \quad & \mathbf{v}_0 \doteq \mathbf{0} \\ \mathbf{z}_t &\doteq \rho_t \!\left(\gamma_t \lambda_t\, \mathbf{z}_{t-1} + \mathbf{x}_t\right), \quad & \mathbf{z}_{-1} \doteq \mathbf{0} \\ \mathbf{z}_t^b &\doteq \gamma_t \lambda_t\, \mathbf{z}_{t-1}^b + \mathbf{x}_t, \quad & \mathbf{z}_{-1}^b \doteq \mathbf{0} \end{aligned}\]
We get
- a 2nd set of weights, $\mathbf{v}_t$.
- a 2nd set of eligibility traces, $\mathbf{z}_t^b$, conventional accumulating traces for the behavior policy.
\[\begin{aligned} \mathbf{z}_t^b = \mathbf{z}_t \text{ if all } \rho_t = 1 &\implies \left(\mathbf{z}_t - \mathbf{z}_t^b\right)^T = \mathbf{0} \\ &\implies \mathbf{w}_{t+1} = \mathbf{w}_t + \alpha\, \delta_t^s\, \mathbf{z}_t \quad \text{(TD(}\lambda\text{))} \end{aligned}\]

Emphatic TD($\lambda$)

Extension of one-step Emphatic TD (Sections 9.11 & 11.8) to eligibility traces.
The resulting algorithm:
- (+) retains strong off-policy convergence guarantees
- (+) enables any degree of bootstrapping
- (-) has high variance
- (-) potentially slow convergence.
Emphatic TD($\lambda$) is defined by:
\[\begin{aligned} \mathbf{w}_{t+1} &\doteq \mathbf{w}_t + \alpha\, \delta_t\, \mathbf{z}_t \\ \delta_t &\doteq R_{t+1} + \gamma_{t+1}\, \mathbf{w}_t^T \mathbf{x}_{t+1} - \mathbf{w}_t^T \mathbf{x}_t \\ \mathbf{z}_t &\doteq \rho_t \!\left(\gamma_t \lambda_t\, \mathbf{z}_{t-1} + M_t \mathbf{x}_t\right), \quad & \mathbf{z}_{-1} \doteq \mathbf{0} \\ M_t &\doteq \lambda_t \mathcal{I}_t + (1 - \lambda_t) F_t \\ F_t &\doteq \rho_{t-1} \gamma_t F_{t-1} + \mathcal{I}_t, \quad & F_0 \doteq \mathcal{i}(S_0) \end{aligned}\] \[\begin{aligned} \text{where} \\ M_t &\geq 0 \equiv \text{emphasis} \\ F_t &\geq 0 \equiv \text{followon trace} \\ \mathcal{I}_t &\geq 0 \equiv \text{interest} \end{aligned}\]
In the on-policy case ($\rho_t = 1$ for all $t$), Emphatic TD($\lambda$) is similar to conventional TD($\lambda$), but still significantly different:
- Emphatic TD($\lambda$) is guaranteed to converge for all state-dependent $\lambda$ functions; TD($\lambda$) is not (TD($\lambda$) is guaranteed only for constant $\lambda$).
- See Yu’s counterexample [Ghassian, Rafiee & Sutton, 2016].

12.12 Implementation Issues

Naive implementation seems expensive: Updating eligibility traces for every state at every time step appears computationally costly on serial computers.
Practical optimization: Most ET are nearly 0; only recently visited states have significant traces, so implementations can track and update only these few states.
Computational cost: With this optimization, tabular methods with traces are only a few times more expensive than one-step methods.
Function approximation reduces overhead: When using neural networks, ET typically only double memory and computation per step (much less overhead than in tabular case).
Tabular is the worst case: The tabular setting represents the highest computational complexity for ET relative to simpler methods.

12.13 Conclusions

Eligibility traces provide an efficient, incremental way to interpolate between TD and MC methods.
ET offer advantages over $n$-step methods in terms of generality and computational trade-offs.
Empirically, an intermediate mix works best: ET should move towards MC but not all the way since pure MC performance degrades sharply.
ET are the first defense against long-delayed rewards and non-Markov tasks, used with TD methods to make them behave more like MC methods without full bootstrapping.
Use traces when data is scarce and online learning is required, as they provide faster learning per sample despite higher computational cost per step.
Avoid traces in offline settings with cheap abundant data (maximum data processing speed matters more than learning efficiency per sample).
True online methods achieve ideal $\lambda$-return performance while maintaining $O(d)$ computational efficiency.
Forward-to-backward view derivations provide computationally efficient, mechanistic, practical implementations of theory.

Citation

If you found this blog post helpful, please consider citing it:

@article{obasi2026RLsuttonBartoCh12notes,
  title   = "Sutton & Barto, Ch. 12: Eligibility Traces (Personal Notes)",
  author  = "Obasi, Chizoba",
  journal = "chizkidd.github.io",
  year    = "2026",
  month   = "Mar",
  url     = "https://chizkidd.github.io/2026/03/13/rl-sutton-barto-notes-ch012/"
}

Chizoba Obasi blog

FlashAttention: Part 2

Table of Contents

Appendix

1. FlashAttention Evolution

1.1 FlashAttention-2: What Changed

1.2 FA2 Parallelism Across Sequence Tiles

1.3 FA2 Work Partitioning and Non-Matmul FLOPs

1.4 FlashAttention-3: The Hopper Generation

1.5 FA3 Asynchrony: Overlap Data Movement, GEMM, and Softmax

1.6 FA3 FP8: Performance Without Pretending Precision Is Free

1.7 FlashAttention-4: The Blackwell Generation

1.8 FA4 and Asymmetric Hardware Scaling

1.9 FA4 Implementation and Current Status

1.10 FlashAttention-1 Through FA-4 Compared

2. Using FlashAttention in Frameworks

2.1 PyTorch Scaled-Dot-Product Attention Today

2.2 Exactness Is Not Bitwise Identity

3. FlashAttention vs. Other Techniques

3.1 FlashAttention vs. PagedAttention

3.2 FlashAttention vs. Sparse and Linear Attention

4. Training vs. Inference

4.1 Training, Prefill, and Decode Are Different Regimes

5. Practical Engineering for FlashAttention

5.1 Common Implementation Mistakes

5.2 Common Misconceptions

6. Conclusion

6.1 The Mental Model to Have

6.2 FlashAttention Summary

Appendix

Reference Map

References

Citation

FlashAttention: Part 1

Table of Contents

Appendix

0. Introduction

0.1 How to Read This Handbook

0.2 Notation

1. The Fundamental Problem

1.1 What FlashAttention Actually Optimizes

1.2 The Attention Equation Is Not the Implementation

1.3 GPU Memory Hierarchy and Why IO Matters

1.4 Why Materializing the Score and Probability Matrices Is Expensive

1.5 Dense Arithmetic Is Still Quadratic

1.6 Memory-Efficient Exact Attention Predates FlashAttention

2. The Mathematical Trick

2.1 Tiling Queries, Keys, and Values

2.2 Softmax Is the Difficult Part of Streaming

Worked example 1: why subtract the max

Worked example 2: streaming two blocks

2.3 Online Softmax from First Principles

2.4 The Blockwise Merge Recurrence

2.5 A Complete Numerical Example

3. Putting the Mathematics onto the GPU

3.1 FlashAttention Forward Pass

3.2 Why Dense FlashAttention Is Exact

3.3 IO Complexity: What the Theorem Actually Says

3.4 Memory Complexity: Linear Auxiliary State, Not Linear Compute

3.5 Causal Masking and Tile Skipping

3.6 Backward Pass: Recompute Instead of Save

3.7 Why More FLOPs Can Still Be Faster

4. Architectural Compatibility

4.1 MHA, MQA, and GQA Compatibility (Shared K/V Heads)

4.2 Variable Lengths, Local Attention, and Dropout

5. Summary

Appendix

Reference Map

References

Citation

Mixture of Experts (MoE): How Transformers Scale Without Activating Everything.

Table of Contents

Appendix

Feed-Forward Networks in Transformers

The decoder-only transformer architecture.

A Useful Intuition for FFN

RMSNorm

The FFN as a Knowledge Store

Why make the FFN bigger?

Why Mixture of Experts?