用Mamba提升连续控制中的算法蒸馏，突破Transformer瓶颈。
Scaling Algorithm Distillation for Continuous Control with Mamba
- 用S6结构的Mamba替代Transformer，实现长序列高效建模。
- 在4个复杂连续环境中，Mamba比Transformer更优且支持超长上下文。
- 长上下文蒸馏使ICRL性能媲美顶尖在线元强化学习方法。
算法蒸馏（AD）是一种新方法，通过因果Transformer自回归地建模跨回合训练历史，实现上下文强化学习（ICRL）。但受注意力机制限制，现有实验受限于Transformer的二次复杂度，仅能在短时程、简单离散环境中进行。本文提出利用近期提出的有选择性结构状态空间序列（S6）模型，该模型在长序列建模中表现优异且计算复杂度线性增长。我们使用基于S6层构建的Mamba模型，在四个复杂连续元强化学习环境中验证其优势：相比Transformer，Mamba在算法蒸馏中整体表现更优。此外，将AD扩展至极长上下文后，可显著提升ICRL性能，甚至达到顶尖在线元强化学习基线水平。
原文摘要 · Abstract (English)
Algorithm Distillation (AD) was recently proposed as a new approach to perform In-Context Reinforcement Learning (ICRL) by modeling across-episodic training histories autoregressively with a causal transformer model. However, due to practical limitations induced by the attention mechanism, experiments were bottlenecked by the transformer's quadratic complexity and limited to simple discrete environments with short time horizons. In this work, we propose leveraging the recently proposed Selective Structured State Space Sequence (S6) models, which achieved state-of-the-art (SOTA) performance on long-range sequence modeling while scaling linearly in sequence length. Through four complex and continuous Meta Reinforcement Learning environments, we demonstrate the overall superiority of Mamba, a model built with S6 layers, over a transformer model for AD. Additionally, we show that scaling AD to very long contexts can improve ICRL performance and make it competitive even with a SOTA online meta RL baseline.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。