训练神经网络时,我们常把优化想象成“小球沿山坡下降”。损失曲线能告诉我们下降有多快,却省略了参数如何移动:一段近乎水平的曲线,既可能对应很小的参数更新,也可能对应沿低损失变化方向的持续移动。损失地形(loss landscape)可视化尝试把这些空间信息补回来。
困难在于,网络参数通常处于高维空间。把轨迹压缩到二维之后,如何给周围的位置赋予有意义的损失值?本文介绍 Ziming Liu 在 Seeing sticky plateau 中提出的局部平面融合思路:用全局坐标展示轨迹形状,用沿轨迹构造的局部平面采样损失,再把这些局部视图融合成一张地图。 下文公式对应本文的 NumPy 实现,交互示例用于辅助理解这一构造。
1. 为什么仅有 PCA 还不够
设网络参数为 \(\theta\in\mathbb{R}^D\),训练产生轨迹 \(\theta_0,\ldots,\theta_T\)。在固定数据集上,损失函数 \(L(\theta)\) 为每个参数位置赋予一个标量。损失曲线只保留序列 \(L(\theta_t)\);我们希望同时看到轨迹的位置及其附近的损失结构。
一种直接的办法是对参数轨迹做主成分分析(PCA)。记轨迹均值为 \(\bar\theta\),前两个单位主方向组成矩阵 \(P\in\mathbb{R}^{D\times2}\),满足 \(P^\top P=I\)。二维坐标为
\[ z_t=P^\top(\theta_t-\bar\theta). \]
这两个方向保留了轨迹中尽可能多的参数变化方差。若要绘制地形,可以在二维网格上取点 \(z\),将其映射回全局 PCA 平面,再计算损失:
\[ q(z)=\bar\theta+Pz,\qquad L_{\mathrm{PCA}}(z)=L(q(z)). \]
问题出在轨迹通常并不完全位于这个平面内。把真实参数分解为
\[ \theta_t=\bar\theta+Pz_t+r_t, \]
其中 \(r_t\) 是 PCA 丢弃的残差,就会发现 \(L_{\mathrm{PCA}}(z_t)\) 实际计算的是 \(L(\theta_t-r_t)\)。即使残差很小,只要损失对这个方向敏感,它与真实训练损失 \(L(\theta_t)\) 仍可能明显不同。于是,投影后的轨迹可能看起来没有沿着图中的谷底移动。
轨迹投影的保真度与损失的保真度是两个问题。 PCA 优化的是参数空间中的重建误差,没有使用损失函数的信息。较高的解释方差只能说明二维坐标保留了较多轨迹变化,不能保证该平面准确反映轨迹附近的损失。
2. 沿轨迹构造局部平面
局部平面融合仍使用 PCA 作为整张图的坐标系统,但在计算损失时引入真实轨迹附近的切片。每个切片以一个真实参数 \(\theta_t\) 为中心,其方向由近期运动决定。
用前进方向和转弯方向确定切片
选定回看步数 \(k\),定义有限差分
\[ v_t=\theta_t-\theta_{t-k},\qquad a_t=(\theta_t-\theta_{t-k})-(\theta_{t-k}-\theta_{t-2k}). \]
\(v_t\) 表示近期位移,\(a_t\) 表示前后两段位移的变化。这里沿用“速度”和“加速度”的称呼;构造平面只需要方向,因此省略了时间间隔的比例因子。
先把速度归一化,再取加速度中与速度正交的部分:
\[ e_{t,1}=\frac{v_t}{\|v_t\|},\qquad b_t=a_t-(a_t^\top e_{t,1})e_{t,1},\qquad e_{t,2}=\frac{b_t}{\|b_t\|}. \]
令 \(E_t=[e_{t,1},e_{t,2}]\),得到局部仿射平面
\[ \Pi_t=\{\theta_t+E_tu:\ u\in\mathbb{R}^2\}. \]
第一个方向沿近期运动,第二个方向捕捉轨迹的弯曲。在轨迹足够光滑、采样间隔足够小时,可以把它理解为对局部运动平面的近似。它选择的是与这次训练有关的方向,并不一定包含损失变化最剧烈的方向,也不是由 Hessian 特征向量定义的曲率切片。
\(k\) 决定观察运动的时间尺度:较小的 \(k\) 更贴近短期变化,较大的 \(k\) 跨越更长的轨迹区间。当位移接近零,或两段位移近乎共线时,无法稳定确定两个独立方向,本文实现会跳过该平面。局部平面只从 \(t\geq2k\) 开始采样,因此最初一段轨迹没有以自身为中心的切片。
把同一个地图位置送到不同切片上
为了让多个切片共享一张二维地图,对任意网格点 \(z\),先得到全局平面上的高维点 \(q(z)\),再将它正交投影到各个局部平面:
\[ \widetilde\theta_t(z) =\theta_t+E_tE_t^\top\bigl(q(z)-\theta_t\bigr). \]
随后用完整网络和固定数据集计算
\[ \ell_t(z)=L\bigl(\widetilde\theta_t(z)\bigr). \]
因此,同一个二维位置 \(z\) 在不同局部平面上对应不同的高维参数。每个 \(\ell_t\) 都提供一份围绕该段轨迹的损失视图,损失值通过网络前向计算获得,没有使用局部二次近似。
这里有一个重要细节:虽然 \(\Pi_t\) 穿过真实参数 \(\theta_t\),但在地图坐标 \(z_t\) 处,投影结果为
\[ \widetilde\theta_t(z_t)=\theta_t-E_tE_t^\top r_t. \]
它仍可能偏离 \(\theta_t\)。局部切片让采样参考真实轨迹的位置和运动方向;这种构造并不保证在轨迹坐标处精确复现训练损失。
3. 将局部视图融合成地形
令 \(\mathcal S\) 为采样得到的有效局部平面索引集合。对地图位置 \(z\),按它到各切片中心的二维距离分配高斯权重:
\[ w_t(z)=\frac{\exp\left(-\|z-z_t\|^2/(2\sigma^2)\right)} {\sum_{s\in\mathcal S}\exp\left(-\|z-z_s\|^2/(2\sigma^2)\right)}, \qquad t\in\mathcal S. \]
权重之和为 1。离查询位置近的轨迹片段贡献更大;\(\sigma\) 较小时,地图更强调最近的局部视图,较大时则混合更宽范围的切片。距离在 PCA 坐标中计算,因此高维中相距很远、投影后相近的轨迹片段也可能被混合。
本文实现先对每份局部损失设置上限 \(C\),再求加权平均:
\[ \widehat L(z)=\sum_{t\in\mathcal S}w_t(z)\min\{\ell_t(z),C\}. \]
截断能防止极高损失掩盖低损失区的细节,也会抹去高处的差异。先截断再平均与最后对融合结果截断不同;前者会直接改变各局部视图的贡献。
整个方法可以概括为:
- 记录训练参数,用 PCA 建立全局二维坐标。
- 沿轨迹采样中心,用位移及其变化构造局部平面。
- 将二维网格点映回高维空间,再分别投影到局部平面并计算损失。
- 按二维距离融合局部损失,绘制二维热力图或三维曲面,并叠加投影轨迹。
得到的 \(\widehat L\) 是一张依赖训练轨迹的融合地图。它既不是完整高维地形的重建,也通常不对应某个固定参数平面上的损失函数。它的价值在于,让全局轨迹和沿途局部损失结构出现在同一幅图中。
4. 示例:一维函数拟合
网络结构为 [1, 10, 1],隐藏层使用 SiLU 激活函数,包含 31 个可训练参数。输入为从 \([-1, 1]\) 均匀采样的 1,000 个点,标签由所选目标函数生成。Adam 优化器最小化这些样本上的均方误差(MSE):
\[ L(\theta) = \frac{1}{N}\sum_{i=1}^{N}\left(f_\theta(x_i)-y_i\right)^2. \]
这个小模型已经需要将 31 维轨迹压缩到二维,足以展示上述方法的基本取舍。默认拟合 \(y=x\);其余目标函数用于观察不同回归任务下的轨迹形状。所有地形点都使用与训练相同的样本计算 MSE。
地形的尺度与采样
下列参数只改变地图的构造,复用同一条训练轨迹。\(k\) 控制局部方向的时间尺度,\(\sigma\) 控制切片融合的空间尺度,\(C\) 控制保留的损失范围。采样间隔决定使用多少个局部平面,网格精度决定二维地图的采样密度;提高后两者的采样密度无法消除降维本身的信息损失。
默认使用 \(k=10\)、\(\sigma=0.1\)、\(C=1\),每 20 步取一个局部平面,在轨迹 PCA 包围框向各侧扩展 20% 的区域上计算 \(24\times24\) 网格。
将损失曲线与空间轨迹并读
损失曲线展示真实的 \(L(\theta_t)\),地形图展示融合后的 \(\widehat L(z)\),叠加轨迹的位置为 \(z_t\)。二维热力图便于观察轨迹相对谷地的走向,三维曲面提供更直观的高度感。时间条把两张图中的同一训练时刻联系起来。
示例由 quarto-marimo 与 Pyodide 在浏览器中计算。初始图表为预计算结果;联网加载运行时后,可提交训练或地形参数。三维视图需要 WebGL。
5. 这张地图帮助我们看见什么
平台期中的参数可能仍在移动
原作者以“sticky plateau”为题,强调这类可视化对观察损失平台期的帮助,并将其机制讨论留给后续文章。这里也把平台期视为待解释的现象:先从真实损失曲线找到下降缓慢的区间,再检查同一时间段的空间轨迹。
若投影轨迹仍明显移动,说明损失变化小的同时,参数可以发生可见的位移;若它在图中也几乎停住,还需要检查完整参数位移,因为运动可能发生在 PCA 丢弃的方向。这样,地形图把“训练变慢了”细化为关于运动方向和空间位置的问题。
在足够小的更新下,损失变化近似满足
\[ L(\theta_{t+1})-L(\theta_t) \approx\nabla L(\theta_t)^\top(\theta_{t+1}-\theta_t). \]
这个内积很小,既可能来自较小的梯度或位移,也可能来自二者接近正交。仅凭损失平台无法区分这些情形;融合地图可以提供线索,机制判断仍需要梯度、完整参数位移等证据。示例使用 Adam,更新还受动量和逐坐标缩放影响,因此“小球沿最陡坡下滚”的直觉也有边界。
转弯与谷地提供局部几何线索
如果一段轨迹沿狭长的低值区域前进,可以进一步询问:沿运动方向损失变化是否较小,横向是否更陡?如果损失下降速率在转弯附近发生变化,可以定位这段轨迹,再检查相应局部切片。地图在这里承担的是发现结构、提出问题和定位训练阶段的作用。
这些观察需要检查对可视化参数的敏感性。在固定训练结果下改变 \(\sigma\),或加密网格与局部平面采样,若某个谷地或脊线随之显著变化,就应谨慎赋予它优化意义。远离轨迹的区域尤其缺少局部依据:归一化权重始终会给出一个值,即使所有切片中心都很远。
地图高度不能替代真实损失
融合同时涉及投影、截断和加权,通常有 \(\widehat L(z_t)\ne L(\theta_t)\)。三维小球的高度取自融合网格的双线性插值,并带有便于显示的偏移;判断训练进展应使用旁边的真实 MSE 曲线。
此外,融合图中的坡度包含权重随位置变化的影响。忽略截断拐点,记 \(\ell_t^C=\min(\ell_t,C)\),则
\[ \nabla_z\widehat L =\sum_t w_t\nabla_z\ell_t^C +\sum_t\ell_t^C\nabla_z w_t. \]
第二项来自切片之间的混合。因此,融合后的谷地与坡度不能直接解释为原网络的梯度场,图中看似平坦的区域也不能证明高维空间中的平坦极小值、良好泛化或优化器收敛。
每次训练都会重新决定 PCA 坐标与局部切片。比较不同学习率或初始化时,地图也随之改变;相同的 \(\sigma\) 数值可能对应不同的相对平滑尺度。跨实验比较仍应以真实损失和明确的参数空间统计量为依据。
小结
局部平面融合把高维优化过程拆成两个层次:PCA 描述整体轨迹,沿途的运动平面提供局部损失视图,再由距离权重将它们组织到同一张地图中。它补充了损失曲线缺少的空间线索,适合探索平台期、转弯和谷地附近的训练行为。理解其采样与融合过程,也就明确了这张地图的用途:帮助形成和检查几何直觉,并为进一步的定量分析定位问题。
参考
- Ziming Liu. Loss landscape visualization 1 — Seeing sticky plateau. 2026. 本文局部平面融合方法的来源;示例为 NumPy 实现。