Dust:无需反向传播的Transformer零阶预训练

TL;DR

Dust是首个在Transformer预训练中能与反向传播相媲美的零阶方法。通过对激活进行逐token扰动(一种虚拟种群),它能在单次前向传播中评估数千个候选,从而获得随种群增长而收敛于反向传播的梯度估计,甚至在大计算量场景下超越反向传播。


1. Dust的不同之处

激活空间扰动

  • Dust为每个线性层的输出逐token添加独立的高斯噪声。因此,每个token充当独立的种群成员,使得单次前向传播能够评估的种群规模与序列长度相当(数千个成员),而不是像权重空间进化策略那样每个成员对应一份权重副本。
  • 这种“虚拟种群”避免了物化权重副本;扰动存在于激活空间中,对于给定层而言,其维度低于完整权重矩阵。

无需反向传播的信用分配

  • 对于线性层(y_t = Wx_t),Dust根据该token处的损失减少量(以及可选地,通过注意力加权的未来token折扣)计算每个token的奖励(r_t)。奖励加权噪声(\hat g_t = -\frac{1}{K\sigma}\sum_{i=1}^K r_t^{(i)} a_t^{(i)})作为估计的输出误差。
  • 权重梯度估计随后是(\hat g_t)与缓存的输入(x_t)的外积: [\hat G_W = \sum_t \hat g_t x_t^{\top}]
  • 注意力内部机制类似处理,但其奖励来源于注意力输出的变化而非直接逐token损失,从而减少干扰。

减少干扰

  1. 逐层传递 – 不同类型层在单独的前向传播中抖动,各自使用独立的噪声尺度。
  2. 注意力特定抖动 – 仅重新计算注意力输出以进行抖动,其余网络保持缓存。
  3. 仅头部扰动 – 语言模型头部在缓存的logits上抖动,每次抽取仅评估词汇表的一小部分。

这些技巧使得每次抽取的成本仅略高于干净的前向传播,同时允许每次更新使用(10^3)–(10^4)次抽取的种群规模。


2. 实证发现

预训练结果(SGD动量)

Token预算 种群(抽取次数) 测试损失(Dust) 测试损失(反向传播) 测试损失(EGGROLL)
100 k 64-1 k 低于反向传播 – –
1 M 1 k-16 k 接近反向传播,差距< 0.02 – –
10 M 4 k-16 k 略高于反向传播,趋势下降 – –
20 M 16 k 外推表明Dust可能低于反向传播(极限≈4.43 vs 反向传播4.63) – –

Dust始终保持在反向传播损失的百分之几以内,并且在最大种群下,拟合的幂律极限低于反向传播。

与权重空间进化策略(EGGROLL)的比较

  • 即使使用**10³–10⁴×**更大的种群,EGGROLL在所有token预算下仍落后Dust 0.02–0.6损失单位。
  • 外推表明,EGGROLL需要10³–10⁴倍的种群才能匹配Dust最小种群下的性能。

模型规模扩展

参数 种群64 256 1 k 4 k 16 k
2 M 5.705 5.486 5.265 5.189 5.171
7 M 5.556 5.362 5.161 5.095 5.065
38 M 5.558 5.358 5.158 5.053 5.036
243 M 5.719 5.419 5.214 5.124 5.086

更大的模型*种群效率更高:即使在最小种群下,243 M参数模型也优于2 M模型,并在小模型达到平台期时继续改进。*

梯度对齐

Dust的估计梯度与真实反向传播梯度的对齐程度随种群增加而提高。余弦相似度遵循: [\cos(K) = \frac{c_{\text{max}}}{\sqrt{1 + c/K}}]

  • 跨层拟合的RMSE < 0.06。
  • 在种群≥16 k时,大多数层的余弦值超过0.9,表明方向几乎一致。
  • 相比之下,EGGROLL的余弦值即使在128 k抽取时仍保持< 0.05,这解释了其训练性能不佳的原因。

3. Hacker News上的社区反应

评论者 关键见解/批评
blt 警告无导数方法历来未能影响实践;认为梯度对于平滑目标仍然必不可少,并引用了最近的复杂度差距证明。
syntacticsalt 质疑“苦涩教训”的框架,指出零阶方法本身并不能解决非凸性,且论文缺乏对凸性假设的讨论。
eriwang915 强调了一个令人惊讶的结果:243 M参数的Dust模型在大多数种群规模下优于120倍更小的模型。
polyomino 建议混合流程:使用Dust对反向传播训练的检查点进行微调,以探索后期训练阶段是否受益于基于搜索的信用分配。
soltanov 呼吁测量墙钟时间、能量和内存以评估实际可行性;指出在Dust取代反向传播之前,当前的效率差距必须缩小。
mkaic 指出零阶方法在梯度不可用的情况下(例如基于模拟器的训练)可能更有价值,而非作为通用替代品。
usernametaken29 认为Dust和反向传播都处于相同的经验风险最小化前沿,但Dust可能绕过阻碍反向传播的Hessian条件限制。
api 观察到Dust每次更新的计算效率较低,但高度可并行化,这可能在大规模硬件上具有优势。

4. Dust的重要性

  1. 基于搜索的信用分配的概念验证 – 证明纯前向算法可以为大规模语言模型生成可用梯度,挑战了零阶方法无法扩展的长期信念。
  2. 在计算丰富场景中超越反向传播的潜力 – 在非常大的种群下,Dust的损失曲线外推低于反向传播,表明只要有足够的计算量,该算法可能找到比链式法则梯度更好的下降方向。
  3. 支持训练不可微架构 – 由于Dust不需要解析梯度,它为包含离散模拟器、外部程序或难以微分的循环的网络打开了大门。
  4. 架构洞察:过参数化的好处 – 更大的模型表现出更好的种群效率,暗示高维激活空间的几何形状可能比先前认为的更有利于基于搜索的优化。

5. 开放问题与未来方向

  • 效率差距 – Dust目前需要比反向传播多几个数量级的FLOPs。需要工程进步(例如专用硬件、更智能的种群分配)才能使其实用。
  • 混合训练机制 – 将反向传播预训练与Dust微调相结合,可以发挥两种方法的优势。
  • 高阶景观探索 – 理解Dust是否隐式捕获曲率信息(超越一阶)可能解释其偶尔的优越性。
  • 更广泛的适用性 – 将Dust应用于梯度不可用的任务(强化学习与模拟器、神经形态硬件)将测试其通用性。
  • 理论基础 – 对种群规模与梯度对齐之间的权衡进行形式化分析,特别是在非凸设置中,将巩固实证观察。

6. 结论

Dust表明,零阶激活扰动可以在与反向传播相当的水平上训练Transformer语言模型,更大的模型甚至变得更具样本效率。尽管由于计算需求,该方法远未在日常实践中取代反向传播,但它提供了一个引人注目的概念验证,即基于搜索的信用分配可以扩展,为不可微架构和计算丰富的优化机制开辟了新的研究途径。

Sources

相关