LooM: Loop-Invariant Memory Makes Loop Transformer Inference Budget Adjustable

Loop Transformers increase effective depth by repeatedly executing shared layers, but memory grows with the number of loops, computation can only be adjusted in whole blocks, and the loop count is hard to change after training—these three limitations have long hindered their adoption. A paper titled "LooM: Loop-Invariant Memory for Causal Encoder-Decoder Language Models" proposes a targeted solution: loop only in the decoder layers of a causal encoder-decoder (CED), and decouple memory writes from the loop count.

Three Bottlenecks of Loop Transformers

The basic idea of a loop Transformer is to have the same set of layers executed repeatedly, achieving the effect of a deeper network with fewer parameters. But the research points out three unavoidable limitations on this path:

• Memory grows with the loop count: each loop round must retain its corresponding cache, so the more loops, the larger the memory footprint;
• Computation granularity is too coarse: the amount of computation can only be increased or decreased in whole blocks, with no fine-grained adjustment;
• Loop count is hard to change after training: once the model is trained, adjusting the loop count is not easy.

Together, these three limitations leave loop Transformers inflexible at deployment—and deployment is precisely where computation and memory most need to be adjusted to fit a budget.

KV cache (Key-Value Cache) is the core memory structure during autoregressive language model inference. When generating each token, the model needs the key (Key) and value (Value) vectors of all previous tokens to compute attention; to avoid recomputation, these vectors are cached. In a standard Transformer, the KV cache size scales linearly with sequence length; in a loop Transformer, if each loop round maintains its own cache, the cache is further multiplied by the loop count—which is exactly the root cause of "memory grows with the loop count." LooM's solution exploits the fact that loop layers share the same set of attention parameters: the KV entries for the same token are identical across loops, so they only need to be written once and can be reused directly in subsequent loops. Understanding this is key to seeing why "loop-invariant memory" can increase effective depth while keeping the cache at the level of a single forward pass.

LooM's Core Idea: Write Memory Only Once

LooM (Loop-invariant Memory) takes a divide-and-conquer approach. In the causal encoder-decoder architecture, it loops only over the decoder layers after the first decoder layer. These layers read from a global memory, and this global memory is written only once per token.

The key lies in KV cache handling: each loop layer publishes a token's key-value (KV) entries only on its first access to that token, so the cache does not grow with the loop count. This is exactly where the name "loop-invariant memory" comes from—the loop count changes, but the memory footprint does not.

This design directly resolves the first limitation. Compared with repeatedly executing the entire model, looping only over decoder layers with memory writes decoupled from the loop count is a lighter-weight path.

Adjustable Budget: From 3 to 21 Layer Executions

LooM's second design is to let these loop layers run independently. This way, each token can get any budget from 3 to 21 layer executions, and computation no longer jumps in whole blocks but can be adjusted continuously.

More practically, a model trained on randomly sampled budgets can cover all the above budgets without an extra router for dynamic selection. This means computation can be chosen on demand at deployment, without training a separate model for each budget. For scenarios that need to switch between different hardware or latency requirements, this flexibility has real value.

Results and the Impact of Training Budgets

The conclusion presented by the research is: when the loop count is sufficient, LooM outperforms a single-forward-pass baseline trained on the same data in language modeling, zero-shot accuracy, and long-context perplexity, while using the same cache size.

The paper's analysis also reveals a noteworthy detail: the budget distribution during training determines what benefits additional depth brings. In other words, more loops are not always better; how budgets are sampled during training affects the type of gains the model obtains from additional depth. This provides an adjustable dimension for future research.

Why CED Is a Natural Place for Looping

The paper's conclusion is an architectural judgment: the causal encoder-decoder is a natural place for looping—memory is written only once, and depth becomes a budget chosen at deployment.

The value of this view lies in turning "depth" from a fixed structure at training time into an adjustable resource at inference time. Once a traditional model is trained, its depth is fixed; LooM makes depth a variable that can be selected according to budget at deployment, without sacrificing cache efficiency. For applications seeking elasticity in inference cost, this is a directional hint.

Of course, this approach presupposes the causal encoder-decoder architecture, with looping applied only to decoder layers. Its applicability to other architectures, as well as specific design strategies for the training budget distribution, still await further verification.

Causal Encoder-Decoder (CED) is an architecture in which both the encoder and decoder are designed to be causal (i.e., each position can only see the left context), distinct from classic encoder-decoders like T5 where the encoder is bidirectionally visible. In CED, the encoder processes the input and produces representations, while the decoder generates output token by token in an autoregressive manner. This structure is particularly friendly to looping: the encoder runs only once and needs no looping, looping occurs only inside the decoder, and the decoder naturally requires a KV cache, so the constraint of "writing memory only once per token" can hold strictly. By contrast, looping all layers in a pure decoder model forces the cache to be reprocessed every round, and in hybrid encoder-decoder models the loop boundary is hard to define. CED thus becomes the architectural premise of LooM's design, and also explains why the paper emphasizes that the approach's applicability is constrained by architecture.

Key Points Recap

• LooM loops only over the decoder layers of CED; global memory is written only once per token, and the KV cache does not grow with the loop count;
• Loop layers run independently, supporting a continuous budget of 3 to 21 layer executions per token; a single model covers all budgets without a router;
• With a sufficient loop count, language modeling, zero-shot accuracy, and long-context perplexity all outperform a single-forward-pass baseline on the same data, with the same cache size;
• The training budget distribution determines the type of gains from additional depth, making depth a budget selectable at deployment.

One-line Summary

LooM loops only over CED decoder layers and writes memory only once per token, so memory does not grow with loops and the computation budget can be adjusted continuously.

Paragraph Summary

Loop Transformers increase effective depth by repeatedly executing shared layers, but suffer from three limitations: memory grows with the loop count, computation can only be adjusted in whole blocks, and the loop count is hard to change after training. LooM proposes a targeted solution: loop only over the decoder layers of a causal encoder-decoder (CED), and write global memory only once per token, so the KV cache does not grow with the loop count. Loop layers run independently, supporting a continuous budget of 3 to 21 layer executions per token, with a single model covering all budgets without a router. Experiments show that with a sufficient loop count, LooM outperforms a single-forward-pass baseline trained on the same data in language modeling, zero-shot accuracy, and long-context perplexity, with the same cache size. The paper also points out that the training budget distribution determines the type of gains from additional depth, thereby turning depth from a fixed structure at training time into a budget selectable at deployment.