用状态调优让小模型在推理时超越大模型,不改权重只调状态。
State Tuning: State-based Test-Time Scaling on RWKV-7
- 用小模型模仿大模型的状态动态,实现高效学习。
- 通过核方法动态扩大状态规模,捕捉更复杂模式。
- 适合资源受限场景,尤其适合想提升RWKV性能的研究者。
测试时扩展已成为机器学习的重要研究方向,使模型在推理阶段增强表达能力。虽然基于Transformer的模型通过扩展键值缓存显著提升性能,但本文针对基于RNN的RWKV-7模型,提出一种全新的基于状态的测试时扩展方法——状态调优。该方法不修改预训练权重,仅通过调节状态矩阵即可实现当前最优表现。核心创新包括:构建观察器框架,使小型模型可复现并学习大型模型的状态动态;采用核方法动态扩展状态尺寸以增强模式捕捉能力;结合去相关反向传播(DBP)优化扩展后的状态矩阵,提升收敛性与表达力。实验表明,仅调整状态矩阵即可使小模型超越更大模型。该方法在保持原始架构效率的同时,借助测试时扩展实现卓越性能,为资源受限场景下的模型增强提供了有效路径。代码已开源:https://github.com/TorchRWKV/flash-linear-attention。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。