技術要約: RADLADS

問題提起

線形アテンションのバリアント（RWKVやMambaなど）は、従来のソフトマックス・アテンション・トランスフォーマーと比較して、顕著な利点（具体的には、KVキャッシュを回避することによる推論時間のトークンあたりの$O(1)$化と、メモリ使用量の一定化）を提供しているが、大規模な線形モデルをゼロから学習させることは極めて高コストである。最先端（SoTA）のトランスフォーマー・モデルは、10兆トークン以上の学習を必要とすることが多く、そのコストは多くの組織にとって手の届かないものである。既存の、事前学習済みのソフトマックス・トランスフォーマーを線形または回帰モデルに変換する手法（T2R、SUPRA、LoLCats、MOHAWKなど）は、歴史的に膨大なトークン数（200億〜1000億以上）を必要としたり、MMLUのようなベンチマークにおいて低いダウンストリーム性能に陥ったりしてきた。さらに、多くの変換の試みはハイブリッド・アーキテクチャ（一部のソフトマックス・アテンションを保持する形式）に依存しているか、あるいは元のティーチャー・モデルの品質に到達できていない。

手法

著者らは、最小限の計算コストでソフトマックス・アテンション・トランスフォーマーを線形アテンション・デコーダー・モデルに変換するために設計されたプロトコルであるRADLADS（Rapid Attention Distillation to Linear Attention Decoders at Scale）を提示する。このプロセスは、3つの主要な段階と特定のアーキテクチャの適応を含んでいる。

1. RADLADS プロトコル

変換は、わずか3.5億〜7億トークン（ティーチャーの事前学習データの0.005%未満）のみを必要とする3ステップのプロセスである：

• セットアップ（アテンション重みの転送）： ティーチャー・モデルのアテンション関連の重み（$W_q, W_k, W_v, W_o$）を、学生モデルのシーケンス混合層に直接転送する。その他の重みは、標準的な事前学習手法によって初期化されるか、即座に効果を発揮することなくティーチャーの挙動を模倣するように設定される。
• ステップ1：アテンション隠れ状態のアライメント： 学生モデルのシーケンス混合層が、対応するティーチャーのアテンション層の隠れ状態出力を近似するように学習される。これは、L2距離（またはMSE）目的関数を用いて行われる。著者らは、Gated Linear Attentionカーネルを使用することで（オフバイワンの減衰やボーナス項を除去）、標準的なRWKV-6よりもティーチャーの隠れ状態により密接に適合できることを見出した。
• ステップ2：知識蒸留： 学生モデル全体が、KLダイバージェンス損失を用いてティーチャーの出力ロジットを近似するように学習される。著者らは、事実的な知識は主にティーチャーのMLPとエンベディングに存在すると仮定している。そのため、破滅的忘却を防ぐために、これらのコンポーネントの学習率は低く保たれるか固定される一方で、シーケンス・ミキサーはより積極的に学習される。
• ステップ3：コンテキスト長の拡張（オプション）： 長いコンテキスト能力を高めるために、標準的なクロスエントロピー損失を用いて長いシーケンス（最大16kトークン）でファインチューニングを行う。あるいは、著者らは、蒸留フェーズ自体の中でシーケンス長を4096に拡張するステップ2aを提案しており、これにより別途ステップ3を設ける必要がなくなる可能性がある。

2. 新しいアーキテクチャ

著者らは、標準的なRWKVアーキテクチャが必ずしも変換に最適ではないことを特定した。彼らは2つの新しいバリアントを導入した：

• RAD-RWKV6 ("RADFinch")： 適合性と安定性を向上させるために、Gated Linear Attentionカーネルと状態バランシング技術を用いたRWKV-6の修正版。
• RAD-RWKV7 ("RADGoose")： （この文脈では恩恵をもたらさなかった）「tokenshift」メカニ면ズムを取り除き、Rotary Positional Embeddings (RoPE) を直接適用したRWKV-7の修正版。このアーキテクチャは、未修正のRWKV-7と比較して、より速い収束と低い蒸留損失を示した。

3. ハイパーパラメータとデータ

• データセット： 著者らは、Qwenモデルの変換において、FineWebやFineWeb-EduよりもDCLM（DataComp-LM）が優れていることを見出し、すべての変換ステップでこれを選択した。
• 学習率： ステップ1では、隠れ状態をアライメントするために高い値（$10^{-3}$）から始まり、ティーチャーの最終的な事前学習学習率（$10^{-5}$）付近で終わるコサイン・アニーリング・スケジュールが使用される。ステップ2とステップ3では、一定の学習率が使用される。

主な貢献

1. RADLADS 蒸留レシピ： 最小限のデータで高品質な変換を可能にする、具体的なハイパーパラメータ、トークン数、データセットの選択を含む詳細なステップ・バイ・ステップのプロトコル。
2. 新しいアーキテクチャ： 変換プロセスに最適化され、未修正の前身モデルよりも高速な推論と優れたティーチャー・モデルへの適合を実現する、RAD-RWKV6およびRAD-RWKV7の導入。
3. 大規模モデル： 人気のあるオープンソースのQwen2.5モデルを、7B、32B、72Bのパラメータスケールで、線形アテンション・バリアントへと変換することに成功。
4. オープンソース・リリース： 他者がプロセスを再現できるように、コードと変換済みモデル（QRWKV6/7-7B/32B/72B）をApache 2.0ライセンス（72BモデルについてはQwenライセンスの制限あり）の下でリリース。

結果

変換されたモデルは、同規模の純粋な回帰モデルの中で最高水準の性能を達成している：

• 効率性： 72Bモデルの変換コストは2,000米ドル未満であり、約7億トークンしか必要としない。
• 性能： 標準的なベンチマーク（Lambada、MMLU、ARCなど）において、RADLADSモデルは他の変換手法（SUPRA、LoLCats、MOHAWK、ARWKVなど）を一貫して上回っている。
  • QRWKV7-7B-Instructは、ティーチャーに対して**92.4%**の相対的MMLUスコアを達成し、他の7B変換モデルを大幅に凌駕している。
  • QRWKV6-72B-Instructは、相対的MMLUスコア**89.9%**を達成し、このスケールにおける純粋なRNN言語モデルの新たなSoTAを確立した。
• 推論速度： 線形アテンション機構により、変換されたモデルはコンテキスト長が増加するにつれてティーチャー・モデルに対して大幅な高速化を示す。例えば、入力8kトークン、出力256トークンの場合、32BのQRWKV7モデルはQwen3-32Bティーチャーよりも1.61倍高速であり、入力6k、出力2kの場合、高速化は3.41倍に達する。

限界と主張

論文では、いくつかの限界についても謙虚に認めている：

• 推論とロングコンテキスト： 標準的なベンチマークでは強力な性能を示す一方で、複雑な推論タスク（Minerva Mathなど）や非常に長いコンテキストタスク（RULERなど）においては、ティーチャー・モデルのように追加の出力トークンに伴って効果的に改善されないという限界がある。
• アーキテクチャの感度： 各新しいアーキテクチャ設計には、RADLADSプロトコルとの互換性を確保するための綿密なテストが必要である。例えば、一部の線形アテンション・バリアントで一般的なGroupNorm/LayerNormは、14B以上のスケールで学習の不安定性を引き起こしたため、プリスケーリング手法に置き換えられた。
• データセットのアライメント： 著者らは、反復的なループ挙動を防ぐために、推論データセットをティーチャー・モデルの分布により密接に合わせる必要があると指摘しており、これはDCLMとOpenThoughtsを混合することで部分的に対処した。

意義

本論文は、RADLADSが大規模な線形アテンション・モデルへのアクセスを民主化するためのコスト効率の高い経路を提供すると主張している。学習に必要なトークン数を兆単位から数億単位へと削減することで、研究者や小規模な組織が、事前学習に伴う極端なコストを負うことなく、新しい種類のRNNアーキテクチャを大規模にテスト、訓練、展開することを可能にする。著者らは、これを次世代の圧縮型状態アテンション・バリアントの開発を加速させるためのツールとして位置づけている。