The Extender: A Log-Structured Transformer

The Extender. A Log-Structured Transformer. By Jakob Eriksson, from the University of Illinois Chicago. Published on the twenty-sixth of September. No code repository is linked to it.

In a standard Transformer, each layer communicates with the layers after it only through the residual, h, which is a superposition channel: every layer's contribution is added into the same shared vector. Attention then has to keep a persistent memory for every layer, sized at twice the layer count times the model width, and that footprint is what the Extender sets out to shrink.

The Extender is a log-structured variant of the Transformer. Alongside the residual h it adds a second channel, x, which works by concatenation rather than addition. Every layer emits two things: a residual update that is added to h, as usual, and a much smaller extension that is appended to x. That is the sense of log-structured. The x channel is a record that grows as each layer appends to it.

The split determines what attention sees. The feed-forward network and the queries both read h, but the key and value projections take only x as input. Once x is fully extended, then, it contains the complete input to the key and value projections of every layer. The persistent attention memory footprint stops depending on the model width per layer, and becomes simply the summed size of the extensions.

With extensions of size thirty-two, the Extender matches Transformer accuracy on short-context CORE tasks, at sizes from one hundred ninety-nine million to nine hundred twenty-four million parameters. On long-context RULER workloads, again at nine hundred twenty-four million parameters, it exceeds Transformer accuracy. For that model, which is one thousand six hundred sixty-four wide, the persistent attention memory footprint is a hundred and four times smaller than MHA.

Two channels, then, with different jobs. The residual keeps carrying the superposed state the feed-forward network and the queries read, and a small appended log carries everything the key and value projections need. The author reports that the memory savings grow with model width.