Songlin Yang Explains Attention Mechanisms from DeepSeek, Kimi, and MiniMax

Zhang Xiaojun invites MIT CSAIL doctoral student Songlin Yang to explain attention papers from DeepSeek, Kimi, and MiniMax. Yang studies hardware-efficient sequence modeling, especially linear attention.

Long chains of thought can increase computation at inference, but ordinary self-attention has quadratic sequence-computation costs and a KV cache that grows with length. The three papers approach these costs differently.

DeepSeek and Kimi use dynamic sparse attention: different queries can select different key-value blocks instead of following a predetermined sparse pattern. MiniMax combines linear attention with softmax attention.

Yang first explains attention: each token produces a query, key, and value. A query scores accessible keys, applies softmax, and uses the resulting weights to combine values. An autoregressive model masks future tokens.

He says both prefill for long inputs and token-by-token decoding need optimization. Long-document processing makes prefill especially relevant for Kimi, while long reasoning processes also make decoding efficiency important.

DeepSeek’s Native Sparse Attention (NSA) aims to make dynamic sparse attention efficient to pretrain and compatible with hardware and inference architectures. Yang highlights training from scratch as a key difference from changing attention only at inference.

Yang introduces query-aware selection of key-value blocks: instead of every position reading the entire history, the current query selects blocks likely to be useful. Actual speed also depends on block organization and memory access.

NSA has three branches: coarse attention over compressed blocks, fine attention over selected blocks, and a sliding window over nearby tokens. Learned gates determine how each branch contributes to the output.

The compression branch represents each block and scores its relevance to the query. The selection branch then attends more finely to high-scoring blocks, preserving a global overview without fine-grained computation over every token.

The sliding-window branch uses the observation that nearby tokens are often important in language. Asked how the branches balance global and local information, Yang says the model learns the gate weights rather than using fixed human-assigned proportions.

He emphasizes that compression, block selection, and sliding windows are not all new ideas. The paper’s contribution lies in combining and implementing them to use hardware effectively during training and inference.

To explain grouped-query attention (GQA), Yang compares multi-head and multi-query attention. Multi-query attention shares keys and values across query heads to shrink the inference KV cache; GQA sits between fully separate and fully shared representations.

In GQA, if heads in a group select different KV blocks, the system must read the union of their selections. Sharing a selection within the group can reduce reads, but limits each head’s freedom to choose independently.

NSA aggregates attention scores within a group and selects shared blocks. Yang sees this as a deliberate tradeoff: fewer memory accesses improve hardware efficiency, while architectural choices seek to preserve sufficient expressive power.

He reviews FlashAttention’s tiled computation: load portions of queries, keys, and values into fast memory and compute within blocks, reducing transfers of intermediates to and from GPU memory. Data movement can bottleneck speed alongside arithmetic.

Dynamic sparsity introduces a problem: different positions in a query block may need different KV blocks. Simply retaining contiguous tiling can require reading their union, wasting bandwidth. Sparsity does not automatically mean faster hardware execution.

NSA packs multiple query heads sharing KV into matrix operations, reducing unnecessary reads while using GPU Tensor Cores. Yang uses this to show why an algorithm’s expression and its hardware execution must be considered together.

Matrix multiplication also has size constraints. To satisfy hardware requirements while retaining selection diversity across groups, the paper increases the number of heads and expands the projection. Yang calls this a bold but purposeful architectural choice.

In the paper’s reported training loss and several evaluations, NSA performs well against its full-attention baseline and is better on some measures. Yang finds this compelling, but the results are comparisons under the paper’s experimental conditions.

He distinguishes two approaches: approximating attention at inference after training a full-attention model can introduce error, while pretraining a sparse architecture lets the model adapt to it. Whether it is better still requires validation across settings.

He also compares retaining the entire KV cache with evicting cache entries: discarded information might be needed later. NSA dynamically selects blocks to read, which is different from permanently deleting unselected history.

The paper’s performance plot compares kernels with similar implementation approaches. Yang points to speedups on long sequences and emphasizes that decoding may be memory-bandwidth-bound; actual gains depend on hardware, length, and implementation.

In the ablation discussion, the compression branch supplies coarse information and also provides a training signal for block selection. Yang considers this useful for end-to-end training; looking only at computation on selected blocks misses part of the design.

Yang describes the contribution mainly as engineering and hardware-software co-design: existing ideas are adapted to shared KV, tiled reads, and matrix multiplication. He particularly admires the kernel design.

Zhang asks why he calls it a “brutal hardware aesthetic.” Yang says shared selections, more heads, and multiple branches fit together to balance speed and model quality under hardware constraints.

The second paper is Kimi’s MoBA. Yang says it also selects blocks based on queries and supports pretraining, but uses a simpler design: mean pooling represents blocks, followed mainly by attention over selected blocks.

MoBA omits NSA’s compression-output and sliding-window-output branches, and different attention heads can select their own KV blocks. Yang says this brings simplicity and flexibility, but different implementation costs from NSA.

To reuse FlashAttention, MoBA must gather queries selecting the same KV block into contiguous tensors. Indexing and reordering have costs of their own, which can become significant when blocks are too small.

Comparing the papers’ block-size choices, he notes that larger blocks can reduce reordering overhead but make selection coarser. Smaller blocks can be more flexible while increasing scheduling and memory-operation costs.

MoBA encounters sparse training signals during supervised fine-tuning: when loss applies mainly to the answer, unselected prefix blocks may receive insufficient gradients. The paper mitigates this by using full attention in the final layers.

Because MoBA’s block selection adds no extra parameters, it can switch relatively easily between sparse and full attention. Yang corrects his initial reading: the mixed-training example is an ablation in the paper.

The Kimi paper also evaluates longer contexts and discusses practical use. Yang frames the designs as different priorities: NSA stresses hardware-software alignment, while MoBA favors simplicity. A universal winner, he says, should not be assumed.

Both approaches seek faster training and inference, but benefits differ by sequence length and implementation. Yang notes that contiguous block reads and matrix operations are common hardware principles; theoretical complexity alone is insufficient.

The third paper is MiniMax-01. Its hybrid architecture places one softmax-attention layer after every seven Lightning Attention layers, repeating the pattern across 80 layers. Linear attention is Yang’s main research area.

To explain linear attention, he starts with a simplified form without softmax. Associativity lets the outer products of past keys and values accumulate into a matrix state, which the current query then reads.

This state can be written as the hidden state of a recurrent neural network: each step incorporates new key-value information into a fixed-size matrix. Unlike a KV cache that retains the full history, the past is compressed into the state.

Yang compares matrix states with traditional vector states in terms of memory capacity and notes that expanding state must also be hardware-efficient. Linear theoretical complexity does not guarantee efficient GPU utilization during training.

Computing token by token limits parallelism and makes efficient matrix multiplication difficult, while computing over the whole sequence can incur quadratic cost. A chunkwise form seeks the benefits of both.

The chunkwise algorithm divides the sequence into chunks: historical information passes through states between chunks while local contributions are computed in parallel. Yang stresses that this reformulation is mathematically equivalent, not necessarily a lossy approximation.

Within each chunk, multiple token updates can be packed into matrix multiplication, and queries can batch their reads of the prior state. This makes linear attention easier to map onto current GPU compute units.

With fixed chunk size, Yang says, this training approach avoids constructing quadratic interactions across the whole sequence. He mentions long video sequences as a possible application, with actual results dependent on the task and implementation.

He introduces the open-source Flash Linear Attention library, which implements several hardware-efficient linear-attention models. Beyond chunk computation, input/output and memory access also require optimization.

Early linear attention was often questioned for its language-modeling performance. Yang says decay and gating have improved it, and compares Lightning Attention’s fixed decay with data-dependent decay.

He explains Mamba-2 through a linear-attention lens, connecting seemingly different recurrent, parallel, and chunkwise expressions. The models nevertheless retain their own parameterizations and design choices.

Hybrid architectures seek to combine two strengths: a fixed-size recurrent state compresses history, while full attention retains a KV cache that can retrieve details on demand. Yang likens the latter to consulting a book when needed.

Pure recurrent or linear attention can be constrained by fixed state capacity on exact retrieval tasks. Adding some softmax-attention layers may compensate; he also discusses how recurrence can provide position-related information.

MiniMax’s hybrid model still has some full-attention layers, so long-sequence costs do not disappear entirely, but the many linear layers can reduce them. Yang sees the paper as evidence that the hybrid approach can scale to a large model.

He credits chunkwise algorithms, expert parallelism, pipelines, and other engineering optimizations with enabling scale. Earlier smaller hybrid models provided precedents; one contribution of MiniMax was extending the approach to a larger scale.

Asked which approach takes greater architectural risk, Yang says dynamic sparsity that retains the KV cache can feel more familiar, while scaling a linear/full-attention hybrid still faces retrieval concerns despite prior examples.

He cautions that models differ in training data, scale, and experimental conditions, making it hard to rank architectures solely from published results. Fairer conclusions require broader comparisons under matched conditions.

On companies publishing papers and opening technology, Yang says this can build technical influence and help challengers create communities. His views about specific commercial benefits are personal assessments.

Of the three papers, he favors NSA: it enables efficient training of dynamic sparse attention and shows competitive results in the paper’s comparisons. Kimi’s minimalist design also impresses him.

Looking ahead in architecture research, he mentions positional encoding, sparse attention, linear attention, and long convolutions. Research topics can rise and fall, so they need continuing experimental evaluation rather than judgments based on a moment’s popularity.

Finally, Yang thinks technical progress by Chinese teams could make them more attractive to researchers overseas. Zhang closes by inviting listeners to continue the discussion.