Inner Thinking Transformer: Leveraging Dynamic Depth Scaling to Foster Adaptive Internal Thinking

📄 arXiv: 2502.13842v2 📥 PDF

作者: Yilong Chen, Junyuan Shang, Zhenyu Zhang, Yanxi Xie, Jiawei Sheng, Tingwen Liu, Shuohuan Wang, Yu Sun, Hua Wu, Haifeng Wang

分类: cs.CL

发布日期: 2025-02-19 (更新: 2025-02-23)

备注: 15 pages, 11 figures

💡 一句话要点

提出Inner Thinking Transformer，通过动态深度缩放提升LLM在关键token上的推理能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 动态深度缩放 自适应计算 Transformer 模型优化 推理加速 资源受限 梯度优化

📋 核心要点

  1. 现有大语言模型在参数受限时，处理复杂推理token的能力不足，梯度在各层间出现突变。
  2. Inner Thinking Transformer (ITT) 将 Transformer 层计算视为隐式思考步骤，动态分配计算资源。
  3. ITT 在小参数规模下实现了接近大模型的性能，减少了训练数据需求，并在多个基准测试中超越现有模型。

📝 摘要（中文）

大型语言模型（LLM）在参数约束下存在固有的性能瓶颈，尤其是在处理需要复杂推理的关键token时。实证分析表明，具有挑战性的token会在各层中引起突发的梯度峰值，暴露了标准Transformer中的架构应力点。基于此，我们提出了Inner Thinking Transformer（ITT），它将层计算重新构想为隐式思考步骤。ITT通过自适应Token路由动态分配计算，通过残差思考连接迭代细化表示，并使用思考步骤编码区分推理阶段。ITT无需扩展参数即可对关键token进行更深入的处理。在1.62亿至4.66亿参数的模型上的评估表明，ITT仅使用1.62亿参数即可达到4.66亿参数Transformer的96.5％的性能，减少了43.2％的训练数据，并且在11个基准测试中优于Transformer/Loop变体。通过在推理过程中实现弹性计算分配，ITT通过架构感知的隐式思考路径优化来平衡性能和效率。

🔬 方法详解

问题定义：现有的大型语言模型在参数量受限的情况下，难以有效处理需要复杂推理的关键token。这些token在Transformer的各层之间传递时，会导致梯度出现剧烈的波动，表明模型在这些层上承受了较大的压力。现有的Transformer架构对于不同token的处理方式缺乏自适应性，导致计算资源的浪费和性能瓶颈。

核心思路：论文的核心思路是将Transformer的每一层计算视为一个“思考步骤”，并允许模型根据token的重要性动态地调整每个token所经历的思考步骤的深度。通过这种方式，模型可以对关键token进行更深入的推理，而对不重要的token则减少计算量，从而在参数量不变的情况下提升模型的整体性能。

技术框架：Inner Thinking Transformer (ITT) 的整体架构基于标准的Transformer，但引入了三个关键模块：自适应Token路由（Adaptive Token Routing）、残差思考连接（Residual Thinking Connections）和思考步骤编码（Thinking Step Encoding）。自适应Token路由负责决定每个token需要经历多少个思考步骤；残差思考连接用于在不同的思考步骤之间传递信息，促进迭代优化；思考步骤编码则用于区分不同的思考阶段，帮助模型更好地理解token的推理过程。

关键创新：ITT最重要的技术创新在于其动态深度缩放机制，即根据token的重要性自适应地调整其在Transformer中经历的层数。这与传统的Transformer固定层数的设计形成了鲜明对比，使得模型能够更加高效地利用计算资源，从而在参数量受限的情况下实现更好的性能。

关键设计：自适应Token路由模块使用一个可学习的门控机制来决定每个token是否需要进入下一层进行更深入的思考。残差思考连接则类似于ResNet中的残差连接，将前一个思考步骤的输出与当前思考步骤的输出相加，以避免梯度消失问题。思考步骤编码则使用一个可学习的嵌入向量来表示每个思考步骤，并将该向量添加到token的表示中，以区分不同的思考阶段。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，仅使用1.62亿参数的ITT模型可以达到4.66亿参数的Transformer模型的96.5%的性能。此外，ITT还能够减少43.2%的训练数据需求，并在11个基准测试中优于Transformer及其Loop变体。这些结果充分证明了ITT在提升模型性能和效率方面的优势。

🎯 应用场景

Inner Thinking Transformer (ITT) 可应用于各种自然语言处理任务，尤其是在资源受限的场景下，例如移动设备上的文本生成、低功耗设备上的机器翻译等。通过动态调整计算深度，ITT可以在保证性能的同时降低计算成本，从而实现更高效的语言模型部署。此外，ITT的设计思想也可以启发其他领域的模型优化，例如计算机视觉和语音识别。

📄 摘要（原文）

Large language models (LLMs) face inherent performance bottlenecks under parameter constraints, particularly in processing critical tokens that demand complex reasoning. Empirical analysis reveals challenging tokens induce abrupt gradient spikes across layers, exposing architectural stress points in standard Transformers. Building on this insight, we propose Inner Thinking Transformer (ITT), which reimagines layer computations as implicit thinking steps. ITT dynamically allocates computation through Adaptive Token Routing, iteratively refines representations via Residual Thinking Connections, and distinguishes reasoning phases using Thinking Step Encoding. ITT enables deeper processing of critical tokens without parameter expansion. Evaluations across 162M-466M parameter models show ITT achieves 96.5\% performance of a 466M Transformer using only 162M parameters, reduces training data by 43.2\%, and outperforms Transformer/Loop variants in 11 benchmarks. By enabling elastic computation allocation during inference, ITT balances performance and efficiency through architecture-aware optimization of implicit thinking pathways.