強化学習による高速重み付けと次系列予測

Reinforced Fast Weights with Next-Sequence Prediction

February 18, 2026

著者: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky

cs.AI

要旨

高速重みアーキテクチャは、文脈長に依存しない一定のメモリオーバーヘッドを維持することで、長文脈モデリングにおける注意機構ベースのトランスフォーマーに代わる有望な代替手段を提供する。しかし、その潜在能力は次トークン予測（NTP）の訓練パラダイムによって制限されている。NTPは単一トークンの予測を最適化するが、プレフィックスに続く複数トークンにわたる意味的コヒーレンスを無視する。結果として、文脈情報を格納するためにパラメータを動的に更新する高速重みモデルは、長距離依存関係を捕捉できない最適ではない表現を学習してしまう。本論文では、強化学習フレームワークであるREFINE（Reinforced Fast weIghts with Next sEquence prediction）を提案する。REFINEは次シーケンス予測（NSP）目標の下で高速重みモデルを訓練し、予測エントロピーに基づいて情報量の多いトークン位置を選択し、複数トークンのロールアウトを生成し、自己監督によるシーケンスレベルの報酬を割り当て、グループ相対方策最適化（GRPO）を用いてモデルを最適化する。REFINEは事前訓練済み言語モデルの訓練ライフサイクル全体（訓練途中、訓練後、テスト時訓練）に適用可能である。LaCT-760MとDeltaNet-1.3Bを用いた実験により、REFINEが干し草の中の針探索、長文脈質問応答、LongBenchの多様なタスクにおいて、NTPを用いた教師ありファインチューニングを一貫して上回ることを実証する。REFINEは高速重みアーキテクチャにおける長文脈モデリングを改善する効果的かつ汎用的なフレームワークを提供する。

English

Fast weight architectures offer a promising alternative to attention-based transformers for long-context modeling by maintaining constant memory overhead regardless of context length. However, their potential is limited by the next-token prediction (NTP) training paradigm. NTP optimizes single-token predictions and ignores semantic coherence across multiple tokens following a prefix. Consequently, fast weight models, which dynamically update their parameters to store contextual information, learn suboptimal representations that fail to capture long-range dependencies. We introduce REFINE (Reinforced Fast weIghts with Next sEquence prediction), a reinforcement learning framework that trains fast weight models under the next-sequence prediction (NSP) objective. REFINE selects informative token positions based on prediction entropy, generates multi-token rollouts, assigns self-supervised sequence-level rewards, and optimizes the model with group relative policy optimization (GRPO). REFINE is applicable throughout the training lifecycle of pre-trained language models: mid-training, post-training, and test-time training. Our experiments on LaCT-760M and DeltaNet-1.3B demonstrate that REFINE consistently outperforms supervised fine-tuning with NTP across needle-in-a-haystack retrieval, long-context question answering, and diverse tasks in LongBench. REFINE provides an effective and versatile framework for improving long-context modeling in fast weight architectures.