State Tuning: State-based Test-Time Scaling on RWKV-7

要約

テスト時間スケーリングは、機械学習における顕著な研究方向として浮上しており、モデルが推論中の表現力のある能力を高めることを可能にします。トランスフォーマーは、効率と表現力の間の微妙なバランスをとることで有名です。
RNNベースのRWKV-7モデルに合わせて、RWKV-7のユニークな強みを活用することにより、この方法は、モデルの事前に訓練された重みを変更することなく、ターゲットタスクの最先端のパフォーマンスを実現します。
私たちのアプローチは、3つの重要なイノベーションに集中しています。
まず、小さなモデルがRWKV-7モデルの状態ダイナミクスを複製および学習できるようにするオブザーバーフレームワークを開発します。
第二に、カーネル法を使用して状態サイズを動的にアップスケールし、複雑なパターンをキャプチャするモデルの能力を高めます。
第三に、逆相関のバックプロパゲーション（DBP）を統合して、アップスケール状態マトリックスを最適化し、それにより収束と表現力を向上させます。
状態マトリックスのみを調整することにより、小さなモデルが指定されたタスク上のより大きなモデルを上回ることができることを示します。
この方法は、テスト時間スケーリングの力を活用して優れた結果を提供しながら、元のRWKV-7アーキテクチャの効率を維持します。
私たちの調査結果は、リソースに制約のある設定でモデルパフォーマンスを進めるための効果的な戦略として、状態調整の可能性を強調しています。
私たちのコードはhttps://github.com/torchrwkv/flash-linear-attentionです。

要約(オリジナル)

Test-time scaling has emerged as a prominent research direction in machine learning, enabling models to enhance their expressive capabilities during inference.Transformers, renowned for striking a delicate balance between efficiency and expressiveness, have benefited from test-time scaling techniques that leverage an expanding key-value (KV) cache to significantly improve performance.In this paper, we introduce a novel state-based approach to test-time scaling, which we term state tuning, tailored to the RNN-based RWKV-7 model.By exploiting the unique strengths of RWKV-7, our method achieves state-of-the-art performance on the target task without altering the model’s pre-trained weights. Our approach centers on three key innovations. First, we develop an observer framework that allows a smaller model to replicate and learn the state dynamics of the RWKV-7 model. Second, we employ a kernel method to dynamically upscale the state size, enhancing the model’s capacity to capture intricate patterns. Third, we integrate Decorrelated Backpropagation (DBP) to optimize the upscaled state matrix, thereby improving convergence and expressivity. By tuning only the state matrix, we demonstrate that a smaller model can outperform larger models on the given task. This method preserves the efficiency of the original RWKV-7 architecture while harnessing the power of test-time scaling to deliver superior results. Our findings underscore the potential of state tuning as an effective strategy for advancing model performance in resource-constrained settings. Our code is https://github.com/TorchRWKV/flash-linear-attention.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google