
中科院、百度提出新架构:突破参数限制,实现高效推理
有不少研究证明,大模型在不断增加参数的情况下,其性能提升却逐渐趋于饱和,特别是在受限的数据可用性和计算资源预算下,这种情况非常明显。
为了解决这一难题,中国科学院信息工程研究所、百度、北京师范大学AI学院的研究人员联合提出了Inner Thinking Transformer架构(简称ITT),通过动态分配计算资源给单个标记,增强了测试性能而无需增加参数。
ITT 的设计基于一个关键假设:Transformer 的每一层可以被视为对输入信息的一次隐式推理。研究人员提出将单个token的生成过程分解为多个内部思考步骤,模型从初始状态开始,通过一系列的思考步骤逐步更新隐藏状态,最终生成输出token。
这种设计允许模型在遇到复杂推理任务时,通过增加思考步骤来提升性能,而无需扩展模型的参数规模。
为了实现这一目标,ITT 引入了 Residual Thinking Connection (RTC) 机制。RTC 的作用是通过迭代累积每个思考步骤的输出,逐步细化token的表示。与传统的循环神经网络相比,RTC 不仅能够实现深度思考,还能通过学习每个步骤的重要性编码来衡量和整合不同思考步骤的结果。这种机制使得模型能够在训练过程中逐步优化思考过程,并在推理时根据需要动态调整思考步骤的数量。
为了进一步提升计算效率,ITT 引入了 Adaptive Token Routing (ATR) 机制。ATR 的核心思想是动态选择需要进一步思考的关键token,避免对所有token进行不必要的深度处理。模型通过一个路由网络为每个token分配一个权重,权重较高的token会被选中进行额外的思考步骤。这种选择性处理不仅减少了计算开销,还能使模型更加专注于对任务至关重要的信息。
在训练过程中,路由网络会根据当前token的表示和思考步骤的编码来决定是否对其进行进一步处理,使得模型能够在训练时学习到最优的思考模式,并在推理时根据需要灵活调整计算资源的分配。通过这种方式,ITT 实现了在不增加模型参数的情况下,对关键token进行深度处理的目标。
除了动态思考步骤和计算资源分配,ITT 还引入了 Thinking Step Encoding 机制。该机制通过为每个思考步骤分配一个独特的编码,帮助模型区分不同的思考阶段。这种编码不仅能够衡量每个步骤的重要性,还能指导模型如何将不同步骤的结果进行整合。
在优化过程中,ITT 的设计使得模型能够通过多步优化来替代传统的单步优化。每个思考步骤的参数更新是基于全局损失的梯度乘以局部映射的导数。这种设计使得模型在反向传播时能够更稳定地收敛,并避免了梯度消失或爆炸的问题。通过这种方式,ITT 实现了在不增加模型参数的情况下,通过动态扩展思考步骤来提升模型的推理能力和性能。
为了测试ITT性能,研究人员使用了 RedPajama 数据集,这是一个包含 500 亿训练标记和 200 万验证标记的多领域数据集。模型训练基于 Sheared-LLaMA 代码框架,在 8 个 NVIDIA A100 GPU 上进行,序列长度为 4096,全局批量大小为 256。
实验结果显示,ITT 在不同参数规模(162M、230M和 466M 参数)的模型上均显著优于基线Transformer 和 Loop 模型。
例如,在 162M 参数规模下,ITT ×4 模型通过在 50% 的层中进行 4 步思考,相比基线模型提升了 1.7% 的性能,而 Loop 模型仅提升了 0.3%。这种性能提升在更大规模的模型中也得到了体现,表明 ITT 的设计能够有效扩展到不同规模的模型。
ITT 的另一个显著优势是其数据效率。在训练过程中,ITT只需使用基线模型 56.8% 的数据量,即可达到与基线模型相同的性能水平。这一结果表明,ITT 通过动态思考步骤和计算资源分配,能够更高效地利用训练数据,减少对大规模数据集的依赖
论文地址:https://arxiv.org/abs/2502.11089
本文转自 AIGC开放社区 ,作者:AIGC开放社区
