OPTScientist：面向Transformer预训练的带类型优化器程序的多智能体发现

# OPTScientist：用于Transformer预训练的类型化优化器程序的多智能体发现
来源：https://arxiv.org/html/2607.20486
李忠正1,2,3，冯天灿3，李文浩3，冉庆松3，冯诗坤3，张晓媛3，王越3，赵晓光1
1中国科学院自动化研究所
2中国科学院大学
3中关村学院
zhangxiaoyuan@bza\.edu\.cn

###### 摘要

为现代深度学习设计优化器仍然是一个具有挑战性的科学问题，需要综合考虑优化几何、状态动力学、数值稳定性、实现约束和泛化经验。现有的自动化优化器发现方法通常要么在无约束的代码空间中搜索，要么在参数化受限的优化器家族内搜索。前者灵活但常常产生无效或不可解释的程序，而后者稳定但限制了新颖性。我们提出OPTScientist，一个理论引导的多智能体框架，用于在类型化领域特定语言（DSL）中发现优化器。OPTScientist将优化器设计形式化为一个受约束的科学搜索过程，其中候选更新通过方向、缩放、预处理、正则化、状态和分组模块来表达。四个角色智能体——理论家、设计师、工程师和评审员——在一个统一的编排循环中协作，提出假设、合成DSL候选、编译和评估优化器，并评论结果。为了克服固定搜索空间的局限性，OPTScientist结合了优化器程序的进化搜索与第二阶段机制，当重复失败揭示表征瓶颈时，该机制会提出小的DSL扩展。利用此框架，我们发现了RS-MR，一种减少状态的矩阵优化器，在我们本地评估协议下，相比强基线改进了Transformer预训练。我们的结果指向一条自动优化器科学的路径，该科学基于理论、类型化程序、编译器验证和闭环实验。

## 1 引言

优化器是现代深度学习系统的核心（Kingma and Ba, 2014 (https://arxiv.org/html/2607.20486#bib.bib33); Loshchilov and Hutter, 2019 (https://arxiv.org/html/2607.20486#bib.bib6)）。它们决定了训练速度、稳定性、内存消耗、最终验证性能以及在固定计算预算下扩展模型的可行性。尽管取得了实质性进展，优化器设计在很大程度上仍由专家驱动（Zhao et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib2); Liu et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib1); Vyas et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib3); Zhang et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib4); Liu et al., 2025a (https://arxiv.org/html/2607.20486#bib.bib5)）。像AdamW、Lion、Muon和Shampoo这样的方法对更新几何、状态累积、缩放、预处理和正则化做出了不同的假设，但这些假设通常是通过直觉、部分理论和大量经验迭代的结合引入的。

困难在于，优化器不仅仅是梯度到更新的映射公式。它是一个涉及随机梯度、历史状态、自适应尺度、数值保护、参数分组规则和实现约束的耦合动力系统。微小的设计变化可能显著改变收敛性和稳定性。因此，优化器设计空间很大，但只有一小部分在科学上有意义、计算上可行且在现实训练协议下鲁棒（Liu et al., 2025b (https://arxiv.org/html/2607.20486#bib.bib10); Wen et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib11)）。

自动优化器发现提供了一种有前景的替代方案，但现有方法面临着表达性与可靠性之间的张力（Boiko et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib34); Baek et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib35)）。无约束程序搜索或基于语言模型的代码生成可以探索广泛的更新规则，但常常产生无效、不稳定、冗余或难以解释的候选。相反，在固定的优化器家族内搜索提高了可靠性，但将发现限制在已知方法的局部变体。缺失的是一个保留专家级理论先验，同时使优化器设计可执行、可搜索、可审计和可扩展的框架（Ghafarollahi and Buehler, 2025 (https://arxiv.org/html/2607.20486#bib.bib36); Hong et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib37)）。

我们提出OPTScientist，一个理论引导的多智能体框架，用于将优化器发现为类型化程序。OPTScientist不生成任意Python代码，而是在类型化优化器领域特定语言（DSL）内搜索，其中候选更新通过方向、缩放、预处理、正则化、状态和分组模块来表达。一个编译器解析每个候选，检查张量和标量兼容性，将其降级为中间表示，并生成可执行优化器。这种编译器支持的表征在训练前拒绝了许多无效候选，并使发现的优化器保持紧凑和可检查。

在搜索层面，OPTScientist将优化器发现组织为一个科学循环。理论家从共享记忆中提出搜索假设和先验，设计师将其转化为多样的DSL候选，工程师编译、修复和评估候选，评审员根据新颖性、有效性、性能、目标阶段行为以及可扩展性对其进行评分。这些角色在一个统一的编排循环中协调，而不是作为独立过程。由此产生的工作流将优化器发现从盲目变异转变为假设生成、程序合成、执行、批评和修订的迭代过程。

OPTScientist的一个关键组件是两阶段进化。在第一阶段，进化算法在当前DSL内搜索优化器程序。在第二阶段，当重复失败、目标阶段停滞或编译器瓶颈揭示表征局限性时，系统会提出保守的DSL扩展，例如安全的宏式组合原语。因此，OPTScientist不仅进化优化器候选，而且在必要时也进化用于表达它们的语言。

利用此框架，我们发现了Reduced-State MAGMA-RowNorm（RS-MR），一种用于Transformer权重矩阵的低状态矩阵优化器。RS-MR结合了行归一化极方向、轻量级全局阻尼和软块级更新门控。与更重的预处理方法不同，它仅用紧凑的块级状态就改善了矩阵更新的选择性。在我们原生Transformer预训练基准中，RS-MR持续优于Muon等强基线，同时保持接近Muon的优化器状态内存占用。

我们的贡献如下：

- • 我们提出了一种类型化程序的自动优化器发现形式化，其中候选优化器在编译器支持的DSL中表达，而不是作为无限制的Python代码生成。
- • 我们开发了OPTScientist，一个理论引导的多智能体框架，在可审计的科学搜索循环内协调假设生成、DSL候选合成、编译、评估、评审员评分和保守的DSL进化。
- • 我们提出了RS-MR，一个发现的用于Transformer预训练的减少状态矩阵优化器，并证明它在原生多阶段验证性能上优于强基线，同时保持了有利的内存-性能权衡。

## 2 相关工作

### 2.1 Transformer预训练的优化器

优化器设计一直是大规模神经网络训练的核心。经典一阶方法如带动量的SGD和Nesterov加速提供了简单而鲁棒的更新规则，而自适应方法如Adam和AdamW因其稳定性和易于调参而被广泛用于Transformer预训练（Kingma and Ba, 2014 (https://arxiv.org/html/2607.20486#bib.bib33); Loshchilov and Hutter, 2019 (https://arxiv.org/html/2607.20486#bib.bib6)）。内存高效变体如Adafactor降低了大型语言模型的二阶矩内存成本（Shazeer and Stern, 2018 (https://arxiv.org/html/2607.20486#bib.bib7)），而最近的优化器如Lion、Sophia、Shampoo式预处理器和Muon式矩阵优化器探索了基于符号的更新、曲率近似和矩阵结构归一化（Zhao et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib2); Liu et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib1); Vyas et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib3); Zhang et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib4); Liu et al., 2025a (https://arxiv.org/html/2607.20486#bib.bib5)）。

尽管取得了这些进展，大多数优化器仍然是通过大量经验验证手动设计和调整的，需要关于更新几何、状态设计、数值稳定性和硬件效率的专家直觉（Shazeer and Stern, 2018 (https://arxiv.org/html/2607.20486#bib.bib7); Yuan et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib8); Pagliardini et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib9)）。这一挑战在Transformer预训练中因异质性参数类型、长视野训练动态以及短视野代理性能与最终训练行为之间的不匹配而加剧（Liu et al., 2025b (https://arxiv.org/html/2607.20486#bib.bib10); Wen et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib11)）。OPTScientist的不同之处在于将优化器设计视为可搜索的类型化程序空间，并通过原生训练验证候选；发现的RS-MR优化器表明，减少状态的矩阵机制（如行归一化极方向和软块级门控）可以在不依赖重二阶状态的情况下提升性能。

### 2.2 自动优化器发现与程序搜索

另一条工作线研究了学习规则和优化算法的自动发现。学习型优化器用神经网络参数化更新规则并通过元学习进行训练（Andrychowicz et al., 2016 (https://arxiv.org/html/2607.20486#bib.bib14); Romera-Paredes et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib23); Novikov et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib24)），但它们往往难以解释、训练成本高且对元训练分布敏感。另一个方向是在符号或程序化空间中搜索，其中遗传编程、进化算法和AutoML风格系统发现更新公式、神经组件或完整算法（Koza, 1994 (https://arxiv.org/html/2607.20486#bib.bib15); Real et al., 2020 (https://arxiv.org/html/2607.20486#bib.bib16); Bello et al., 2017 (https://arxiv.org/html/2607.20486#bib.bib17)）；符号优化器发现也产生了如Lion这样的实用方法（Chen et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib12)）。

现有的自动发现方法面临表达性与可靠性之间的权衡。无约束代码搜索灵活但容易产生无效程序、审计性差和实现行为脆弱（Moudgil et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib18); Marfinetz, 2025 (https://arxiv.org/html/2607.20486#bib.bib19)），而参数化受限的空间更容易验证但将发现限制在已知机制的局部变体（Yang et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib20); Liu et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib21); Ye et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib22)）。OPTScientist通过在类型化优化器DSL内搜索来解决这一权衡，该DSL强制执行显式状态类型、张量操作和更新语义，同时仍然支持动量、行/列统计、极更新、块掩码和结构化预处理等机制。

### 2.3 用于科学发现的LLM智能体

大型语言模型越来越多地被用作推理、调用工具、编写代码和迭代改进解决方案的智能体。早期的智能体方法将语言模型推理与外部动作或反馈相结合，包括工具使用智能体、ReAct风格推理、Reflexion风格自我改进和自主探索系统（Schick et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib27); Yao et al., 2022 (https://arxiv.org/html/2607.20486#bib.bib28); Shinn et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib29); Wang et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib30)）。多智能体框架将复杂任务分解为规划、实现、批评或验证等专门角色（Wu et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib31); Li et al., 2023 (https://arxiv.org/html/2607.20486#bib.bib32)），而基于LLM的科学发现系统已被用于生成假设、提出代码、搜索程序以及改进数学或算法构造（Romera-Paredes et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib23); Lu et al., 2024 (https://arxiv.org/html/2607.20486#bib.bib25); Schmidgall et al., 2025 (https://arxiv.org/html/2607.20486#bib.bib26)）。

OPTScientist与这条工作线相关，但专注于Transformer预训练优化器设计这一受约束且经验扎实的领域。它并非允许智能体自由编辑任意代码，而是为每个角色分配特定的科学功能：理论家提出假设，设计师实例化DSL程序，工程师编译并评估它们，评审员对候选进行评分并决定是否应最小限度地扩展DSL。与通用LLM智能体相比，OPTScientist强调编译器支持的合法性、原生训练评估、显式搜索记忆和受控的DSL进化。

## 3 预备知识

#### 优化器作为模块化更新程序。

现代优化器的表面形式差异很大，但其中许多可以理解为在 θt+1 = θt - ηtUt 中构建更新张量 Ut 的不同方式。例如，SGD选择原始梯度或动量作为更新方向，Adam类方法通过二阶矩统计添加自适应缩放，Lion使用基于符号的方向，Muon应用矩阵正交化方向，Shampoo类方法引入结构化预处理。我们利用这一观察来定义结构化的优化器设计空间：优化器更新可以分解为可重用模块，涉及方向、缩放、预处理或几何、正则化、状态和参数分组。这种分解并非旨在成为唯一的数学分解；相反，它为搜索优化器机制提供了一个实用的坐标系统。

#### 优化器DSL。

OPTScientist将每个候选优化器表示为一个类型化领域特定语言（DSL）程序。系统不是要求语言模型生成任意的Python优化器代码，而是在由元数据声明、超参数声明、状态更新、中间表达式和最终更新表达式组成的紧凑DSL程序中进行搜索。在高级层面，DSL程序指定如何从当前梯度、参数张量、超参数和持久优化器状态计算Ut。这种设计使优化器发现更接近于在结构化科学语言上的程序合成，而非无约束的代码生成。

#### 原子优化器模块。

DSL在多个层级公开优化器构建模块。