Hybrid Recurrent-Attention

The comparison that settles the shelf is not recurrence against attention. It is that each fails at exactly what the other does well, and the failures are complementary enough that mixing them is not a compromise.

Recurrence decodes in constant memory and cannot recall a specific distant token. Its state is a fixed-size summary; retrieving a phone number from forty thousand tokens back requires that the number survived the summarisation, and nothing guarantees it did.

Attention recalls anything at all and pays a cache that grows with context — the thing that makes long-context serving expensive.

The ratio

At $r = 8$ the KV cache is an eighth of a transformer’s, which is the difference between 256k context fitting on one device and not. And the quality curve is flat over a wide range of $r$ — everything between 1:4 and 1:8 performs about the same, so the ratio is chosen by the memory budget rather than by ablation.

Jamba’s 1:7 with MoE in the recurrent layers reaches 256k context on a single 80GB device. Griffin interleaves gated linear recurrences with local attention, so even its attention layers have bounded cache, and matches Llama-2 quality at 6B while training on six times fewer tokens.

Where the attention layers go

Not evenly, in the designs that ablate it. The first block is attention in most of them, on the reasoning that the model needs unrestricted mixing before the recurrent state has anything worth summarising; the last block is usually recurrent. Beyond those two the placement within the middle appears not to matter much, which is consistent with what the layer-ablation literature finds generally.

from torch import Tensor, nn


class HybridStack(nn.Module):
    """One attention layer in every `ratio`; the rest recurrent."""

    def __init__(self, dim: int, layers: int, ratio: int = 8, window: int = 4096):
        super().__init__()
        self.blocks = nn.ModuleList(
            # Attention first, then every ratio-th layer. The last stays recurrent.
            SlidingWindowAttentionBlock(dim, window=window)
            if i == 0 or (i % ratio == 0 and i != layers - 1)
            else MambaBlock(dim)
            for i in range(layers)
        )

    def forward(self, x: Tensor) -> Tensor:
        for block in self.blocks:
            x = block(x)
        return x

The serving consequence is that two kinds of state now have to be managed together: a KV cache that grows and is paged, and a recurrent state that is fixed per sequence. Paged attention handles the first and was not designed for the second, so hybrid inference stacks carry a second allocator — small, but a separate thing to get right, and the usual source of bugs when a sequence is pre-empted and resumed.

Every entry here has been a route to the same destination: linear-time sequence mixing with a constant decode state. All of them arrive, and all of them arrive with the same limitation, because it is a property of fixed-size state rather than of any particular parameterisation.

The honest summary of five years of work is that recurrence did not replace attention and was never going to. It replaced most of it, which is a smaller claim and a more useful one — the cache is where the money goes, and cutting it eightfold while keeping the recall is worth more than an architecture with a better story.