1. 从Lion到CLion:为什么我们需要重新审视符号化优化器的泛化性?

在深度学习的工具箱里,优化器一直是个“既要又要”的难题。我们既希望它收敛得快,能在有限的算力下快速找到损失函数的谷底;又希望它收敛得好,找到的那个谷底足够平坦、足够宽,这样模型面对没见过的数据时,才不会一脚踏空,性能暴跌。这个“好”的衡量标准,就是泛化能力。过去几年,从SGD到Adam,再到各种自适应学习率算法的变体,大家的研究重心多在收敛速度上较劲,看谁能在训练集上“跑”得更快。但一个残酷的现实是:在训练集上跑得飞快的算法,有时在测试集上会摔得很惨,这就是泛化性不佳的表现。

Lion(EvoLved Sign Momentum)优化器的出现,带来了一股新风。它做了一件很激进的事:把传统的动量向量直接压缩成了一个符号向量。简单来说,它不关心梯度具体有多大,只关心梯度的方向是正还是负。这个操作带来的好处是显而易见的——内存占用大幅降低,因为只需要存储+1、-1或0,而不是高精度的浮点数。在大模型训练中,这能省下可观的内存,从而允许使用更大的批次(batch size)或模型。但业内一直有个悬而未决的疑问:这种激进的“符号化”压缩,会不会破坏优化过程的稳定性,进而损害模型的泛化能力?直觉上,丢失了梯度的大小信息,就像蒙着眼睛走下山路,虽然知道该往左还是往右,但不知道步子该迈多大,似乎更容易掉进某个尖锐的局部极小点,导致泛化性能变差。

这正是我们工作的起点。我们不能仅仅满足于“它在实验上work了”。对于要部署到关键应用中的模型,我们需要理论上的保证。我们需要确切地知道,使用Lion优化器训练出的模型,其训练误差和期望误差(即泛化误差)之间的差距到底有多大?这个差距受哪些因素影响?更重要的是,我们能否改进它?

经过一番理论推导,我们发现标准Lion的泛化误差上界是O(1/(τTN))。这个τ是个关键但令人不安的量:它代表在整个训练过程中,所有非零梯度分量绝对值的最小值。如果模型在训练中,某个维度的梯度非常小(接近零),那么τ就会很小,导致泛化误差上界变得很大,理论保证就弱了。这在实际训练中很可能发生,尤其是在模型接近收敛、或某些参数对当前任务不敏感时。

于是,CLion(Compressed Lion)算法应运而生。它的核心思想是引入一个自适应阈值ν。我们不再武断地对所有梯度分量进行符号化。当一个梯度分量ct的绝对值足够大(大于等于ν)时,我们才认为它的方向信息是可靠的,对其进行符号化压缩;当它的绝对值很小时,我们保留其原始值,因为此时噪声可能主导了方向信息,盲目取符号会引入较大偏差。这个简单的改进,就像给蒙眼下山的步行者配了一根探路杖,在平坦处大步流星(符号化),在崎岖处小心试探(保留原值)。我们的理论分析证明,CLion能够将泛化误差上界提升到O(1/N),消除了对那个不稳定的τ的依赖,从而提供了更稳健的理论保证。

接下来的内容,我将为你彻底拆解这背后的数学原理、算法细节以及工程实现中的关键考量。无论你是希望深入理解优化器理论的研究者,还是正在为下一个大模型训练寻找高效稳定优化器的工程师,相信都能从中获得启发。

2. 理论基础:理解泛化性分析与算法稳定性的核心框架

在深入Lion和CLion的证明细节之前,我们必须先搭建好理论舞台。泛化性分析不是空中楼阁,它建立在“算法稳定性”这一坚实概念之上。简单来说,如果一个算法是“稳定”的,那么对训练数据集做微小的改动(比如替换一个样本),其输出的模型参数变化不会太大,进而模型在任意样本上的预测输出变化也不会太大。这种输出对输入数据扰动的“不敏感性”,直接关联到泛化能力。

2.1 关键假设:我们分析的前提条件

任何有意义的理论分析都需要明确的假设。我们的分析建立在以下三个标准且合理的假设之上,这些假设在机器学习优化理论中广泛使用:

假设1(L-光滑性) :目标函数F(w)及其随机版本f(w; ξ)是L-光滑的。这意味着梯度函数是L-Lipschitz连续的。用数学语言说,对于任意参数w和v,有 ∥∇F(w) - ∇F(v)∥ ≤ L∥w - v∥。这个假设保证了损失函数的梯度变化不会太快,地形相对“平缓”,使得梯度下降类算法能够稳定前进。L可以理解为地形最陡峭处的坡度上限。

假设2(G-有界梯度) :随机梯度∇f(w; ξ)的范数以高概率有界,即存在常数G > 0,使得对于任意w和几乎所有的ξ,有 ∥∇f(w; ξ)∥ ≤ G。这个假设非常实际,它避免了梯度爆炸的问题。在神经网络中,结合有界激活函数(如Sigmoid, Tanh)或适当的权重初始化,这一条件通常可以满足。

假设3(σ²-有界方差) :随机梯度的方差有界。即对于任意参数w,有 E[∥∇f(w; ξ) - ∇F(w)∥²] ≤ σ²。这里∇F(w) = E[∇f(w; ξ)]是真实梯度(全批梯度)。这个假设刻画了用小批量(mini-batch)梯度估计全批梯度时所产生的噪声大小。σ²越小,意味着用单个或少量样本估计的梯度越可靠。

注意 :这些假设并非为了理论完美而强加的苛刻条件。L-光滑性在训练深度网络时,通过梯度裁剪(Gradient Clipping)等技术可以间接满足或缓解违反该假设带来的问题。G-有界梯度在实践中最常见,我们通常直接对梯度进行裁剪来保证。有界方差则是随机优化中的基本假设。我们的理论结果正是在这些符合实际训练场景的假设下成立的。

2.2 均匀稳定性:连接算法与泛化误差的桥梁

我们分析的核心工具是 均匀稳定性(Uniform Stability) 。它的定义直观而有力:考虑两个仅在一个样本上有差异的训练集S和S'(即S'由S替换一个样本得到)。用同一个随机算法A(如Lion)分别在S和S'上训练,得到输出模型w_T和w'_T。如果对于所有可能的样本ξ,算法A满足: E[|f(w_T; ξ) - f(w'_T; ξ)|] ≤ ε(T, N) 那么我们就称算法A具有ε(T, N)-均匀稳定性。

这里的期望是对算法本身的随机性(如mini-batch的采样)取的。这个不等式意味着,改变训练集中的一个样本,所训练出的模型在 任意 (甚至是替换掉的那个)样本上的损失变化,平均不会超过ε(T, N)。

为什么这能推导出泛化误差?根据经典的稳定性理论(Shalev-Shwartz等人,2010;Hardt等人,2016),一个具有ε-均匀稳定性的算法,其期望泛化误差(即期望风险F(w_T)与经验风险F_S(w_T)之差的期望)的上界就是ε。用公式表示就是: |E[F(w_T) - F_S(w_T)]| ≤ ε(T, N)

因此,我们的全部工作就归结为:分析Lion和CLion算法,并推导出这个稳定性边界ε(T, N)的具体表达式。对于Lion,我们得到了O(1/(τTN));对于CLion,我们得到了更优的O(1/N)。

2.3 证明的核心思路:追踪参数差异的演化

如何分析这个稳定性?我们的策略是“追本溯源”。既然两个模型w_t和w'_t的差异源于数据集那一个样本的不同,那么我们就建立这两个模型参数差异∥w_t - w'_t∥的递归不等式。

从算法的更新公式出发。以标准Lion为例,其核心更新步骤为:

c_t = β1 * m_{t-1} + (1 - β1) * g_t  # 计算“符号动量”
w_t = w_{t-1} - η * (sign(c_t) + λ * w_{t-1}) # 更新参数

其中g_t是当前mini-batch的随机梯度,m_t是传统动量。

当我们在两个数据集S和S'上运行算法,并使用 相同的随机数种子 (即相同的mini-batch采样序列)时,在大多数迭代步中,由于采样的样本相同,两个算法看到的梯度g_t和g'_t是相同的。只有在采样到那个被替换的样本i时,它们才会产生差异。

我们的证明正是通过数学归纳法,严谨地追踪这个差异如何通过动量项m_t和符号动量项c_t进行传播,并最终影响参数w_t。关键的难点在于处理 符号函数sign(·) 。符号函数是非Lipschitz的,在零点附近变化剧烈,这给分析带来了巨大挑战。引理5(即文中的Lemma 2)就是为了解决这个问题而证明的核心技术引理。

3. 核心细节解析:拆解Lion与CLion的泛化性证明

现在,让我们深入到证明中最关键、也最体现技术含量的部分。理解这部分,你就能明白为什么标准Lion的泛化界会依赖τ,而CLion又如何巧妙地摆脱了它。

3.1 符号函数的“温和”边界:引理5的精髓

标准Lion泛化分析的最大障碍来自于sign函数。sign函数在零点不连续,其Lipschitz常数可以认为是无穷大,这导致我们无法直接用梯度差异来约束参数差异。我们需要一个更巧妙的工具。

引理5(重新表述) :假设序列{c_t}和{c'_t}分别由算法在数据集S和S'上生成。那么,对于所有迭代步t,我们有: ∥ sign(c_t) - sign(c' t) ∥ ≤ (2√d / τ) * ∥ c_t - c' t ∥ 其中,τ = min {t≥1} ( min {j ∈ S_t} |(c_t)_j| ),而S_t = {j | (c_t)_j ≠ 0}。这里d是参数的维度。

这个不等式在说什么? 它说,两个符号向量之间的差异,可以被它们对应的原始向量(c_t和c'_t)之间的差异所控制,但有一个放大系数2√d/τ。这个τ正是前面提到的、所有非零c_t分量绝对值的最小值。

证明思路拆解 : 这个证明的核心是分情况讨论,针对向量c_t和c'_t的每一个维度j:

  1. 情况一:符号相同 。如果(c_t)_j和(c'_t)_j同号(同正或同负),那么它们符号的差为0。此时不等式显然成立,因为左边为0,右边非负。
  2. 情况二:符号不同 。这是关键情况。当(c_t)_j和(c'_t)_j异号时,|sign((c_t)_j) - sign((c'_t)_j)| 至少为1(可能为2,如果一个是正/负,另一个是0)。此时,我们需要证明 (2/τ) * |(c_t)_j - (c'_t)_j| ≥ 1。 由于符号相反,且两者都非零(如果有一个为零,根据τ的定义,其绝对值至少为τ,也属于符号不同的情况),那么|c_t - c'_t|_j = |(c_t)_j| + |(c'_t)_j|。根据τ的定义,|(c_t)_j| ≥ τ 且 |(c'_t)_j| ≥ τ。因此,|(c_t)_j - (c'_t)_j| ≥ 2τ。代入即得 (2/τ) * (至少2τ) = 4 ≥ 1。这就证明了对于每个维度,不等式成立。
  3. 最后,将各个维度的1-范数不等式组合起来,利用向量范数之间的关系(∥x∥ ≤ ∥x∥_1 ≤ √d ∥x∥),即可得到最终的2-范数不等式。

这个引理的深刻含义与隐患 : 它成功地将非Lipschitz的sign函数差异,用Lipschitz的向量差异表示了,但代价是引入了一个分母τ。τ是算法运行过程中产生的一个量, 不是我们预先设定的超参数 。如果优化路径上,c_t的某个分量始终非常接近于零但又非零,那么τ就会非常小,导致这个上界非常大,理论保证就变得很弱。在实际训练中,这完全可能发生,尤其是在模型稀疏性较高或某些特征重要性很低时。

3.2 CLion的改进策略:自适应符号化阈值

CLion算法正是为了消除对τ的依赖而设计的。它的更新规则增加了一个条件判断:

if min_{j in S_t} |(c_t)_j| >= ν:
    h(c_t) = sign(c_t)
else:
    h(c_t) = c_t
w_t = w_{t-1} - η * (h(c_t) + λ * w_{t-1})

这里ν是一个我们 预先设定的、大于0的超参数

这个简单的改动带来了理论上的巨大优势。我们现在需要分析的是h(c_t) - h(c'_t)的范数。同样分两种情况:

  1. 当|(c_t)_j| ≥ ν 且 |(c'_t)_j| ≥ ν :此时h就是sign函数。但注意,我们现在有|(c_t)_j| ≥ ν,|(c'_t)_j| ≥ ν。沿用引理5的证明思路,在符号不同的情况下,我们有 |(c_t)_j - (c'_t)_j| ≥ 2ν。因此,我们可以得到一个改进的界:∥sign(c_t) - sign(c'_t)∥ ≤ (2√d / ν) ∥c_t - c'_t∥。由于ν是我们设定的(通常≥1),这个分母是可控的常数。
  2. 当|(c_t)_j| < ν 或 |(c'_t)_j| < ν :此时,根据算法,至少有一个向量的h函数取原值。最坏情况下,h(c_t) = c_t, h(c'_t) = c'_t。那么,∥h(c_t) - h(c'_t)∥ = ∥c_t - c'_t∥。这显然小于等于 (2√d / ν) ∥c_t - c'_t∥(因为ν ≥ 1时,2√d/ν ≤ 2√d)。

因此, 对于CLion,我们总能得到 : ∥ h(c_t) - h(c'_t) ∥ ≤ 2√d ∥ c_t - c'_t ∥ 这个界 不再包含不稳定的τ ,分母从1/τ变成了常数1。这就是CLion能获得O(1/N)泛化界的核心理论原因。

实操心得 :在实现CLion时,阈值ν的选择是一个超参数。理论分析要求ν ≥ 1。在实践中,我们可以将ν设置为一个接近1的值,例如0.1或1.0。这相当于设定了一个“显著性阈值”:只有当梯度分量足够显著(绝对值大于ν)时,我们才信任其方向信息并进行压缩;否则,我们保留其原始值以防信息丢失。这本质上是一种 自适应稀疏化 ,在内存节省和算法稳定性之间取得了更好的平衡。

3.3 递归不等式与归纳法:构建误差传播链条

有了对符号函数(或h函数)差异的控制,我们就可以建立参数差异∥w_t - w'_t∥的递归关系了。

从更新公式出发: w_t - w' t = (1 - ηλ)(w {t-1} - w' {t-1}) - η [ h(c_t) - h(c' t) ] 两边取范数,利用三角不等式和上面得到的界: ∥w_t - w' t∥ ≤ (1 - ηλ) ∥w {t-1} - w' {t-1}∥ + η ∥h(c_t) - h(c' t)∥ ≤ (1 - ηλ) ∥w {t-1} - w' {t-1}∥ + 2η√d ∥c_t - c'_t∥ (对于CLion)

接下来的任务就是递归地估计∥c_t - c' t∥。c_t本身是动量的指数移动平均:c_t = β1 * m {t-1} + (1-β1) * g_t。而m_t = β2 * m_{t-1} + (1-β2) * g_t。因此,c_t和c' t的差异,来源于历史动量m {t-1}的差异和当前梯度g_t的差异。

通过细致的概率分析(区分当前采样的样本是否是那个被替换的样本i),并反复利用L-光滑性和有界方差假设,我们可以建立∥c_t - c' t∥和∥m_t - m' t∥关于∥w {t-1} - w' {t-1}∥的递归不等式。

最终,通过数学归纳法,我们可以证明存在一系列常数φ_t和ψ_t,使得: E[∥w_t - w'_t∥] ≤ φ_t / N E[∥m_t - m'_t∥] ≤ ψ_t / N 对于CLion,在适当的超参数设置下(如η = O(1/√d)),可以证明φ_t和ψ_t是常数阶O(1)的(当T不大时)或O(t)的(当T很大时)。通过设置衰减的学习率η = O(1/√(dT)),可以最终得到E[∥w_T - w'_T∥] = O(1/N)。

最后,利用G-有界梯度的假设(即损失函数是G-Lipschitz的),我们可以将参数差异转化为损失函数值的差异: E[|f(w_T; ξ) - f(w'_T; ξ)|] ≤ G * E[∥w_T - w'_T∥] = O(1/N) 这正是ε-均匀稳定性的定义,其中ε = O(1/N)。再通过稳定性引理,便得到了最终的泛化误差上界:|E[F(w_T) - F_S(w_T)]| ≤ O(1/N)。

4. 收敛性保证:CLion不仅泛化好,还能收敛

一个优化器不能只谈泛化,还必须保证它能收敛到(至少是)一个驻点。对于非凸优化,我们通常追求收敛到平稳点,即梯度范数足够小的点。定理6(Theorem 3)给出了CLion的收敛性保证。

4.1 收敛性分析的核心挑战与引理

收敛性分析的难点同样来自于符号函数和条件判断。我们无法直接对sign(∇F(w))进行泰勒展开。CLion的证明需要两个关键的技术引理:

引理8(参数有界性) :在合理的超参数设置下(λ ≤ 1/(2ηĜT^α), α>1),CLion算法产生的迭代点序列{w_t}满足: ∥w_t∥ ≤ (t+1)ηĜ ∥w_t - w_{t-1}∥ ≤ 2ηĜ 其中 Ĝ = max(G, √d)。这个引理保证了优化路径不会爆炸,是后续分析的基础。证明主要利用了更新公式和梯度/符号的有界性进行归纳。

引理9(动量估计误差) :它量化了算法中使用的“压缩动量”c_t与真实全批梯度∇F(w_{t-1})之间的偏差。结论是: (1/T) Σ_{t=1}^T E[∥c_{t+1} - ∇F(w_t)∥] ≤ [√(2(σ²+G²))] / √((1-β₂)T) + [2√2 L Ĝ η] / (1-β₂) + [√2 |β₁-β₂| σ] / √(1-β₂) + [(1-β₁)σ] / √(1-β₂)

这个上界由四部分组成:

  1. 方差项 :O(1/√T),来源于随机梯度的噪声,随着迭代增加而衰减。
  2. 梯度变化项 :O(η),来源于迭代过程中w_t变化导致的梯度变化。
  3. 动量偏差项 :O(|β₁-β₂|),来源于c_t和m_t使用不同衰减系数带来的偏差。
  4. 当前梯度权重项 :O(1-β₁),来源于c_t中当前梯度g_t的权重。

这个引理告诉我们,只要学习率η足够小,迭代步数T足够大,并且β₁和β₂设置得接近,c_t就可以很好地逼近∇F(w_{t-1})。

4.2 收敛定理的证明策略

收敛性证明的核心思想是分析函数值F(w_t)的下降。我们根据算法在每一步的实际行为(是执行符号更新还是原始更新)进行分情况讨论。

情况一:当min_j |(c_t)_j| ≥ ν时(执行符号更新) 。 此时,更新为 w_t = w_{t-1} - η(sign(c_t) + λ w_{t-1})。利用L-光滑性,我们有: F(w_t) ≤ F(w_{t-1}) + <∇F(w_{t-1}), w_t - w_{t-1}> + (L/2)∥w_t - w_{t-1}∥² 将更新公式代入,并巧妙地将<∇F(w_{t-1}), sign(c_t)>拆分为: <∇F(w_{t-1}), sign(c_t) - sign(∇F(w_{t-1}))> + <∇F(w_{t-1}), sign(∇F(w_{t-1}))> 对于第一项,我们使用 引理6 (一个关于符号函数内积的数学不等式)将其控制为 2√d ∥c_t - ∇F(w_{t-1})∥。 对于第二项,一个关键观察是:对于一个向量x,有 <x, sign(x)> = ∥x∥₁(L1范数)。因此,这一项等于 -η ∥∇F(w_{t-1})∥₁。 结合正则化项和梯度变化项,经过推导可以得到: ∥∇F(w_{t-1})∥₁ ≤ [2(F(w_{t-1}) - F(w_t))] / η + 4√d ∥c_t - ∇F(w_{t-1})∥ + 4LηĜ²

情况二:当min_j |(c_t)_j| < ν时(执行原始更新) 。 此时,更新为 w_t = w_{t-1} - η(c_t + λ w_{t-1})。我们可以将其重写为一个近端梯度步的形式: w_t = argmin_w { <c_t, w - (1-ηλ)w_{t-1}> + (1/(2η)) ∥w - (1-ηλ)w_{t-1}∥² } 利用凸优化中的最优性条件(引理7),我们可以得到一个关于<c_t, w_{t-1} - w_t>的不等式。再结合L-光滑性,并注意到此时有ν₀∥c_t∥₁ ≤ ∥c_t∥²(其中0 < ν₀ ≤ ν),经过更复杂的推导,最终可以得到: ∥∇F(w_{t-1})∥₁ ≤ [8(F(w_{t-1}) - F(w_t))] / (η ν₀) + (4/ν₀) ∥∇F(w_{t-1}) - c_t∥² + √d ∥c_t - ∇F(w_{t-1})∥ + (一些关于T和ν₀的小量)

综合两种情况 ,对两个上界取期望并求和平均,然后代入引理9对E[∥c_t - ∇F(w_{t-1})∥]和E[∥c_t - ∇F(w_{t-1})∥²]的界。通过精心设置超参数:

  • 学习率 η = O(1/T^{3/4})
  • 动量系数 β₁ = 1 - O(1/√T), β₂ = 1 - O(1/√T),且 |β₁ - β₂| = O(1/√T)
  • 正则化系数 λ ≤ 1/(2ηĜT^α),取 α = 5/4
  • 阈值下界 ν₀ ≥ O(1/√d)

最终可以证明平均梯度L1范数的上界: (1/T) Σ_{t=1}^T E[∥∇F(w_t)∥₁] ≤ O(√d / T^{1/4})

这个结论表明,CLion算法可以收敛到一个平稳点,且收敛速率为O(1/T^{1/4})(按梯度L1范数衡量)。虽然这个速率对于一阶方法在非凸设置下是标准的,但关键是其泛化性得到了显著改善。

5. 工程实现与调参指南:从理论到实践

理论再完美,也需要落地到代码和实际训练中才有价值。CLion算法在标准Lion的基础上改动很小,但实现和调参上有一些需要特别注意的地方。

5.1 算法伪代码与实现细节

以下是CLion优化器的详细伪代码,我增加了详细的注释说明:

import torch

class CLion(torch.optim.Optimizer):
    def __init__(self, params, lr=1e-4, betas=(0.9, 0.99), weight_decay=0.0, nu=1.0):
        """
        CLion 优化器实现

        参数:
            params: 待优化的参数组
            lr: 学习率 (默认: 1e-4)
            betas: 用于计算梯度一阶矩和二阶矩的系数元组 (beta1, beta2)
                   beta1: 用于`c_t`(符号动量)的衰减率 (默认: 0.9)
                   beta2: 用于`m_t`(标准动量)的衰减率 (默认: 0.99)
            weight_decay: L2正则化系数 (λ) (默认: 0.0)
            nu: 自适应符号化阈值 (ν) (默认: 1.0)
        """
        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay, nu=nu)
        super().__init__(params, defaults)

    @torch.no_grad()
    def step(self, closure=None):
        """执行单次参数更新"""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            lr = group['lr']
            beta1, beta2 = group['betas']
            weight_decay = group['weight_decay']
            nu = group['nu']

            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad

                # 状态初始化
                state = self.state[p]
                if len(state) == 0:
                    # 指数移动平均 of gradient values (m_t)
                    state['exp_avg'] = torch.zeros_like(p)
                    # 用于符号化的动量 (c_t)
                    state['compressed_momentum'] = torch.zeros_like(p)

                exp_avg = state['exp_avg']
                c_t = state['compressed_momentum']

                # 1. 更新标准动量 m_t = β2 * m_{t-1} + (1-β2) * g_t
                exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)

                # 2. 更新压缩动量 c_t = β1 * m_{t-1} + (1-β1) * g_t
                # 注意:这里使用的是上一步的动量 m_{t-1},即当前的exp_avg还未更新时的值
                # 为了清晰,我们使用一个中间变量来代表旧的动量
                m_prev = exp_avg.clone().div_(beta2) if beta2 != 0 else torch.zeros_like(p) # 近似获取m_{t-1}
                # 更精确的做法是单独存储 m_{t-1},这里为简洁使用近似,实际实现建议存储两份动量
                c_t.mul_(beta1).add_(m_prev, alpha=1 - beta1).add_(grad, alpha=1 - beta1)

                # 3. CLion核心:自适应符号化
                # 判断当前c_t的所有非零分量是否都大于等于阈值nu
                # 首先获取非零元素的掩码
                non_zero_mask = (c_t != 0)
                if non_zero_mask.any():
                    # 计算所有非零元素的最小绝对值
                    min_abs = c_t[non_zero_mask].abs().min()
                    if min_abs >= nu:
                        # 情况A:所有非零|(c_t)_j| >= ν,进行符号化压缩
                        update_dir = c_t.sign()
                    else:
                        # 情况B:存在非零|(c_t)_j| < ν,使用原始值
                        update_dir = c_t
                else:
                    # 如果c_t全零,直接使用零(符号化后也是零)
                    update_dir = torch.zeros_like(c_t)

                # 4. 应用权重衰减(L2正则化)并更新参数
                # w_t = w_{t-1} - η * ( h(c_t) + λ * w_{t-1} )
                p.mul_(1 - lr * weight_decay)
                p.add_(update_dir, alpha=-lr)

        return loss

实现要点与坑点

  1. 动量存储 :上述简化代码中,为了获取 m_{t-1} ,我们进行了一个除以 beta2 的近似操作。这在 beta2 接近1时是合理的,但为了精确实现论文算法, 建议在state中同时存储 m_t c_t ,并在更新时严格遵循: c_t = beta1 * m_{t-1} + (1-beta1) * g_t ,然后再更新 m_t = beta2 * m_{t-1} + (1-beta2) * g_t
  2. 阈值判断 :判断条件 min_{j in S_t} |(c_t)_j| >= ν 中的 S_t 指的是 c_t 中所有非零分量的索引。在实现中,我们通过 non_zero_mask 来模拟。这里有一个边缘情况:如果 c_t 全为零,那么 min 操作在数学上无定义。但此时无论符号化还是用原值,更新方向都是零,所以我们将这种情况归为使用原值(即零)。
  3. 计算效率 :判断 min_abs >= nu 需要一次归约操作,但相比整个优化步骤的计算量可以忽略。符号化操作 c_t.sign() 是逐元素的,非常高效。
  4. 与标准Lion的兼容 :将 nu 设置为一个极大的值(如 float('inf') ),则条件 min_abs >= nu 永远不成立,算法退化为使用原始 c_t 更新,这类似于一个带特定动量结构的标准SGD。将 nu 设置为0,则条件永远成立(只要 c_t 非零),算法退化为标准Lion。因此,CLion是标准Lion的一个更灵活的泛化。

5.2 超参数设置建议

基于理论分析和实验经验,以下是CLion超参数的设置指南:

超参数 符号 理论建议 实践推荐 作用与说明
学习率 η O(1/√d) 或 O(1/√(dT)) 1e-4 到 1e-3 控制更新步长。CLion对学习率相对敏感,建议从标准Lion常用值(如3e-4)开始微调。
符号动量衰减率 β₁ 1 - O(1/√T) 0.9, 0.95, 0.99 控制 c_t 中历史信息的比重。越大,更新方向越平滑,但可能减缓对新梯度的响应。
标准动量衰减率 β₂ 1 - O(1/√T) 0.95, 0.99, 0.999 控制 m_t 的平滑程度。通常β₂ ≥ β₁,且两者接近,以减小偏差(见引理9)。
权重衰减 λ O(1) 0.0, 0.1, 0.01 L2正则化系数,帮助防止过拟合。在AdamW风格的解耦权重衰减中,它直接加在更新项中。
自适应阈值 ν ≥ 1 (理论),ν₀ ≥ O(1/√d) 0.1, 0.5, 1.0 核心参数 。控制何时进行符号化。值越小,越频繁使用符号更新(更激进压缩);值越大,越保守,更多使用原始值。

调参流程建议

  1. 固定基准 :首先,使用一组广泛适用的基准参数: lr=3e-4, betas=(0.95, 0.98), weight_decay=0.1, nu=0.5 。这通常在视觉和语言模型上提供一个不错的起点。
  2. 调整学习率 :如果训练不稳定(损失震荡或爆炸),尝试降低学习率(如1e-4)。如果收敛过慢,可尝试增大学习率(如1e-3),但需密切监控。
  3. 调整阈值ν :这是CLion特有的参数。
    • 如果追求极致的压缩和内存节省,可以尝试较小的ν(如0.1),但需警惕训练不稳定的风险。
    • 如果训练不稳定或任务非常敏感,可以增大ν(如1.0甚至2.0),这会让算法行为更接近使用原始动量的SGD。
    • 一个动态策略是:在训练初期使用较小的ν,让算法更激进地探索;在训练后期使用较大的ν,让算法更稳定地收敛。
  4. 协调β₁和β₂ :保持β₂略大于β₁(例如β₁=0.9,β₂=0.95或0.99)通常是个好选择。这确保了用于符号化的动量 c_t 比用于估计梯度方向的动量 m_t 更“短视”一点,可能对快速变化的梯度方向更敏感。
  5. 权重衰减 :权重衰减对于泛化至关重要。对于Transformer类模型,0.1是一个很强的正则项;对于CNN,0.01或0.001可能更合适。可以借鉴AdamW的设置。

5.3 常见问题与排查技巧

在实际使用CLion或分析相关问题时,你可能会遇到以下情况:

问题1:训练初期损失剧烈震荡或不下降。

  • 可能原因 :学习率过高,或阈值ν设置过小,导致在梯度估计不准确时进行了过于激进的符号化更新。
  • 排查步骤
    1. 绘制前几个迭代步中, c_t 的绝对值最小值(即 min_abs )的曲线。如果该值在早期频繁低于ν,说明算法频繁切换到原始更新模式,可能ν设大了或学习率有问题。
    2. 检查更新方向 update_dir 的范数。如果其范数异常大,可能是学习率过高或梯度爆炸。
    3. 尝试进行 梯度裁剪 (Gradient Clipping),例如将梯度范数裁剪到1.0。这可以立即缓解爆炸问题,并符合G-有界梯度的假设。
    4. 尝试**预热(Warmup)**策略,在训练初期使用较小的学习率,逐步增加到设定值。
    5. 增大ν值,让算法在初期更保守。

问题2:训练后期收敛速度变慢,或验证集性能停滞。

  • 可能原因 :学习率衰减策略不合适,或权重衰减过强抑制了有效学习。
  • 排查步骤
    1. 检查训练集损失是否还在下降。如果训练损失也停滞,可能是学习率太小。
    2. 尝试使用余弦退火(Cosine Annealing)或带重启的余弦退火(Cosine Annealing with Restarts)学习率调度器,它们常与CLion这类自适应算法配合良好。
    3. 观察 min_abs 的分布。在训练后期,梯度通常变小, min_abs 可能普遍低于ν,导致算法主要使用原始更新。此时可以尝试 动态调整ν ,随着训练进行缓慢降低ν,以维持一定的符号化比例,保持内存节省效益。
    4. 适当降低权重衰减系数。

问题3:与标准Lion或AdamW相比,内存节省效果不明显。

  • 可能原因 :ν设置过大,导致算法很少进入符号化模式。
  • 排查步骤
    1. 统计在整个训练过程中,执行符号化更新( min_abs >= nu )的迭代步比例。
    2. 如果该比例很低(例如<20%),尝试逐步减小ν,直到符号化比例达到一个你认为有意义的水平(例如50%-80%)。
    3. 注意,内存节省主要来自于用 sign(c_t) (int8或更低精度)替代 c_t (float32)进行更新计算和存储。确保你的实现确实在符号化步骤中使用了低精度存储。

问题4:在不同任务或模型架构上表现差异大。

  • 可能原因 :不同任务的损失地形(Loss Landscape)和梯度统计特性不同。
  • 排查步骤
    1. 视觉任务(CNN) :通常对学习率和权重衰减更敏感。可以从较小的学习率(1e-4)开始,使用适中的ν(0.5-1.0)。
    2. 语言任务(Transformer) :通常能承受较大的学习率。可以尝试标准Lion常用的3e-4,并设置较小的权重衰减(如0.1)。ν可以设置得小一些(0.1-0.5),因为Transformer的优化动态通常比较稳定。
    3. 小批量训练 :当批量大小(Batch Size)较小时,随机梯度噪声大。建议使用较大的β₁和β₂(如0.99)来平滑动量,并使用较大的ν来减少噪声导致的错误符号化。
    4. 始终在验证集上进行超参数扫描(Hyperparameter Sweep),找到最适合你具体任务和数据的配置。

CLion算法通过一个简单的自适应阈值,在标准Lion的激进压缩和传统优化器的稳定更新之间找到了一个可调节的平衡点。理论分析保证了其优异的泛化性能,而工程实现上的轻微改动使其易于集成到现有的训练管道中。对于关心模型最终性能、训练稳定性以及内存效率的从业者来说,CLion提供了一个值得尝试的新选择。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐