Bottlenecked Transformers: Periodic KV Cache Abstraction for Generalised Reasoning

要約

印象的な能力にもかかわらず、大規模な言語モデルは、トレーニングの分布を超えて一般化に苦労し、多くの場合、真の抽象的な推論（外挿）ではなく洗練されたパターン補間を示します。
この作業では、ボトルネック（IB）理論のレンズを通じてこの制限にアプローチします。これは、モデルの一般化が、入力圧縮と潜在的な表現における予測情報の保持との最適なバランスから生じると仮定しています。
IB理論を使用して、デコーダーのみの変圧器がタスクと最適なシーケンス表現を形成する能力に本質的に制約されていることを証明します。
次に、この結果を使用して、内部シーケンスレベルの表現（KVキャッシュ）の定期的なグローバル変換が、推論タスクにおけるトランスの一般化を改善するために必要な計算ステップであることを示します。
これらの理論的洞察に基づいて、定期的な間隔でKVキャッシュをグローバルに書き換える追加モジュールの形で、トランスアーキテクチャの変更を提案し、入力のプレフィックスの記憶からエンコード機能への容量をシフトし、将来のトケンを予測するのに最も役立ちます。
私たちのモデルは、数学的推論ベンチマークに大きな利益をもたらし、最大3.5倍のパラメーターとキャッシュ圧縮のためのヒューリスティック駆動型の剪定メカニズムで両方のバニラ変圧器よりも優れています。
私たちのアプローチは、既存のKVキャッシュ圧縮法の原則的な一般化と見なすことができます。
このような方法は、入力表現の圧縮のみに焦点を当てていますが、多くの場合、予測情報を保持することを犠牲にしてそうします。したがって、それらの能力は本質的に制約のないモデルの機能によって制限されます。
これにより、情報理論を使用してトランスメモリを操作するための原則的なフレームワークが確立され、スケーリングだけでは克服できない基本的な推論の制限に対処します。

要約(オリジナル)

Despite their impressive capabilities, Large Language Models struggle with generalisation beyond their training distribution, often exhibiting sophisticated pattern interpolation rather than true abstract reasoning (extrapolation). In this work, we approach this limitation through the lens of Information Bottleneck (IB) theory, which posits that model generalisation emerges from an optimal balance between input compression and retention of predictive information in latent representations. We prove using IB theory that decoder-only Transformers are inherently constrained in their ability to form task-optimal sequence representations. We then use this result to demonstrate that periodic global transformation of the internal sequence-level representations (KV cache) is a necessary computational step for improving Transformer generalisation in reasoning tasks. Based on these theoretical insights, we propose a modification to the Transformer architecture, in the form of an additional module that globally rewrites the KV cache at periodic intervals, shifting its capacity away from memorising input prefixes and toward encoding features most useful for predicting future tokens. Our model delivers substantial gains on mathematical reasoning benchmarks, outperforming both vanilla Transformers with up to 3.5x more parameters, as well as heuristic-driven pruning mechanisms for cache compression. Our approach can be seen as a principled generalisation of existing KV-cache compression methods; whereas such methods focus solely on compressing input representations, they often do so at the expense of retaining predictive information, and thus their capabilities are inherently bounded by those of an unconstrained model. This establishes a principled framework to manipulate Transformer memory using information theory, addressing fundamental reasoning limitations that scaling alone cannot overcome.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google