RePo: コンテキスト再配置による言語モデル

要旨

文脈内学習は現代の大規模言語モデル（LLM）の中核をなす技術である。しかし、従来のアーキテクチャでは、線形あるいは固定の位置インデックスを割り当てることで、硬直的で固定された文脈構造が課せられてきた。認知的負荷理論（CLT）に基づき、我々はこの情報量の乏しい構造が、余分な認知的負荷を増大させ、深い推論や注意配分に割り当てられるべき有限の作業記憶容量を消費していると論じる。この問題に対処するため、我々は文脈の再配置によって余分な負荷を軽減する新規メカニズム「RePo」を提案する。標準的な手法とは異なり、RePoは微分可能なモジュールf_φを利用して、事前定義された整数範囲に依存するのではなく、文脈的依存関係を捉えるトークン位置を割り当てる。OLMo-2 1Bバックボーンを用いた継続事前学習により、RePoが、雑音を含む文脈、構造化データ、長い文脈長を伴うタスクにおける性能を大幅に向上させながら、一般的な短文脈タスクでも競争力のある性能を維持することを実証する。詳細な分析により、RePoが、遠く離れた関連情報により高い注意を割り当て、密で非線形な空間内で位置を割り当て、入力文脈の内在的構造を捉えることに成功していることが明らかとなった。コードはhttps://github.com/SakanaAI/repo で公開している。

English

In-context learning is fundamental to modern Large Language Models (LLMs); however, prevailing architectures impose a rigid and fixed contextual structure by assigning linear or constant positional indices. Drawing on Cognitive Load Theory (CLT), we argue that this uninformative structure increases extraneous cognitive load, consuming finite working memory capacity that should be allocated to deep reasoning and attention allocation. To address this, we propose RePo, a novel mechanism that reduces extraneous load via context re-positioning. Unlike standard approaches, RePo utilizes a differentiable module, f_φ, to assign token positions that capture contextual dependencies, rather than replying on pre-defined integer range. By continually pre-training on the OLMo-2 1B backbone, we demonstrate that RePo significantly enhances performance on tasks involving noisy contexts, structured data, and longer context length, while maintaining competitive performance on general short-context tasks. Detailed analysis reveals that RePo successfully allocate higher attention to distant but relevant information, assign positions in dense and non-linear space, and capture the intrinsic structure of the input context. Our code is available at https://github.com/SakanaAI/repo.

RePo: コンテキスト再配置による言語モデル

RePo: Language Models with Context Re-Positioning

要旨

Support