跳转至

layerwise progressive freezing a training scaffold for depth-scalable binary net

会议: NeurIPS 2026
arXiv: 2606.27759
领域: 模型压缩
关键词: 二值神经网络, 直通估计器, 逐层渐进冻结, 梯度阻断, 随机掩码

一句话总结

针对全二值神经网络(BNN)因直通估计器(STE)前后向不匹配导致网络加深时性能严重崩溃的痛点,提出了从输入到输出逐层硬化二值参数与激活值的 StoMPP 渐进冻结框架,彻底消除了激活值引发的梯度阻断并实现了可深度扩展的二值网络训练。

研究背景与动机

深度神经网络在算力受限边缘设备部署时,其巨大的内存和计算开销成为核心瓶颈。二值神经网络(BNN)将权重甚至激活值同时极化约束到 \(\{-1, +1\}\),能够将推理代价和内存消耗降低一个数量级以上。然而,从头开始训练全二值神经网络一直受困于其前向硬符号函数 \(\mathrm{sign}(\cdot)\) 不可微的数学本性。当前领域的主流解法是直通估计器(Straight-Through Estimator, STE),即在前向计算中应用硬二值符号映射,在反向传播中代入恒等或平滑的替代梯度。然而,STE 固有的前后向不匹配(forward/backward mismatch)使得模型随着网络深度的增加出现极其严重的性能衰退,常规全二值网络难以向深层 ResNet 或现代复杂网络架构扩展。

学界此前应对深度崩溃的思路大多聚焦于“改良替代梯度本身”,例如优化符号函数的平滑逼近形式、引入梯度整形或设计专用优化器,但这始终未跳出“全网络自始至终承受前后向失配”的范式。渐进式量化(Progressive Quantization)虽然在权重二值化(BWN)中被证明能平滑参数离散化过程,但在全二值网络(BNN)中却遭遇了无法逾越的屏障:如果直接在全网随机选取激活值或权重进行渐进冻结,网络在深度增加时依然会彻底崩溃或停滞。

导致这一困境的核心矛盾在于激活值诱导的梯度阻断(Activation-Induced Gradient Blockades)。在全二值网络中,一旦某个激活值被冻结为离散硬二值符号,其真实导数几乎处处为零;若在全网全局随机硬化激活值,被冻结的神经元将化作一道道阻断墙,彻底切断下游损失流向前序浅层的有效梯度路径。权重二值网络(BWN)由于激活值始终保持连续而天然免疫此问题,全 BNN 则脆弱不堪。本文的切入角度是改变训练中施加二值化的时空顺序,而非在全局失配的泥潭里微调替代梯度。核心 idea:提出自输入端向输出端逐层推进的随机掩码渐进二值化框架 StoMPP,让当前正在过渡的层始终拥有未冻结且梯度畅通的下游后缀网络,在机理上阻断梯度隔绝并支持纯无 STE 训练。

方法详解

整体框架

StoMPP(Stochastic Masked Partial Progressive Binarization)维持权重与激活值的实数连续底态变量,并通过动态二值掩码调度每个元素处于连续代理状态还是离散冻结状态。整个网络在训练进程中沿前向网络深度方向划分为三个互斥且滑移的区域:完全离散冻结的前缀层、正在通过局部随机掩码与软重置逐渐提高二值比例的过渡层,以及保持完全可微连续状态的后缀层。

由于二值化沿着“输入 \(\to\) 输出”单向传播,正处于过渡状态的层始终通过连续的下游后缀直接连通最终损失函数,杜绝了前向路径被中途冻结激活值切断的危险。在过渡层内部,StoMPP 采用三次多项式单调提高二值化占比,并配合软重置机制动态更新冻结子集,防止早期参数过早僵化陷入次优局部极值。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据与浅层特征"] --> B["前缀网络:已冻结二值层<br/>权重与激活值全硬化/无梯度回传需求"]
    B --> C["过渡网络:当前调度层<br/>混合连续代理与离散符号"]
    C --> D["单层软重置与三次步进调度<br/>Bernoulli局部重采样/避免局部极值"]
    D --> E["后缀网络:未冻结连续层<br/>保持Clip激活/构建通畅梯度回传通路"]
    E --> F["任务损失与反向传播更新"]

关键设计

1. 混合前向映射与双模式反向梯度解耦:确立无 STE 独立训练与兼容基准 为了无缝衔接连续优化与硬二值约束,StoMPP 对权重或预激活值变量 \(u\) 统一定义了连续平滑代理 \(\mathrm{Smooth}(u)\) 与硬离散映射 \(\mathrm{sign}(u)\)。激活值采用硬截断连续函数 \(\mathrm{clip}(u) = \max(-1, \min(1, u))\),权重采用恒等映射。给定形状相同的二值掩码 \(M \in \{0, 1\}\),前向计算表述为两者的凸组合: $\(u' = M \odot \mathrm{sign}(u) + (1 - M) \odot \mathrm{Smooth}(u)\)$ 在反向传播中,StoMPP 将穿透已冻结区域(\(M=1\))的梯度形式解耦为一个独立设计空间 \(\mathrm{Surr}(u)\): $\(\nabla_{u} = \nabla_{u'} \odot \left[ (1 - M) \odot \frac{\partial \mathrm{Smooth}(u)}{\partial u} + M \odot \mathrm{Surr}(u) \right]\)$ 当令 \(\mathrm{Surr}(u) = 0\) 时,模型构成了纯粹的“无 STE(STE-free)”训练范式,所有已硬化二值单元不接收任何人为伪造的替代梯度,学习信号完全依赖未冻结分量;当令 \(\mathrm{Surr}(u) = 1\) 时,则与经典 STE 替代梯度复合(StoMPP+STE)。这种设计彻底剥离了渐进调度与替代梯度估算,使得机制本身的消融与度量不再受 STE 超参数的纠缠。

2. 单层软重置与三次步进调度:平衡探索稳定性与离散收敛 在处于过渡阶段的单层内,直接将部分参数永久固化会导致不可逆的累积量化误差。StoMPP 引入了带软重置(Soft Refresh)的随机掩码演化机制。对于包含 \(n\) 个参数的层,设定过渡进程时间步 \(\tau \in \{1, \dots, T\}\),目标冻结比例按三次多项式平滑上升: $\(p(\tau) = \left( \frac{\tau}{T} \right)^3\)$ 在每个训练步,算法并不重新生成整张掩码,而是仅仅抽取 \(k = \lfloor n/r \rfloor\) 个随机索引(刷新率基准设为 \(r=100\),即单步仅扰动约 1% 的位置),仅对这部分位置按 \(\mathrm{Bernoulli}(p(\tau))\) 独立重采样,其余 99% 的掩码保持原状。这种保守更新不仅维持了整体冻结比例在期望上严格符合 \(p(\tau)\),还为网络提供了高度的时间连贯性,使未冻结参数有充分时间适应当前的二值拓扑,避免了全量重抽样造成的剧烈震荡。

3. 正向逐层调度杜绝梯度阻断:破解全二值深度崩溃的核心机理 在全二值网络中,若采用全局无序渐进冻结,一旦深层或中间层的某组激活值变为硬 \(\mathrm{sign}(z)\),由于 \(\partial \mathrm{sign}(z)/\partial z = 0\) 处处成立,通往浅层的反向求导链将直接被多处阻断断绝。StoMPP 采用从输入到输出(input-to-output)的严格单向逐层硬化流水线:当网络前进至第 \(\ell\) 层时,所有前序层 \(1 \le i < \ell\) 虽已被完全冻结为二值且不可导,但这些前置层早已完成收敛、不再需要参数更新;而后续层 \(\ell < i \le N\) 全体处于未冻结的连续代理状态(激活值导数在 \((-1, 1)\) 内为 1),为当前第 \(\ell\) 层提供了从损失函数直达该层的完全畅通的导数梯度高速公路。该设计从根本上规避了激活硬化引发的梯度截断,使得网络深度的增加不再等价于梯度消失的灾难。

损失函数 / 训练策略

训练采用标准交叉熵损失函数。为了纯粹验证量化规则本身的演化效能,排除了权重衰减(Weight Decay,因为 \(L_2\) 正则化会拉扯二值权重远离 \(\pm 1\)),并在主对比中使用固定学习率(SGD 带 Nesterov 动量,学习率 \(\eta=0.1\),动量 0.9)。每个待量化层分配均等训练轮数,首尾层及下采样投影层按常规惯例保持全精度,其余层全部施加逐层渐进硬化。

实验关键数据

主实验

在 CIFAR-10、CIFAR-100、ImageNet(ILSVRC2012)、SST-2(BERT-base)及 MobileNetV2 上进行了全面评测。在全二值网络(BNN,权重和激活值均为 1-bit)设定下,随着网络深度增加(ResNet-18 \(\to\) ResNet-34 \(\to\) ResNet-50),传统 STE 发生灾难性断崖下跌,而 StoMPP 展现出显著的深度可扩展性优势。

数据集 模型架构 量化配置 评估指标 Vanilla STE StoMPP (STE-free) StoMPP + STE 复合提升 (vs STE)
CIFAR-10 ResNet-18 BNN (1W/1A) Top-1 准确率 (%) 77.8 80.9 86.1 +8.3
CIFAR-10 ResNet-34 BNN (1W/1A) Top-1 准确率 (%) 61.5 76.0 82.1 +20.6
CIFAR-10 ResNet-50 BNN (1W/1A) Top-1 准确率 (%) 51.5 69.5 78.6 +27.1
CIFAR-100 ResNet-18 BNN (1W/1A) Top-1 准确率 (%) 49.1 53.8 58.0 +8.9
CIFAR-100 ResNet-34 BNN (1W/1A) Top-1 准确率 (%) 33.7 39.8 47.8 +14.1
CIFAR-100 ResNet-50 BNN (1W/1A) Top-1 准确率 (%) 26.7 40.2 46.5 +19.8
ImageNet ResNet-18 BNN (1W/1A) Top-1 准确率 (%) 41.9 39.5 45.5 +3.6
ImageNet ResNet-34 BNN (1W/1A) Top-1 准确率 (%) 37.0 39.7 48.1 +11.1
ImageNet ResNet-50 BNN (1W/1A) Top-1 准确率 (%) 30.8 34.2 48.5 +17.7
CIFAR-100 MobileNetV2 BNN (1W/1A) Top-1 准确率 (%) 40.5 42.0 48.0 +7.5
SST-2 BERT-base BNN (1W/1A) 准确率 (%) 50.9 (崩溃) 79.0 82.8 +31.9

在权重二值化(BWN,1W/32A)设定下,由于连续激活值天然保障了梯度回传,三种方法表现相近:在 CIFAR-100 上 ResNet-50 的 STE、StoMPP、StoMPP+STE 准确率分别为 64.3%、69.0%、68.3%,证实了传统 BNN 深度退化的根源确实由激活值量化引发。

消融实验

在 CIFAR-100 上对冻结顺序(正向逐层 vs 全局 vs 逆向逐层)和掩码策略进行了彻底消融(全量消融均在纯无 STE 环境下完成,以排除替代梯度的掩饰作用)。

量化体制 掩码调度顺序 掩码策略 ResNet-18 训练集 (%) ResNet-18 测试集 (%) ResNet-50 训练集 (%) ResNet-50 测试集 (%) 现象说明
BNN 正向逐层 (Layerwise) 随机 (Stochastic) 83.3 53.8 54.6 40.0 基准最优,未发生阻断
BNN 全局调度 (Global) 随机 (Stochastic) 62.3 53.2 39.5 35.9 随深度增加性能显著恶化
BNN 逆向逐层 (Reverse) 随机 (Stochastic) 28.4 28.4 (崩溃) 9.6 8.6 (崩溃) 早期输出层截断浅层,彻底瘫痪
BWN 正向逐层 (Layerwise) 随机 (Stochastic) 99.9 69.5 99.9 69.0 激活连续,无梯度阻塞
BWN 逆向逐层 (Reverse) 随机 (Stochastic) 99.6 69.7 99.6 66.6 逆向不引发网络瘫痪
BWN 全局调度 (Global) 确定性 (Deterministic) 99.5 70.2 96.9 65.9 传统权重渐进固化依然有效

此外,在超参数敏感度测试中,StoMPP 在 Cosine、Linear、Quadratic、Cubic 等多种递增调度曲线以及刷新率 \(r \in [10, 1000]\) 范围内均展现出极强的稳定性;在与二值专用网络 Bi-Real Net 结合时,从头训练在 BiReal-18 / 34 上分别取得 56.6% / 59.1% 的准确率(优于 STE 的 56.0% / 55.8%)。

关键发现

  • 冻结顺序是全 BNN 的生杀决断点:在 BNN 中逆向逐层冻结(从输出到输入)导致模型彻底坍塌(ResNet-50 在 CIFAR-100 仅 8.6% 测试准确率,接近随机猜测),而在 BWN 中逆向仅有微弱扰动(66.6%)。这无可辩驳地证明了“激活值诱导的梯度阻断”是造成全二值网络深度优化崩溃的根本源头。
  • 锯齿状动态表征层级过渡与重构:StoMPP 在训练精度曲线上呈现独特的“锯齿状(sawtooth)”轨迹——每当新的一层开始渐进硬化时,网络精度短暂下挫,随后未冻结参数迅速自适应并拉升精度;这与 STE 单调平滑但深层早衰的曲线形成鲜明对比。
  • 与二值残差架构的协同互补:在融入 Bi-Real Net 的密集跳跃连接后,锯齿状震荡被显著平抑,证实残差旁路在物理层面上为逐层渐进二值化提供了额外的梯度兜底通路。

亮点与洞察

  • 首次确立了“逐层正向渐进调度”能够彻底摆脱直通估计器(STE)而独立训练出高精度深度全二值网络,打破了 BNN 长期对人工启发式替代梯度的绝对依赖。
  • 深刻揭示了全 BNN 与权重 BNN 之间的非对称优化机理:精准定位了“激活值二值化”所诱发的上游梯度阻断,为后续离散网络架构设计提供了坚实的理论解释支点。
  • 软重置(Soft Refresh)机制通过期望不变的小样本随机置换,巧妙平衡了离散探索与拓扑稳定性,使得整套训练框架对调度形状与刷新超参极不敏感,工程鲁棒性极高。

局限与展望

  • 作者承认的局限:为了在严格受控条件下隔离量化规则的真实效果,实验统一剥离了学习率衰减调度、知识蒸馏以及复杂的网络修改,导致绝对精度与当前堆砌多种技巧的非受控 SOTA 相比存在差距。
  • 探索与延伸空间:逐层均匀分配训练步长尚未达到全局最优,浅层与深层特征对量化误差的容忍度不同,未来可探索自适应层级步长调度;此外,激活阻断机制尚需设计更为微观的梯度干预实验以给出更加严格的理论形式化界定。

相关工作与启发

  • vs Vanilla STE (BinaryConnect / Binarized Neural Networks): 传统 STE 采用全局固定硬前向与替代反向,前后向失配严重随深度放大;StoMPP 通过渐进正向掩码保留连续下游后缀,不仅能作为 STE-free 独立方案,与 STE 复合(StoMPP+STE)更能彻底解决深度衰减。
  • vs INQ (Incremental Network Quantization): INQ 仅适用于权重剪枝和权重离散化(BWN),若直接应用于含二值激活的全 BNN 会因全局随机阻断而崩溃;StoMPP 首次通过正向时空流向将渐进固化成功扩展到包含二值激活的全离散网络。
  • vs Bi-Real Net / ReActNet: 此类工作依靠引入额外的浮点捷径分支或可学习偏移来缓解精度损失;StoMPP 专注于训练规则本身,二者相互正交且可无缝叠加,在 Bi-Real 架构上展现出了出色的通用适配性。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 颠覆了依赖替代梯度的传统思维,从量化介入的时空顺序创新性破解了 BNN 深度崩溃。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 ResNet、MobileNet 与 Transformer,全量消融剥离 STE 干扰,对比严密扎实。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑链条极为清晰,梯度阻断物理图景明确,主次机制剖析深入浅出。
  • 价值: ⭐⭐⭐⭐⭐ 为超低比特量化与离散神经网络的规模扩展性(Depth Scalability)树立了极具启发性的训练范式。