2025_NIPS_Recursive Transformer: Boosting Reasoning Ability with State Stack

文章核心总结与翻译

一、主要内容

该研究针对Transformer架构在形式语言推理（如正则表达式、确定性上下文无关文法）中的固有局限性，受下推自动机的栈结构启发，提出了名为STACKTRANS的新型Transformer变体。核心思路是在Transformer层间引入可微分的隐藏状态栈，通过软推入（push）、弹出（pop）和无操作（no-op）等可学习操作，为模型注入乔姆斯基层级语法的归纳偏置，同时保持与flash-attention等现有框架的兼容性。

研究通过多维度实验验证模型效果：在形式语言任务（涵盖正则语言、确定性上下文无关语言、上下文敏感语言）中，STACKTRANS相比标准Transformer准确率提升至少30%，部分任务达到近100%；在自然语言任务（常识推理、问答、数学推理等）中，360M参数的STACKTRANS-360M模型（预训练于约1T tokens）性能优于参数规模2-3倍的开源LLM。此外，模型成功扩展至7B参数，展现出良好的扩展性和训练效率。

二、创新点

1. 层间可微分栈设计：在Transformer层间嵌入隐藏状态栈，而非修改注意力机制，既保持了原有架构兼容性，又通过可微分的软栈操作实现端到端训练。
2. 栈操作优化：设计多头部低秩栈机制提升表征能力，采用全局读取操作（而非仅读取栈顶元素）稳定训练过程，引入栈掩码和正则化避免操作退化。
3. 兼顾形式与自然语言推理：首次实现栈增强Transformer在乔姆斯基层级