近日，上海交通大学自然科学研究院/数学科学学院深度学习基础理论团队许志钦、张耀宇及其学生系统研究了复杂度控制对 Transformer 模型复合任务泛化能力的影响机制。研究发现，较低模型复杂度（例如较小的初始化尺度或较大的权重衰减系数）能够显著促进模型学习具备推理结构的组合规则，从而提升分布外泛化能力。Transformer 作为大语言模型的核心架构，已在自然语言处理与计算机视觉领域取得突破性进展。然而，在结构简单但具有组合性质的任务上，如何实现真正的组合泛化（compositional generalization）仍然是关键挑战。本研究在理论与机制层面填补了这一重要空白。相关成果以 “Complexity Control Facilitates Reasoning-Based Compositional Generalization in Transformers” 为题，发表于人工智能领域顶级期刊 IEEE Transactions on Pattern Analysis and Machine Intelligence (TPAMI)。

研究背景

在语言与图像任务中表现卓越的 Transformer，在面对分布外的简单复合任务时，仍可能出现泛化失效。这一现象背后的原因及其解决路径，长期以来是 Transformer 泛化研究的重要议题。针对该问题，研究团队创新性地以“小初始化”为突破口，系统分析了初始化因子与模型泛化行为之间的关系，并进一步将该分析框架推广至常见正则化手段（如权重衰减），提出了统一的复杂度控制（complexity control）理论框架。

研究方法与结果

具体而言，研究首先基于 anchor function 构造了可解释的合成复合任务，并通过精细的数据划分设计，为模型预设三类潜在解：泛化解（低复杂度）：模型学习到基本函数结构，并能组合推导复合函数，实现真正的分布外泛化；对称解（中复杂度）：模型记忆复合函数整体映射，但未抽象到底层基本函数；记忆解（高复杂度）：模型仅通过过拟合记忆训练数据中的所有组合。三类解的复杂度依次递增。实验表明：在较大初始化下模型最终会演化为高复杂度的记忆解，在正常初始化下选择中复杂度的对称解，在小初始化下选择泛化解。进一步地，作者通过对部分 token 进行掩码分析，揭示不同阶段模型内部机制的差异。降维可视化结果显示：从记忆解到对称解，模型开始识别复合函数的可交换性，在隐空间中对等价结构进行聚类；从对称解到泛化解，模型进一步学会对基本函数进行抽象表示，从而实现逐步推理式计算。

团队进一步发现小初始化的 transformer 会出现明显的凝聚现象，凝聚现象是神经网络非线性学习过程中的一个普遍特征，它使网络的有效神经元数目远小于实际神经元数目，从而使得高度复杂的神经网络在保证拟合数据的前提下，仍保持尽可能低的模型复杂度。此外，作者将小初始化推广到常见的正则化方法，如权重衰减并统称为复杂度控制。这些方法在简单任务、真实语言任务和图像复合任务上均表现出更强的泛化能力。

上图提供了在简单复合任务（左）和图像复合任务（右）不同的复杂度下模型的表现。The figure above illustrates model performance under different complexity settings on simple compositional tasks (left) and image compositional tasks (right).

团队简介

学生张众望和林鹏潇为本工作的共同第一作者，通讯作者为许志钦。许志钦为上海交通大学自然科学研究院/数学科学学院教授，张耀宇为上海交通大学自然科学研究院/数学科学学院副教授，他们也是交大致远学院2008级理科班首届毕业生。张众望为该团队五年级博士生，林鹏潇为三年级博士生。

此外，该团队已在 NeurIPS 2024 发表本工作的前期成果Initialization is critical to whether transformers fit composite functions by reasoning or memorizing。

目前上海交通大学自然科学研究院和数学科学学院已经形成了一批从事深度学习基础研究的科研人员并发表了一系列相关工作。该团队主要有两个系列工作，分别是频率原则和参数凝聚。

本研究得到国家重点研发计划青年科学家项目（2022YFA1008200）、国家自然科学基金项目（92270001、12371511、1242211、12101402）、临港实验室（LG-QS-202202-0）、上海市科技重大专项（2021SHZDZX0102）等支持，并获得上海交通大学思源一号超算、数学科学学院超算中心及学生创新中心的计算资源支持。

关于TPAMI

IEEE TPAMI是人工智能、模式识别、计算机视觉等领域的国际顶尖期刊，2023年度最新发布影响因子为18.6。该期刊谷歌指数（H-Index）在计算机科学和工程技术两个大类学科里均列首位。IEEE TPAMI以严苛的审稿过程、深刻的理论分析著称，在计算机科学与人工智能领域具有权威影响力。