状态调优:RWKV-7 上基于状态的测试时缩放
计算与语言
2025-04-08 v1 机器学习
摘要
测试时缩放已成为机器学习领域的一个重要研究方向,使模型能够在推理过程中增强其表达能力。Transformer 因在效率和表达性之间取得了微妙的平衡而闻名,它受益于利用不断扩展的键值(KV)缓存以显著提升性能的测试时缩放技术。在本文中,我们引入了一种新颖的基于状态的测试时缩放方法,我们称之为状态调优,专为基于 RNN 的 RWKV-7 模型量身定制。通过利用 RWKV-7 的独特优势,我们的方法在不改变模型预训练权重的情况下,在目标任务上实现了 SOTA 性能。我们的方法围绕三个关键创新展开。首先,我们开发了一个观察者框架,允许较小的模型复制并学习 RWKV-7 模型的状态动态。其次,我们采用核方法动态放大状态尺寸,增强了模型捕捉复杂模式的能力。第三,我们集成了去相关反向传播(DBP)来优化放大的状态矩阵,从而改善了收敛性和表达性。通过仅调优状态矩阵,我们证明了较小的模型在给定任务上可以优于较大的模型。该方法在保持原始 RWKV-7 架构效率的同时,利用了测试时缩放的力量以提供卓越的结果。我们的发现强调了状态调优作为在资源受限环境中提升模型性能的有效策略的潜力。我们的代码见 https://github.com/TorchRWKV/flash-linear-attention。
引用
@article{arxiv.2504.05097,
title = {State Tuning: State-based Test-Time Scaling on RWKV-7},
author = {Liu Xiao and Li Zhiyuan and Lin Yueyu},
journal= {arXiv preprint arXiv:2504.05097},
year = {2025}
}