Recall with Reasoning: Chain-of-Thought Distillation for Mamba’s Long-Context Memory and Extrapolation

要約

マンバの理論的無限コンテキストの可能性は、シーケンスがトレーニングの長さをはるかに上回る場合、実際には制限されています。
この作業では、教師モデルからのチェーン（COT）の要約を蒸留することにより、推論（RWR）でリコールする単純な効果的な方法で、マンバの長いコンテストメモリ能力のロックを解除します。
具体的には、RWRは、これらの要約を微調整中にCOTとしてプロンプトとして準備し、マンバに長いコンテキストで積極的に思い出し、推論するように教えます。
Longmemeval and Helmetの実験では、RWRがRWRを上回っており、Mambaの長いコンテキストパフォーマンスは、同様の前脱条件下での同等の変圧器/ハイブリッドベースラインに対するパフォーマンスを向上させ、すべてが建築的な変更なしで短いコンテキスト機能を維持します。

要約(オリジナル)

Mamba’s theoretical infinite-context potential is limited in practice when sequences far exceed training lengths. This work explores unlocking Mamba’s long-context memory ability by a simple-yet-effective method, Recall with Reasoning (RwR), by distilling chain-of-thought (CoT) summarization from a teacher model. Specifically, RwR prepends these summarization as CoT prompts during fine-tuning, teaching Mamba to actively recall and reason over long contexts. Experiments on LONGMEMEVAL and HELMET show RwR boosts Mamba’s long-context performance against comparable Transformer/hybrid baselines under similar pretraining conditions, while preserving short-context capabilities, all without architectural changes.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google