Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 4: Attention Alternatives — Annotated Lecture Notes & Key Takeaways

Course Information

Institution: Stanford

Original Course: Stanford CS336 Language Modeling from Scratch | Spring 2026 | Lecture 4: Attention Alternatives

Instructor Bio: This lecture is delivered by **Percy Liang**, Associate Professor of Computer Science at Stanford University and core faculty of the Stanford Institute for Human-Centered Artificial Intelligence (HAI). Percy Liang leads the Stanford Natural Language Processing Group, with research spanning the theoretical foundations and practical systems of language modeling, machine learning alignment, compositional semantics, and NLP evaluation. He received his PhD in Computer Science from the University of California, Berkeley and his BA in Mathematics from Harvard University. His work has been recognized with the NSF CAREER Award, Google Faculty Research Award, and multiple best paper awards at top-tier machine learning and NLP conferences. He is also widely known for creating influential benchmarks and open-source frameworks for language model research and assessment.

Course Description: This lecture explores alternatives to standard full self-attention, addressing the quadratic computational and memory bottlenecks that limit long-context language modeling. It covers a range of efficient attention and sequence modeling approaches, including linear attention, sliding window attention, sparse attention patterns, and state-space models (SSMs) such as Mamba. The session systematically compares each approach in terms of computational complexity, memory footprint, and modeling quality, and discusses their respective suitability for different context lengths, training scenarios, and deployment targets.

These are my own notes from the recording, paraphrased throughout — not a transcript, and definitely not the lecturer's words. If you missed class, the short version is that the announced topic, attention alternatives, turns out to be the smaller half of the session. The professor spends the first stretch on ways to break the quadratic cost of attention, then pivots for the longer stretch into mixture of experts. He flags the split explicitly at the top: the first part modifies the attention block, the second modifies the MLP block.
 

Why Longer Context Is a Cost Problem First and a Modeling Problem Second

 

The framing is blunt. People want longer context because they want to shove more knowledge into the window, or because they're running agents that operate over a lot of material. He points at a plot of context window size against release date, drawn on a log scale, and describes it as a rush among the major model vendors to keep offering bigger windows.
 

The more interesting slide is the second one: the ratio of compute spent in the feed-forward part versus the attention part as sequence length grows. Feed-forward starts out as the dominant cost and grows only linearly. Attention is an all-to-all comparison between every pair of positions, so it grows quadratically and eventually swallows everything else. That's the whole motivation in one picture — at short lengths attention is a rounding error, at long lengths it is the problem.
 

Two toolkits get named for controlling it. The first is architectural: use local attention most of the time and mix in global attention sparingly, say once every eight layers. The second is systems engineering, and he gets mildly worked up about how underrated it is. There's a dig here at people with a theory background who fixate on big-O. Flash attention is his evidence: mathematically it does nothing about the quadratic, all it does is reorder the computation to cut memory traffic, and yet it produces roughly two-fold speedups over the stock PyTorch path and, at lengths where the naive version simply won't fit in memory, makes the computation possible at all. He keeps repeating that constant factors matter enormously.
 

But constant factors run out. At five or ten million tokens, he says, the tricks probably aren't enough, which is what motivates the hunt for genuinely linear scaling. A throwaway remark worth catching: he notes this is the first year he's teaching linear-time attention at all, because only in the last couple of years have these methods been validated at production scale.
 

Linear-Time Attention Alternatives: One Algebraic Trick and Everything That Grew From It

 

Reassociating the product

 

Everything in this section comes from one idea, and he says so: the associativity of multiplication.
 

Write attention compactly as queries, keys, and values with a softmax in the middle. Now pretend the softmax isn't there. Once it's gone, you're free to move the parentheses. Instead of multiplying Q by K-transpose first and then hitting the result with V, you multiply K-transpose by V first and apply Q afterward. The cost changes character completely: the bad term goes from N-squared-times-d to N-times-the-product-of-the-two-feature-dimensions. N is the context length and runs into the millions. The feature dimensions are in the thousands or tens of thousands. Swapping one dependency for the other is the entire trick.
 

The same math wearing an RNN costume

 

The second observation, and the one he credits with launching most of the follow-on research, is that the right-hand multiplication in the reassociated form is just a recurrence. Sweep left to right across the sequence, fold each position's key-value outer product into a running state, carry that state forward, and read out with the current query. This is equivalent to the dense form.
 

That equivalence is what makes the whole family practical. The dense form parallelizes, so it's what you want for training. The recurrent form carries a fixed-size state, so it's what you want at inference. You get to pick per use case. He then says something I had to replay — that the result "is linear, which is not very good." I read that as a slip for "the linear approximation is a weak model," given everything he says afterward about hybrids, but the sentence as delivered is genuinely confusing.
 

It does work in production, though. MiniMax M1 uses a 7:1 pattern, seven linear attention layers per one full softmax layer, and lands in the same neighborhood as models like OpenAI's o3 or DeepSeek R1. And the caveat he attaches is important: nobody has demonstrated a fully linear attention stack at scale. Every example in the lecture is a hybrid.
 

Mamba 2 adds a forget gate

 

If plain linear attention feels too crude, the next step is to gate the recurrence. Mamba 2, from Albert Gu, Tri Dao and collaborators, came out of state-space theory, but he insists the mechanics are just a small elaboration of what's already on the board. You add a gate, gamma, that scales how much of the previous state survives into the next step. The gate reads only the current input, never the state, and that single restriction is what keeps the parallel training form available. The intuition is straight out of LSTMs: sometimes you should carry information forward, sometimes you should let it decay to zero.
 

Nemotron 3 interleaves Mamba 2 layers with occasional softmax attention and holds up reasonably against Qwen3-thinking and GPT-OSS, with good throughput at long contexts. He tempers this: these are small frontier models, not the largest ones.
 

Gated DeltaNet adds an erase step

 

Push the gating idea one step further and you get Gated DeltaNet, which he calls the most widely used state-space design right now and the one scaled up in Qwen 3.5. Relative to Mamba 2 there's a second gate, beta, which decides whether the current input gets written into the state at all.
 

The distinctive part is what happens on the write. Rather than blending new information in, the update includes a projection term, identity minus beta times the outer product of the current key with itself, which clears out whatever was previously stored along that key direction before the new value goes in. He's candid that this isn't quite a proper projection since there's no normalization, but the picture holds. He also notes that the same update has been rediscovered independently through meta-learning least-squares formulations, fast weight programming, and test-time training — different starting points, identical destination.
 

Qwen 3.5 and its Qwen Next predecessors run a 3:1 Gated DeltaNet hybrid, decode faster than Qwen 3 as context grows, and give up very little in quality.
 

How far you can push the hybrid ratio

 

This is the part of the attention half that I found most useful and most unsettled. Good controlled comparisons of hybrid architectures are scarce. He cites essentially one, from ByteDance Seed and UC Santa Cruz, sweeping the ratio of recurrent to full-attention layers. His own assessment of the results is that they're messy.
 

The pattern he pulls out: at low ratios of recurrent layers the best designs lose almost nothing, then past some threshold long-context performance degrades, and a fully recurrent stack degrades badly across the board. He also warns that some evaluation tasks, single-key retrieval in particular, are things these architectures are explicitly built to ace, so treat those panels skeptically; the QA curves tell the same story more honestly.
 

A question from the room sharpens the picture. Where exactly is the loss? His answer: the step from softmax attention to linear attention is the lossy one. Everything after that — between the linear form and the recurrent form — is exact.
 

A Different Bet: Let a Cheap Indexer Throw Most of the Context Away

 

Sparse attention is presented as a genuinely different philosophy rather than a variant of the above. DeepSeek introduced it in V3.2 (he makes an accidental pun about the acronym and moves on). A lightweight indexer scans the long context and picks a small subset of positions, and full attention then runs only over that subset. Mechanically the indexer takes query-key inner products, applies a ReLU, weights them by the preceding tokens, and takes the top K.
 

Two things make this cheap in practice. First, you don't have to pretrain with it — you train a normal transformer, then bolt the indexer on during the long-context extension phase that everyone runs anyway, between short-context pretraining and post-training. He admits it's surprising that grafting on a non-differentiable top-K module at that stage works. Second, K is bounded and chosen to look more like a short-context length than the full input length.
 

The validation is that V3.2 is competitive with frontier models of its moment, and that GLM 5 adopted the same approach with published ablations showing full DSA training costs little even on long-context retrieval, the exact regime where recurrent hybrids struggle.
 

Crucially, this is not linear time. The indexer does brute-force inner products over everything, so it's quadratic. The savings come from making it low-precision, low-dimensional, and small, and from running the second, precise stage over a much shorter sequence. That's the constant-factors argument again, and a student pushes back on exactly this point in the Q&A; the answer is essentially that quadratic-but-tiny beats your intuitions about asymptotics.
 

Mixture of Experts: More Parameters Than You Pay For

 

He opens by deflating the concept: an MoE is just a more efficient MLP. Replace one feed-forward block with several, add a mechanism that picks one per input, and you hold per-token FLOPs flat while multiplying parameter count. The original motivation was bluntly parameter-centric — you want more parameters, you don't want to pay for them.
 

The argument for why anyone bothers is empirical. Work from Fedus and colleagues in 2022 showed test loss falling as expert count rises with active parameters held fixed, and faster progress at fixed training compute. The AI2 Olmo study reproduced it, reporting something like a two-fold training speedup over dense. DeepSeek V2 was the moment the industry noticed: fewer active parameters, comparable or better MMLU.
 

Routing: token choice won, and nobody fully knows why

 

Routing happens per token, and the router itself is embarrassingly simple — a single matrix multiply. Tokens pick their top K experts, and that's what nearly every deployed model does. Expert choice, where each expert selects its favorite tokens, also trains fine but scores worse on validation loss and downstream benchmarks in the Olmo comparison.
 

Two alternatives come up that never made it into production. Hashing inputs to experts gives real if smaller gains, and he admits he finds it a little mysterious that it works at all — it's a common baseline in papers and essentially unused in deployment. Reinforcement learning over routing decisions is the theoretically natural framing, going back to Bengio's work in 2013, and it does function, but the overhead and gradient variance aren't worth it. There's also global assignment via solving a linear assignment problem, which is elegant and far too expensive at scale.
 

Fine-grained and shared experts

 

DeepSeek's contribution here is now close to universal. Cut the experts into smaller, finer-grained pieces, and designate a subset as shared experts that bypass the router and process every token. The reasoning is that in a classic design, routed experts waste capacity relearning common processing; a shared expert absorbs that and lets the rest specialize. Their ablations show gains from both changes, with pronounced jumps on TriviaQA and Natural Questions.
 

He's careful to report a disagreement: the Olmo team's controlled study agrees that fine-grained experts help but concludes shared experts don't help much. Most current open models follow the DeepSeek recipe anyway. The parallel he draws is that Llama's dense design became the default template, and DeepSeek's MoE design did the same for sparse models.
 

Training: non-differentiable routing held together by one balancing loss

 

Here's the crux. You need sparsity during training or you lose the efficiency that justified the architecture. But sparsity means the gating decision is non-differentiable and you never observe the experts you didn't pick. It smells like a bandit problem, and he says so, and then says nobody solves it that way.
 

Two historically interesting attempts come first. Shazeer's early work injects input-scaled noise into the routing logits before top-K, which acts as a stochastic tiebreak and lets the softmax learn a ranking. Fedus and colleagues applied a uniform multiplicative perturbation to harden the experts. Later Google work dropped these, and ablations suggest removing the stochastic machinery actually improves both stability and final quality.
 

What actually runs is a balancing loss bolted onto the language modeling objective. The Switch Transformer version multiplies, for each expert, the fraction of tokens dispatched to it by the router's total probability mass allocated to it, summed over experts. He admits this isn't something you'd derive from first principles and recommends reading its gradient instead: the derivative with respect to an expert's probability mass is that expert's token fraction, so the term pushes probability mass down on experts in proportion to how popular they already are.
 

The failure mode it prevents is what he calls expert collapse — a rich-gets-richer spiral where experts that get picked get stronger, get picked more, and eventually absorb everything. DeepSeek's early models used the same per-expert loss plus a second device-level loss, so that the GPUs holding the experts stay evenly utilized. V3 moved toward a per-expert bias updated by an online-learning-style trick, pitched as auxiliary-loss-free, though he notes they still end up adding some auxiliary losses to prevent extreme imbalance, and no one has eliminated them entirely.
 

The Olmo ablation is the money shot. Strip the balancing loss out and losses get worse, and the utilization panels show almost every token collapsing onto two experts while the rest sit idle — a large fraction of the model's parameters doing nothing. Put it back and utilization evens out.
 

What he finds genuinely surprising, and I agree it's the strangest part of the lecture, is that this is enough. A hard top-K selection, non-differentiable, plus one hand-written balancing penalty, and then you just push gradients through as if nothing were wrong, and it trains well. His explanation is that two dynamics cancel: useful experts get reinforced, and the balancing loss spreads traffic out.
 

He also closes the loop with the earlier section — the same top-K-plus-auxiliary-loss pattern shows up in DeepSeek's sparse attention and in work on removing tokenizers. It's becoming a general architectural ingredient.
 

Systems, Stability, and Where the Lecture Felt Thinnest

 

MoEs add a third axis of parallelism alongside data and model parallelism, and sparsity maps onto structured matrix multiplications that hardware handles well. A recent Nemotron 3 trick addresses the communication cost of shipping activations to the expert that owns them: down-project the residual stream before the collective communication, while leaving the shared expert in the higher dimension since it doesn't need shipping. Older infrastructure had a nastier problem — overloaded experts would build queues and silently drop tokens, returning zeros, so that another user's traffic could change your output. Dropless implementations like MegaBlocks have fixed this.
 

Stability gets a short treatment. The router adds another softmax, and softmaxes are where exponentials and divisions live. Running the router in float32 and adding a Z-loss are the standard mitigations, and the Olmo curves show visibly spikier training without it. Fine-tuning is its own headache: the parameter count invites overfitting, with a large train-validation gap on the GLUE example he shows, and the common workaround is to fine-tune only the attention or only the non-MoE layers.
 

The closing walkthrough traces DeepSeek V1 through V3. V1 is the archetype: shared and fine-grained experts, top-K routing, auxiliary balancing. V2 scales it and adds routing and communication losses aimed at the hardware — his comment is that successful training isn't just deep learning, it's respecting your system. V3 changes the balancing mechanism and the expert weighting, and adds multi-head latent attention, which caches a compressed latent instead of the full keys and values, with a complication around positional encoding. Multi-token prediction gets a mention as both a statistical idea and a free speculative decoder. Upcycling — copying a trained dense MLP into several experts and letting them specialize — is presented as a cool idea that has fallen out of favor, since nobody trains a dense model first anymore.
 

Two places left me unconvinced. The hybrid-ratio evidence rests on what he himself calls one messy study, yet the recommendation to mix in a softmax layer every few layers is stated with a lot of confidence. And the balancing loss is presented as working "for some reason," with the mechanistic explanation — reinforcement versus equalization canceling out — offered as narrative rather than measurement. Neither is a reason to dismiss the lecture; both are places where I'd want to read the underlying papers before trusting the rule of thumb.
 

Still, the through-line holds up well. Attention alternatives and sparse Mixture-of-Experts layers are the two places where modern models buy capacity they don't pay for at inference, and the lecture's real message is that most of the wins come from simple tricks applied with a lot of systems discipline, not from cleverer math.
 

Content Disclaimer:

This article is for general reference only and does not constitute professional R&D guidance, production process advice or quality certification. All material performance data has specific test premises; readers should verify parameters against actual equipment and working conditions.