Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning

• URL: http://arxiv.org/abs/2505.16950v3
• Date: Fri, 26 Sep 2025 14:35:04 GMT
• Title: Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning
• Authors: Adnan Oomerjee, Zafeirios Fountas, Haitham Bou-Ammar, Jun Wang,
• Abstract summary: Transformer LLMs have been shown to exhibit strong reasoning ability that scales with inference-time compute.<n>We give a theoretical justification as to why memory (re)consolidation via KV cache rewrites is beneficial for improved reasoning.<n>Our model sees consistent performance gains over vanilla Transformers and pause-token augmented baselines, with gains of up to +6.6pp for selected tasks/backbones.
• Score: 16.35681450323654
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformer LLMs have been shown to exhibit strong reasoning ability that scales with inference-time compute, most prominently through token-space "thinking" chains of thought. A growing line of work pushes extra computation into the model's latent space, which we term Auxiliary Latent-Space Computation (ALSC). Existing ALSC methods largely fall into three buckets: (i) token-mediated latent rollouts, (ii) residual/activation steering, and (iii) memory (KV) compression. An underexplored alternative is memory consolidation/reconsolidation, two processes in the brain that are responsible for stabilising newly formed memory traces, and, upon recall, transiently rendering established traces plastic such they can integrate new contextual information before restabilising. In Transformer LLMs, this can be seen as analogous to performing in-place rewrites of new KV segments, and rewrites of recalled past segments. In this work, we give a theoretical justification as to why memory (re)consolidation via KV cache rewrites is beneficial for improved reasoning. We do this through the lens of Information Bottleneck (IB) theory, which posits that model generalisation emerges from an optimal balance between input information compression and retention of predictive information in latent representations. We then introduce the Bottlenecked Transformer, which augments a backbone LLM with a Cache Processor, an auxiliary Transformer that performs periodic, non-causal, in-place KV rewrites at newline-delimited reasoning step boundaries. The Processor consolidates recently written KV entries and reconsolidates a small, top-k attention-selected set of prior entries. We evaluate our Bottlenecked Transformer architecture on math reasoning benchmarks. Our model sees consistent performance gains over vanilla Transformers and pause-token augmented baselines, with gains of up to +6.6pp for selected tasks/backbones.

Related papers

• DeltaKV: Residual-Based KV Cache Compression via Long-Range Similarity [50.52392445266824]
  We propose a residual-based KV cache compression framework motivated by long-range inter-token similarity and highly shared latent components in KV representations.<n>Instead of discarding tokens, DeltaKV encodes semantic residuals relative to retrieved historical references, preserving fidelity while substantially reducing storage.<n>Experiments show that DeltaKV reduces KV cache memory to 29% of the original while maintaining near-lossless accuracy on LongBench, SCBench, and AIME.
  arXiv  Detail & Related papers  (2026-02-08T15:14:36Z)
• QuantSparse: Comprehensively Compressing Video Diffusion Transformer with Model Quantization and Attention Sparsification [67.15451442018258]
  Diffusion transformers exhibit remarkable video generation capability, yet their prohibitive computational and memory costs hinder practical deployment.<n>Model quantization and attention sparsification are two promising directions for compression, but each alone suffers severe performance degradation under aggressive compression.<n>We propose textbfQuantSparse, a unified framework that integrates model quantization with attention sparsification.
  arXiv  Detail & Related papers  (2025-09-28T06:49:44Z)
• OjaKV: Context-Aware Online Low-Rank KV Cache Compression with Oja's Rule [54.37983890753086]
  We introduce OjaKV, a framework that integrates a strategic hybrid storage policy with online subspace adaptation.<n>OjaKV preserves crucial first and most recent tokens in full-rank, maintaining high-fidelity anchors for attention.<n>It applies low-rank compression by incrementally adapting the projection basis using Oja's algorithm for online principal component analysis.
  arXiv  Detail & Related papers  (2025-09-25T21:42:27Z)
• KV-Latent: Dimensional-level KV Cache Reduction with Frequency-aware Rotary Positional Embedding [72.12756830560217]
  Large language models (LLMs) based on Transformer Decoders have become the preferred choice for conversational generative AI.<n>Despite the overall superiority of the Decoder architecture, the gradually increasing Key-Value cache during inference has emerged as a primary efficiency bottleneck.<n>By down-sampling the Key-Value vector dimensions into a latent space, we can significantly reduce the KV Cache footprint and improve inference speed.
  arXiv  Detail & Related papers  (2025-07-15T12:52:12Z)
• Learning Modular Exponentiation with Transformers [0.0]
  We train a 4-layer encoder-decoder Transformer model to perform modular exponentiation.<n>We find that reciprocal training leads to strong performance gains, with sudden generalization across related moduli.<n>These results suggest that transformer models learn modular arithmetic through specialized computational circuits.
  arXiv  Detail & Related papers  (2025-06-30T10:00:44Z)
• Unified Scaling Laws for Compressed Representations [69.72517034565467]
  We investigate whether a unified scaling framework can accurately predict model performance when training occurs over various compressed representations.<n>Our main finding is demonstrating both theoretically and empirically that there exists a simple "capacity" metric.<n>We extend our formulation to directly compare the accuracy potential of different compressed formats, and to derive better algorithms for training over sparse-quantized formats.
  arXiv  Detail & Related papers  (2025-06-02T16:52:51Z)
• Transformers Are Universally Consistent [14.904264782690639]
  We show that Transformers equipped with softmax-based nonlinear attention are uniformly consistent when tasked with executing Least Squares regression.<n>We derive upper bounds on the empirical error which, in the regime, decay at a provable rate of $mathcalO(t-1/2d)$, where $t$ denotes the number of input tokens and $d$ the embedding dimensionality.
  arXiv  Detail & Related papers  (2025-05-30T12:39:26Z)
• PiT: Progressive Diffusion Transformer [50.46345527963736]
  We propose a series of Pseudo textbfProgressive Dtextbfiffusion textbfTransformer (textbfPiT)<n>Our proposed PiT-L achieves 54%$uparrow$ FID improvement over DiT-XL/2 while using less computation.
  arXiv  Detail & Related papers  (2025-05-19T15:02:33Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• LoRC: Low-Rank Compression for LLMs KV Cache with a Progressive Compression Strategy [59.1298692559785]
  Key-Value ( KV) cache is crucial component in serving transformer-based autoregressive large language models (LLMs) Existing approaches to mitigate this issue include: (1) efficient attention variants integrated in upcycling stages; (2) KV cache compression at test time; and (3) KV cache compression at test time. We propose a low-rank approximation of KV weight matrices, allowing plug-in integration with existing transformer-based LLMs without model retraining. Our method is designed to function without model tuning in upcycling stages or task-specific profiling in test stages.
  arXiv  Detail & Related papers  (2024-10-04T03:10:53Z)
• Block Transformer: Global-to-Local Language Modeling for Fast Inference [41.361135177431336]
  Self-attention requires the key-value cache of all previous sequences to be retrieved from memory at every decoding step to retrieve context information. We introduce the Block Transformer which adopts hierarchical global-to-local modeling to autoregressive transformers to mitigate the inference bottlenecks associated with self-attention. We demonstrate that Block Transformers reach 10--20x inference throughput compared to vanilla transformers with equivalent perplexity and zero-shot task performance.
  arXiv  Detail & Related papers  (2024-06-04T17:45:26Z)
• Beyond Scaling Laws: Understanding Transformer Performance with Associative Memory [11.3128832831327]
  Increasing the size of a Transformer does not always lead to enhanced performance.<n>We present a theoretical framework that sheds light on the memorization during pre-training of transformer-based language models.
  arXiv  Detail & Related papers  (2024-05-14T15:48:36Z)
• Folded Context Condensation in Path Integral Formalism for Infinite Context Transformers [0.0]
  We present a generalized formulation of the Transformer algorithm by reinterpreting its core mechanisms within the framework of Path Integral formalism.<n>We obtain a more compact and efficient representation, in which the contextual information of a sequence is condensed into memory-like segments.<n>We validate the effectiveness of this approach through the Passkey retrieval task and a summarization task, demonstrating that the proposed method preserves historical information while exhibiting memory usage that scales linearly with sequence length.
  arXiv  Detail & Related papers  (2024-05-07T19:05:26Z)
• A Survey on Transformer Compression [84.18094368700379]
  Transformer plays a vital role in the realms of natural language processing (NLP) and computer vision (CV) Model compression methods reduce the memory and computational cost of Transformer. This survey provides a comprehensive review of recent compression methods, with a specific focus on their application to Transformer-based models.
  arXiv  Detail & Related papers  (2024-02-05T12:16:28Z)
• On the Convergence of Encoder-only Shallow Transformers [62.639819460956176]
  We build the global convergence theory of encoder-only shallow Transformers under a realistic setting. Our results can pave the way for a better understanding of modern Transformers, particularly on training dynamics.
  arXiv  Detail & Related papers  (2023-11-02T20:03:05Z)
• Context Compression for Auto-regressive Transformers with Sentinel Tokens [37.07722536907739]
  We propose a plug-and-play approach that is able to incrementally compress the intermediate activation of a specified span of tokens into compact ones. Experiments on both in-domain language modeling and zero-shot open-ended document generation demonstrate the advantage of our approach.
  arXiv  Detail & Related papers  (2023-10-12T09:18:19Z)
• RRWKV: Capturing Long-range Dependencies in RWKV [0.0]
  The paper devises the Retrospected Receptance Weighted Key Value architecture via incorporating the retrospecting ability into the RWKV to effectively absorb information. RWKV has exploited a linearly tensor-product attention mechanism and achieved parallelized computations by deploying the time-sequential mode.
  arXiv  Detail & Related papers  (2023-06-08T13:17:06Z)
• DIFFormer: Scalable (Graph) Transformers Induced by Energy Constrained Diffusion [66.21290235237808]
  We introduce an energy constrained diffusion model which encodes a batch of instances from a dataset into evolutionary states. We provide rigorous theory that implies closed-form optimal estimates for the pairwise diffusion strength among arbitrary instance pairs. Experiments highlight the wide applicability of our model as a general-purpose encoder backbone with superior performance in various tasks.
  arXiv  Detail & Related papers  (2023-01-23T15:18:54Z)
• Cost Aggregation with 4D Convolutional Swin Transformer for Few-Shot Segmentation [58.4650849317274]
  Volumetric Aggregation with Transformers (VAT) is a cost aggregation network for few-shot segmentation. VAT attains state-of-the-art performance for semantic correspondence as well, where cost aggregation also plays a central role.
  arXiv  Detail & Related papers  (2022-07-22T04:10:30Z)
• Resource-Efficient Separation Transformer [14.666016177212837]
  This paper explores Transformer-based speech separation with a reduced computational cost. Our main contribution is the development of the Resource-Efficient Separation Transformer (RE-SepFormer), a self-attention-based architecture. The RE-SepFormer reaches a competitive performance on the popular WSJ0-2Mix and WHAM! datasets in both causal and non-causal settings.
  arXiv  Detail & Related papers  (2022-06-19T23:37:24Z)
• Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers [42.93754828584075]
  We present a new Transformer architecture, Performer, based on Fast Attention Via Orthogonal Random features (FAVOR) Our mechanism scales linearly rather than quadratically in the number of tokens in the sequence, is characterized by sub-quadratic space complexity and does not incorporate any sparsity pattern priors. It provides strong theoretical guarantees: unbiased estimation of the attention matrix and uniform convergence.
  arXiv  Detail & Related papers  (2020-06-05T17:09:16Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.