ICML 2026 的论文标题《ED:量化并优化神经网络的简单性,多项式表示》指向一个经常被忽视的问题:神经网络的表达能力不能只看参数量,真实复杂度隐藏在激活函数、权重分布和输入输出映射关系中。ED 方法的核心是借助多项式表示建立输入到输出的低维近似,再从近似系数中定义一个简单性指标,最后把指标变成可微正则项直接参与训练。这篇文章围绕这条技术主线展开:先解释为什么需要量化简单性,再给出一个基于 PyTorch 的最小可实现案例,之后补充运行验证、常见坑和生产落地建议。适合正在做模型可解释性、正则化、模型压缩或者论文复现的读者。
1. 为什么神经网络需要单独量化“简单性”
1.1 参数量与函数复杂度之间存在偏差
一个直接的直觉是参数少的网络更简单,参数多的网络更复杂。实际并不总是这样。两个参数量完全相同的 MLP,一个权重分布接近退化矩阵,实际只使用了少数几个隐藏神经元;另一个权重均匀非零,每个神经元都对输出有贡献。前者在参数数量上并没有优势,但表达出的函数可能更简单。
这种偏差带来的直接问题是:只看参数量无法判断模型是否容易被扰动、是否容易解释、是否适合部署到资源受限环境。ED 方法把“简单性”从“参数数量”中剥离出来,回到函数本身的表达方式上。这样做的价值在于,简单性和泛化能力、可解释性、剪枝可行性直接相关,而不只是一个理论概念。
另一个常见场景是模型压缩。用知识蒸馏或剪枝压缩模型时,经常遇到“压缩后精度没有下降,但部署后的行为变得不可解释”的情况。原因在于压缩只减少了存储和计算量,并没有改变网络内部的复杂耦合关系。如果从一开始就量化网络映射的复杂度,压缩的目标会更清晰。
1.2 简单性为什么与泛化、可解释性绑定
简单性之所以重要,是因为它和多个工程指标绑定。一个使用大量高次项去拟合数据的模型,在训练集边界处往往会出现剧烈震荡;一个只使用少数低次项的模型,在同样误差范围内通常更平滑。平滑意味着对输入扰动更不敏感,也就更容易在测试集上保持稳定。
可解释性方面,多项式表示有一个特殊优势:多项式系数是人类可以直接阅读的。系数中的某项足够大,说明输入和输出之间存在对应阶次的影响;某项趋近于零,说明这个阶次可以被忽略。如果把神经网络的输入输出关系投影到多项式基函数上,得到的系数分布就是一张“函数复杂度体检报告”。
所以 ED 方法不是追求“让模型变笨”,而是希望在不显著损失拟合能力的前提下,让模型用更少的表达资源完成任务。这一点与 L1 正则化、Dropout、权重衰减的思路不同,后三者主要约束参数本身,ED 约束的是函数表达空间中的有效维度。
1.3 ED 要解决的三个具体问题
把 ED 落地到工程时,至少需要解决三个问题:
- 可计算:简单性必须能在训练过程中实时计算,不能依赖人工观察网络结构。
- 可比较:不同架构、不同初始化的模型,必须能用同一个指标比较复杂度。
- 可优化:指标必须可微,能够作为正则项参与反向传播,在训练中持续降低。
这正是“多项式表示”发挥作用的地方。无论网络内部结构如何,最终都可以把输入输出关系映射到一组多项式基函数上,然后从系数分布中计算出有效维度。有效维度越低,函数表达越集中,也就是越简单。
2. ED 的核心机制:多项式表示与简单性度量
2.1 多项式表示如何近似神经网络
神经网络可以被看作一个复杂的函数映射器。理论上,只要基函数选择得当,任何连续函数都可以用一组基函数的线性组合去近似。多项式基是最经典的基函数之一,任何足够平滑的函数,都可以在局部用 Taylor 展开的思想表示为:
f(x) ≈ w0 + w1 * x + w2 * x^2 + ... + wp * x^p对神经网络来说,x 是输入,f(x) 是网络输出。ED 方法不做全局等价转换,而是用多项式基去近似网络在当前输入分布下的行为。这样得到的一组系数,就构成了网络局部行为的紧凑描述。
这里的“多项式表示”不是把整个网络替换成多项式,而是额外建立一个多项式投影。工程上可以有两种落地方式:
- 端到端可微方式:把网络输出层的一部分设计成多项式特征层,训练时直接优化系数分布。
- 拟合后评估方式:训练完成后,在网络输入上采样一批数据,用多项式回归拟合网络输出,再计算系数分布。
两种方式各有作用。前者适合把 ED 作为正则项;后者适合对已有模型做复杂度评估。论文的完整版本通常两者会结合,但理解核心机制可以从第一种出发。
2.2 从系数分布定义 ED 指标
拿到多项式系数后,如何定义“简单性”?一个自然的方法是看系数的集中程度。
假设求出的多项式系数是:
w = [w0, w1, w2, w3, w4]如果只有 w0 和 w1 明显非零,其余接近零,说明网络在这个输入范围内基本是线性行为,函数表达简单。如果所有系数都显著非零,说明网络使用了各个阶次的信息,函数表达复杂。
把系数做绝对值归一化,得到一组概率分布:
p_i = |w_i| / sum(|w_j|)然后计算这组分布的熵:
H = - sum(p_i * log(p_i))熵越小,系数越集中;熵越大,系数越分散。为了更直观,可以把它转换成有效维度:
ED = exp(H)有效维度的最小值是 1,表示所有表达压力集中在一个基函数上;最大值是多项式特征总数,表示每个基函数都被均匀使用。ED 越小,函数越简单。这个指标既可以在不同模型之间比较,也可以作为训练过程中的优化目标。
注意:ED 在不同论文中可能有不同全称,常见解读是 Effective Dimension,也可能被解释为 Entropic Diversity。无论采用哪种定义,核心都是通过系数集中度刻画函数复杂度。落地前需要以论文正式版本的公式为准。
2.3 为什么 ED 可以参与梯度优化
ED 由多项式系数计算得到,而多项式系数来自网络权重,因此 ED 对网络权重是可微的。把它加入 loss 后,梯度会同时优化拟合误差和表达复杂度:
loss = MSE + alpha * EDalpha 控制简单性在总目标中的权重。alpha 越大,模型越倾向使用更少的多项式基函数;alpha 越小,模型越优先拟合训练数据。
这里有一个容易被忽略的细节:直接最小化 ED 时,网络可能通过把所有权重逼近零来让熵趋向均匀分布,反而无法得到稀疏系数。因此实际实现中通常先对系数做绝对值归一化,再加一个很小的 epsilon 防止 log(0),必要时配合 L1 正则让系数真正稀疏化。
这个机制的另一种理解方式:ED 正则等价于在函数空间上施加“表达预算”。网络可以随便设计内部结构,但最终输入输出关系要尽量用少数多项式维度说清楚。
3. 环境准备与最小实验框架
3.1 依赖与版本建议
本文的示例代码基于 PyTorch,目标是 CPU 环境即可运行,不需要 GPU。建议使用以下依赖版本:
| 依赖 | 版本建议 | 用途 |
|---|---|---|
| Python | 3.9 或更高 | 运行环境 |
| PyTorch | 2.0 或更高 | 张量计算与自动微分 |
| NumPy | 1.24 或更高 | 数据处理 |
| Matplotlib | 3.7 或更高 | 绘制训练曲线 |
如果原始论文发布了官方仓库,建议优先在官方仓库提供的虚拟环境中运行,再回到这里理解每一步的作用。以下示例代码用于说明 ED 的核心流程,实际项目要结合自己的包名、数据格式和版本调整。
3.2 项目结构
建议按下面的目录组织代码:
ed_demo/ ├── train_poly.py # 训练入口 ├── ed_utils.py # 多项式基和 ED 指标 ├── model.py # 模型定义 └── requirements.txt这个结构足够简单,也能方便后面扩展成更完整的论文复现工程。
3.3 构造合成回归数据
为了快速验证 ED 的效果,使用一个带噪声的 sin 函数作为目标:
import torch torch.manual_seed(42) x = torch.linspace(-1, 1, 200).reshape(-1, 1) y = torch.sin(3 * x) + 0.02 * torch.randn_like(x) indices = torch.randperm(x.shape[0]) train_idx, test_idx = indices[:160], indices[160:] x_train, y_train = x[train_idx], y[train_idx] x_test, y_test = x[test_idx], y[test_idx]输入归一化到 [-1, 1] 不是可选项。多项式的高次幂对输入范围非常敏感,如果输入范围是 [0, 1000],x 的 10 次方很容易超过浮点数表示范围。这也是后面排查 NaN 时的第一检查点。
4. PyTorch 实现:多项式基、ED 计算与正则训练
4.1 多项式特征层
先实现一个多项式特征构造函数。输入形状是(batch_size, in_features),输出是所有阶次的拼接结果:
import torch import torch.nn as nn def polynomial_basis(x, degree=3): # x 形状: (B, D) parts = [torch.ones_like(x[:, :1])] for d in range(1, degree + 1): parts.append(torch.pow(x, d)) return torch.cat(parts, dim=-1) # 形状: (B, 1 + D * degree)这里用单项式基做演示,工程上更推荐 Legendre 正交基。单项式基的优点是实现简单,缺点是高阶项与低阶项之间高度相关,容易造成数值不稳定。正交基虽然在代码上多几行,但在输入范围较大或者阶数较高时稳定很多。
4.2 ED 指标计算
定义有效维度指标。输入是多项式系数层权重,输出是一个标量 tensor:
def effective_dimension(weight, eps=1e-8): w = torch.abs(weight.squeeze()) total = w.sum() + eps p = w / total entropy = -(p * torch.log(p + eps)).sum() return torch.exp(entropy)这里有几处设计意图:
- 先取绝对值再归一化,避免正负系数相互抵消。
- 加上
eps是为了防止log(0),但eps不能太大,否则会压低真实熵的区分度。 - 使用
torch.exp(entropy)把熵转换成可解释的有效维度。如果系数完全集中在一个维度上,结果接近 1;如果均匀分布在 N 个维度上,结果接近 N。
如果只需要做复杂度评估,不需要反向传播,也可以直接输出熵值。两种方式等价。
4.3 完整模型与训练闭环
这里使用最简单的方式:模型本身就是一个多项式线性层。虽然这不是完整神经网络,但足够演示 ED 的计算和梯度传递:
class PolyNet(nn.Module): def __init__(self, in_features=1, degree=3): super().__init__() self.degree = degree self.poly_dim = 1 + in_features * degree self.linear = nn.Linear(self.poly_dim, 1, bias=False) def forward(self, x): phi = polynomial_basis(x, self.degree) out = self.linear(phi) return out, phi训练循环需要同时计算 MSE 和 ED:
def train(model, x_train, y_train, epochs=500, lr=1e-2, alpha=0.01): opt = torch.optim.Adam(model.parameters(), lr=lr) mse_loss = nn.MSELoss() for epoch in range(epochs): model.train() opt.zero_grad() pred, phi = model(x_train) loss_mse = mse_loss(pred, y_train) ed = effective_dimension(model.linear.weight) loss = loss_mse + alpha * ed loss.backward() opt.step() if (epoch + 1) % 100 == 0: print( f"epoch {epoch + 1:3d} | " f"mse {loss_mse.item():.6f} | " f"ed {ed.item():.4f} | " f"loss {loss.item():.6f}" )要同时跑一组基线对比。alpha 设为 0 时,模型只优化 MSE;alpha 设为 0.01 时,模型额外优化简单性:
model_base = PolyNet(in_features=1, degree=3) model_ed = PolyNet(in_features=1, degree=3) train(model_base, x_train, y_train, alpha=0.0) print("---") train(model_ed, x_train, y_train, alpha=0.01)这里需要说明两点。第一,由于这个演示模型本身只有一层多项式层,ED 直接作用于唯一的权重层,效果最明显。换成有多层网络的场景时,可以把多项式特征层作为输出端的一部分,或者对网络输出做局部拟合。第二,训练时每次迭代都重新计算 ED,这会增加少量计算开销,但通常可以接受。
4.4 关键参数含义与推荐范围
| 参数 | 含义 | 示例值 | 影响 |
|---|---|---|---|
| degree | 多项式最高阶数 | 3 | 越大表达能力越强,但更容易过拟合 |
| alpha | ED 正则权重 | 0.01 | 越大模型越简单,但可能欠拟合 |
| lr | 学习率 | 1e-2 | 影响收敛速度和稳定性 |
| epochs | 训练轮数 | 500 | 需要足够让 ED 正则生效 |
alpha 的选择是 ED 方法里最影响结果的操作。alpha 太小,简单性约束可以忽略;alpha 太大,模型会优先压低有效维度,测试误差可能反而上升。推荐先在验证集上做一个 0、0.001、0.01、0.1、1 的稀疏网格搜索,观察 MSE 和 ED 的平衡点。
5. 运行验证:看三条可量化证据
5.1 训练曲线对比
运行训练脚本后,预期看到类似下面的输出:
epoch 100 | mse 0.130872 | ed 2.8431 | loss 0.159301 epoch 200 | mse 0.040251 | ed 2.1120 | loss 0.061371 epoch 300 | mse 0.026980 | ed 1.9033 | loss 0.046013具体数值会随随机种子变化,但趋势应该是稳定的:ED 从初始接近均匀分布的值逐渐下降到较低值,MSE 同步下降。如果 ED 不下降,通常说明 alpha 设置过小,或者系数归一化逻辑写错了。
5.2 多项式系数稀疏度
训练完成后,直接检查模型权重:
for name, param in model_ed.named_parameters(): print(name, param.data)可以看到系数分布。以 degree=3、输入维度为 1 为例,模型有 4 个系数。alpha=0 的模型,系数通常都比较分散,每个项都被使用;alpha=0.01 的模型,系数会更集中在一次项或二次项附近,高次项可能接近零。
不要只看系数是否非零,更合理的检查是打印归一化比例:
w = torch.abs(model_ed.linear.weight.squeeze()).detach() p = w / w.sum() print(p.numpy())如果某个维度的比例超过 0.7,说明模型确实在用少数基函数表达函数,ED 正则生效了。
5.3 测试集误差与有效维度
训练完成后,在测试集上评估:
model_ed.eval() with torch.no_grad(): pred_test, _ = model_ed(x_test) test_mse = nn.MSELoss()(pred_test, y_test).item() ed_final = effective_dimension(model_ed.linear.weight).item() print(f"test mse: {test_mse:.6f}") print(f"ed: {ed_final:.4f}")做对比实验时,应该同时记录三组数据:baseline 的测试 MSE、ED 的测试 MSE、两者的最终有效维度。常见结果是 alpha=0.01 的模型测试误差与 baseline 相近或略低,但 ED 明显更小。如果测试误差明显变差,说明 alpha 过大,需要调小。
6. 常见问题排查
6.1 训练早期出现 NaN
| 现象 | 可能原因 | 检查方式 | 处理建议 |
|---|---|---|---|
| loss 或 ED 出现 NaN | 输入范围过大导致高次幂溢出 | 打印 x 的 min/max | 归一化到 [-1, 1];改用 Legendre 基 |
| loss 出现 NaN 但 ED 正常 | 学习率过大导致梯度爆炸 | 打印每层梯度范数 | 降低 lr;加梯度裁剪 |
| 多项式回归解报 Singular | 特征矩阵不可逆 | 检查特征维度是否重复 | 增加 ridge 系数 |
单项式基最容易出的问题就是高次幂。x=0.5 时,x 的 10 次方还能接受;x=10 时,x 的 10 次方已经超出常规浮点精度。先把输入归一化范围定下来,再选择多项式阶数。
6.2 ED 保持常数不下降
如果训练过程中 ED 一直维持在高位,可以从三个方向排查:
- alpha 是否太小,MSE 梯度完全主导了优化方向。
- 权重是否在零附近,导致归一化后接近均匀分布。
- 是否需要先预训练一段时间,再开启 ED 正则。
这里的深层原因是:系数越接近零,绝对值归一化后各个维度的比例差别越小,熵越接近最大值,梯度也越微弱。这不是 ED 失效,而是归一化方式带来的优化困难。解决办法是配合 L1 正则使用,或者启用 ED 前先跑几百轮普通训练。
6.3 高阶多项式导致严重过拟合
把 degree 从 3 调到 8 后,测试误差可能不降反升。原因是多项式阶数越高,模型拟合边界的灵活性越强,也越容易把噪声特征当成真实模式。
建议不要盲目提高 degree。先用 degree=3 跑通,再逐步增加,并在验证集上观察误差。如果加入 ED 正则后高阶系数明显下降,说明 ED 正在发挥作用;如果高阶系数仍然很大,说明 alpha 不够。
6.4 复现结果与论文不一致
论文中的网络结构、多项式基函数、归一化方式、alpha 调度策略都会影响结果。单项式基和正交基在数值表现上差异很大,单纯对比 loss 曲线没有意义。
复现时优先确认论文里 ED 的精确定义。若使用 Effective Dimension,需要确认是 exp(entropy) 还是别的归一化版本;若使用 Entropic Diversity,计算公式可能完全不同。先把指标定义对齐,再谈结果对比。
7. 从复现到生产:建议与扩展
7.1 论文复现与工程落地的差异
论文复现可以只关心指标趋势,工程落地则需要考虑更多约束。ED 作为复杂度指标,进入生产环境前需要做几件事:
- 数据归一化策略要固化,不能训练时和预测时不一致。
- 多项式最高阶数要设上限,否则线上数据分布一旦偏移,高次项会迅速放大异常。
- ED 指标需要记录到监控系统里,与 MSE 一起观察。如果模型上线后 ED 突然升高,说明输入分布漂移导致网络使用了更多表达维度。
- 模型回滚判断不能只看精度,还要看 ED。一个 ED 显著升高但精度变化不大的模型,更可能是记住噪声而不是学到规律。
7.2 可复用清单
落地 ED 相关功能时,可以复用以下检查清单:
- 输入数据是否已经归一化到合理范围?
- 多项式基是否选择了数值稳定的版本?
- ED 定义是否与目标论文一致,是 exp(entropy) 还是原始熵?
- 是否包含 alpha=0 的基线对比?
- 是否在训练中同时记录 MSE 和 ED?
- 是否在测试集上验证,而不是只观察训练 loss?
- 多项式特征是否做了缓存,避免重复计算?
- 是否设置多项式阶数上限?
- 是否在监控中加入了 ED 指标?
7.3 扩展方向
ED 方法的思路可以延伸到多个方向。第一个方向是模型选择:训练多个候选模型时,除了比较测试精度,还可以比较 ED,选择精度相近但表达更简单的模型。第二个方向是知识蒸馏:教师模型的 ED 可以作为蒸馏目标的一部分,指导学生模型不仅学输出,还学输出函数的复杂度结构。第三个方向是局部可解释性:把全局多项式投影改成局部多项式拟合,可以对输入区间分别评估复杂度,定位模型在哪个区域行为最不稳定。
对刚接触这个方向的新手,建议先把这个最小多项式层跑通,再把多项式特征层接到一个两层的 MLP 后面,观察 ED 对隐藏层的影响。对比不同 alpha 下的系数分布,比直接读论文更容易理解“用多项式表示量化简单性”到底解决什么问题。