FlashLoop: Fast and Memory-Efficient Looped Transformers via Lazy Updates

• Looped Transformers use repeated shared blocks to add computational depth without increasing the stored model-weight footprint.This provides a way to increase test-time compute while keeping model weights fixed.Source: §1 · p. 1
• Looped models apply the shared block repeatedly, increasing computation and KV-cache memory linearly with loop depth; at 32K, Ouro-2.6B’s prefill takes 27 seconds and its cache uses 48 GiB.The passage compares these costs with 3 seconds and 4 GiB for LLaMA-3.1-8B.Source: §1 · p. 1
• As recurrence proceeds, changes become increasingly sparse, predictable, and compressible across token states, attention outputs, and adjacent-loop KV residuals.Most token rows change little, attention differences concentrate on predictable columns, and residuals reconstruct with lower quantization error than full KV states.Source: §1 · p. 2
• FlashLoop is a training-free, loop-native framework combining token-sparse updates, sparse attention, and KV-residual quantization.It reuses converged-token states, selects important keys using preceding-loop attention statistics, and stores a quantized base with low-bit residual updates.Source: §1 · p. 2
• Up to 1.64× end-to-end speedup and 6× lower KV-cache memory, with largely preserved accuracy, bring inference costs to comparable or lower levels than performance-matched non-looped models.The comparison is illustrated at an 8K context length.Source: §1 · p. 2, Figure 1 · p. 2

• Related recurrent models reuse experts or shared modules, or incorporate latent reasoning through weight-shared recurrence, while Universal Transformer introduced depth recurrence and adaptive per-position computation.Source: §2 · p. 2
• Prior loop-specific cache methods use learned gated updates or trained encoders and decoders, whereas FlashLoop exploits recurrence redundancy at inference without modifying pretrained model weights or requiring training.Source: §2 · p. 3
• General inference methods address quantization, cross-layer KV reuse, token eviction, and sparse attention, but primarily target non-looped Transformers rather than recurrence-axis redundancy.FlashLoop’s loop-specific optimizations are described as complementary to these methods.Source: §2 · p. 3

3 CROSS-LOOP REDUNDANCY IN LOOPED TRANSFORMERS

• The original execution applies the shared Transformer block to all tokens at every loop and stores a separate key-value pair for each loop, so computation and cache grow linearly with loop count.The section frames these full updates as overlooking the increasingly structured changes between successive loops.Source: §3 · pp. 3–4

3.1 TOKEN-UPDATE REDUNDANCY

• Tokens with substantial hidden-state changes become sparser across loops, remain distributed through the context, and are largely nested within the active tokens of the preceding transition.Because keys and values are computed from token representations, converged tokens can reuse preceding-loop KV states.Source: §3.1 · p. 4

3.2 ATTENTION-COMPUTATION REDUNDANCY

• Attention-output changes concentrate on a sparse subset of key columns whose importance remains stable across adjacent loops, allowing the preceding loop to identify likely important columns.The section describes approximating the current-loop output by selectively updating columns and reusing the rest.Source: §3.2 · p. 4
• In Ouro-1.4B at loop 3/4, updating only 10% of attention columns reconstructs over 90% of exact attention output.Selected columns are recomputed from current-loop states while the remaining contributions are reused from the preceding loop.Source: §3.2 · p. 5

3.3 QUANTIZATION-FRIENDLY KV RESIDUAL

• Adjacent-loop KV residuals shrink with loop depth and yield lower relative reconstruction error under quantization than full KV states, motivating a quantized base with residual updates.In long contexts, caching active-token KV states remains a major memory overhead.Source: §3.3 · p. 5

4 FLASHLOOP

• FlashLoop combines token-sparse updates, sparse attention, and KV-residual quantization to reduce structural redundancy in pretrained looped Transformers without training.Source: §4 · p. 5

4.1 CROSS-LOOP TOKEN-SPARSE UPDATES

• Later loops update only the top-k tokens retained from the preceding loop, while inactive tokens reuse hidden and KV states.Nested active-token sets reduce computation and avoid storing a complete KV cache for each loop.Source: §4.1 · p. 5

4.2 LOOP-AWARE SPARSE ATTENTION

• Sparse attention recomputes top-K key columns ranked using the preceding loop, then rescales selected-column probabilities by their cached full-attention mass.Later loops can prune the selected set further using statistics from the preceding sparse loop.Source: §4.2 · pp. 5–6

4.3 CROSS-LOOP KV RESIDUAL QUANTIZATION

• KV-residual quantization stores a quantized first-loop base and quantized residual updates relative to each preceding reconstructed loop.Inactive tokens reuse their K/V states, so residuals are generated and stored only for active tokens, reducing the number and cost of updates.Source: §4.3 · p. 6

4.4 FAST GPU SUPPORT

• GPU support gathers selected tokens and keys into compact buffers and fuses cross-loop KV updates with quantization in the attention kernel.These choices address irregular memory access, fragmented computation, data movement, and kernel-launch overhead.Source: §4.4 · p. 6

5.1 EXPERIMENTAL SETUP

• The evaluation covers five Looped Transformers and five benchmarks, with task accuracy and KV-memory reduction reported relative to the original models.Source: §5.1 · p. 6, Table 1 · p. 7

5.2 MAIN RESULTS

• Across five benchmarks, FlashLoop largely preserves task performance while substantially reducing inference costs; base models tend to be more robust than thinking variants.Source: Table 1 · p. 7, §5.2 · p. 7

5.3.1 CONTRIBUTION OF EACH COMPONENT

• Component ablations find that sparse attention slightly improves average performance, token-sparse updates largely preserve that gain, and full FlashLoop remains comparable to the original; full-state KV quantization degrades accuracy more than residual quantization.Source: §5.3.1 · p. 7, Table 2 · p. 7
• Token-sparse updates reduce prefill computation and cache storage, sparse attention reduces decode costs without added overhead, and 4-bit residual quantization further reduces cache storage and decode latency.Source: §5.3.1 · pp. 7–8, Figure 4 · p. 8

5.3.2 COMPARISON WITH ALTERNATIVE KV-CACHE STRATEGIES

• Compared with H2O and Last-step KV Reuse, FlashLoop maintains performance closer to original models by preserving useful cross-loop information while selectively compressing redundant updates.Source: Table 3 · p. 7, §5.3.2 · p. 8

5.3.3 SPARSITY AND QUANTIZATION CONFIGURATIONS

• Accuracy is robust across most attention-column settings; retaining 25% of tokens in loop 3 and 10% in loop 4 preserves comparable accuracy, while 4-bit quantization has negligible impact and 2-bit degrades performance.Source: §5.3.3 · p. 8, Figure 5 · p. 9

5.4 END-TO-END EFFICIENCY ANALYSIS

• At 32K context, FlashLoop reduces KV-cache storage by 20 GiB (82.8%), achieves approximately 1.6× end-to-end speedup, and reduces decode latency by about 37.5%.Source: §5.4 · pp. 8–9, Figure 6 · p. 9
• As loop count rises, baseline prefill FLOPs and KV-cache memory grow approximately linearly, while FlashLoop slows their growth in later loops, especially for Huginn-3.5B with 32 loops.Source: §5.4 · p. 9, Figure 7 · p. 9

6 CONCLUSION

• FlashLoop uses token-sparse updates, loop-aware sparse attention, and cross-loop KV residual quantization to exploit three forms of redundancy in looped Transformers.Source: §6 · p. 10

AI USE STATEMENT

• Generative AI assisted with readability and minor code editing, but not with research hypotheses, methodology, results interpretation, or main scientific claims.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT

• The work studies inference efficiency using publicly available models and benchmarks, but does not address inherited biases or harmful-generation risks.The statement calls for appropriate safety safeguards in deployment.Source: ETHICS STATEMENT · p. 10

A PERFORMANCE-MATCHED LOOPED AND NON-LOOPED MODELS

• Ouro-2.6B with four loops achieves comparable accuracy on several benchmarks, while FlashLoop addresses its inference-efficiency challenge without training.The comparison includes performance-matched non-looped models.Source: A PERFORMANCE-MATCHED LOOPED AND… · p. 13

B CROSS-LOOP REDUNDANCY IN MODEL WITH MORE LOOPS

• In Huginn-3.5B with 32 recurrent loops, important attention columns stabilize in later loops as adjacent-loop KV residuals shrink.The authors report similar cross-loop redundancy and attention concentration in this deeper-loop model.Source: B CROSS-LOOP REDUNDANCY IN MODEL WITH… · p. 13

C ATTENTION PATTERNS ACROSS LOOPS

• Ouro-1.4B prefill attention maps show recurring, stable key-column patterns across loops, motivating selection from the preceding loop’s attention probabilities.The figure visualizes attention probabilities across four loops.Source: C ATTENTION PATTERNS ACROSS LOOPS · p. 13, Figure 9 · p. 14

D MATHEMATICAL FORMULATION OF FLASHLOOP

• The formulation uses H, K, and V for hidden states, keys, and values; r indexes loops, while i and j index token positions.The Transformer block is shared across recurrence, and the original implementation applies it to all tokens at every loop.Source: D MATHEMATICAL FORMULATION OF FLASHLOOP · p. 13

D.1 CROSS-LOOP TOKEN-SPARSE UPDATES

• After dense early loops, later loops recompute only nested active-token sets, reusing hidden states and KV cache for the remaining tokens.A(r) is the active-token set and k(r) its budget.Source: §D.1 · p. 14

D.2 LOOP-AWARE SPARSE ATTENTION

• Loop-aware sparse attention selects and caches key indices and their global mass from the preceding loop, then reuses both for current-loop attention.Updated probabilities prepare the nested selection and cached mass for the next loop.Source: §D.2 · pp. 14–15

D.3 CROSS-LOOP KV RESIDUAL QUANTIZATION

• KV residual quantization stores a quantized base and cross-loop residuals; only active-token residuals need new payload, while column sparsity limits keys and values read.Token sparsity selects updates, quantization compresses them, and column sparsity selects reads.Source: §D.3 · p. 15

Sparsity Setting

• The same loop-wise sparsity configuration is used for each model across all benchmarks.Table 5 summarizes token and column retention settings.Source: E DETAILED EXPERIMENTAL SETTINGS · p. 15

Quantization Setting

• KV base states and residuals use INT4 asymmetric group-wise quantization with group size 64, while the most recent 64 tokens remain in BF16.Post-RoPE keys use per-channel quantization and values use per-token quantization.Source: E DETAILED EXPERIMENTAL SETTINGS · p. 15

Speedup and KV Cache Memory

• Speedup and KV-cache memory are measured using end-to-end wall-clock time and peak cache memory over the benchmark experiments.Source: E DETAILED EXPERIMENTAL SETTINGS · p. 15

Hyperparameter

• Each model’s smallest retention ratio is calibrated to keep hidden-state or attention-output relative reconstruction error below 5%, then fixed across downstream tasks.Calibration uses 128 randomly sampled WikiText-2 sequences.Source: E DETAILED EXPERIMENTAL SETTINGS · p. 15

F LONG-CONTEXT QUALITY EVALUATION

• Across the tested long-context settings, FlashLoop does not bring significant performance degradation for Ouro-1.4B.The comparison uses WikiText-2 at 8K, 16K, and 32K context lengths and a Needle-in-a-Haystack benchmark, with the main-experiment hyperparameters.Source: F LONG-CONTEXT QUALITY EVALUATION · p. 15
• Table 6 reports WikiText-2 perplexity on 128 continuation tokens and Needle-in-a-Haystack accuracy over full single-needle and multi-needles samples.Source: Table 6 · p. 16

G EFFICIENT GPU IMPLEMENTATION

• The optimized engine avoids direct PyTorch overheads from kernel launches, memory traffic, and intermediate tensors to realize runtime and memory savings.The paper also provides a PyTorch reference implementation.Source: G EFFICIENT GPU IMPLEMENTATION · p. 16
• GPU optimizations include active-token execution, fused projections, packed 4-bit KV storage, custom attention kernels, and fused sparse decode selection.Attention kernels read packed anchor and residual KV streams directly, avoiding full BF16 KV reconstruction.Source: G EFFICIENT GPU IMPLEMENTATION · p. 16
• The engine targets inference with BF16 activations on NVIDIA GPUs and is intended as a reference implementation of FlashLoop’s techniques.Source: G EFFICIENT GPU IMPLEMENTATION · p. 16