arXiv:2507.15343cs.SEcs.AI2025-07

让大模型学会用栈结构理解语法,提升逻辑推理能力。

StackTrans: From Large Language Model to Large Pushdown Automata Model

  • 在Transformer层间引入可学习的隐藏状态栈,实现类下推自动机行为。
  • 在语法和自然语言任务中均超越标准Transformer,360M模型性能超2-3倍参数的开源模型。
  • 兼容现有框架,支持从3.6亿到70亿参数的规模扩展,适合需要强逻辑推理的应用。

Transformer架构虽推动了大语言模型(LLM)的发展,但其仍无法有效捕捉乔姆斯基层级中的正则表达式或确定性上下文无关语法。受下推自动机利用栈处理上下文无关语法的启发,本文提出StackTrans,通过在Transformer层间显式引入可学习的隐藏状态栈来解决该问题。该设计保持与flash-attention等现有框架兼容,栈操作如压栈、出栈均为可微且端到端可训练。在涵盖乔姆斯基层级与大规模自然语言的多个基准测试中,StackTrans持续优于标准Transformer及其他基线模型。我们成功将模型规模扩展至360M至7B参数。特别地,从头预训练的StackTrans-360M在多项任务中表现超越多个参数量为2-3倍的开源大模型,验证了其高效性与强大推理能力。

原文摘要 · Abstract (English)

The Transformer architecture has emerged as a landmark advancement within the broad field of artificial intelligence, effectively catalyzing the advent of large language models (LLMs). However, despite its remarkable capabilities and the substantial progress it has facilitated, the Transformer architecture still has some limitations. One such intrinsic limitation is its inability to effectively capture the Chomsky hierarchy, such as regular expressions or deterministic context-free grammars. Drawing inspiration from pushdown automata, which efficiently resolve deterministic context-free grammars using stacks, we propose StackTrans to address the aforementioned issue within LLMs. Unlike previous approaches that modify the attention computation, StackTrans explicitly incorporates hidden state stacks between Transformer layers. This design maintains compatibility with existing frameworks like flash-attention. Specifically, our design features stack operations -- such as pushing and popping hidden states -- that are differentiable and can be learned in an end-to-end manner. Our comprehensive evaluation spans benchmarks for both Chomsky hierarchies and large-scale natural languages. Across these diverse tasks, StackTrans consistently outperforms standard Transformer models and other baselines. We have successfully scaled StackTrans up from 360M to 7B parameters. In particular, our from-scratch pretrained model StackTrans-360M outperforms several larger open-source LLMs with 2-3x more parameters, showcasing its superior efficiency and reasoning capability.

大模型语法理解栈机制推理能力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。