文章主要内容总结本文针对Transformer架构在捕获乔姆斯基层级（如正则表达式、确定性上下文无关文法）方面的固有局限，提出了一种名为STACKTRANS的新型模型。该模型受下推自动机（使用栈高效处理确定性上下文无关文法）启发，在Transformer层间引入隐藏状态栈，通过可微分的栈操作（推入、弹出、无操作）实现端到端训练，同时保持与现有框架（如flash-attention）的兼容性。STACKTRANS的核心设计包括：可微分栈状态更新与掩码维护、全局读取机制（通过可学习的栈上查询注意力解决仅读取栈顶元素导致的训练不稳定问题）、多头低秩栈（提升表示能力并降低计算成本）。实验表明，在形式语言任务（乔姆斯基层级相关）和大规模自然语言任务中，STACKTRANS均优于标准Transformer及其他基线模型。其中，360M参数的STACKTRANS模型性能超过了参数规模为其2-3倍的开源大语言模型，展现出更高的效率和推理能力。文章创新点架构创新：在Transformer层间引入隐藏状态栈，而非修改注意力计算，既保留原有框架兼容性，又赋予模型处理乔姆斯基层级文法的能力。可微分栈操作：设计软推入、弹出和无操作，使栈操作可微分，支持端到端训练，解决了传统离散栈操作阻碍梯度传播的问题。多头低秩与全局读取：采用多头栈机制提升表示灵活性，通过全局读取（栈上注意力加权求和）替代仅读取栈顶元素，增强训练稳定性和模型表达能力。性能与效率