跳转至

Layerwise Progressive Freezing: A Training Scaffold for Depth-Scalable Binary Networks

会议: NeurIPS 2026
arXiv: 2606.27759
OpenReview: https://openreview.net/forum?id=iJdFhf3jyg
领域: 模型压缩
关键词: 二值神经网络, 逐层渐进冻结, 激活梯度阻断, 直通估计器, 深度扩展性

一句话总结

针对全二值神经网络在深层网络中随深度加剧退化甚至崩溃的难题,本文揭示了二值激活导致的“梯度阻断”机制,提出了输入至输出的逐层随机掩码渐进冻结算法(StoMPP),在完全不依赖直通估计器(STE-free)以及与 STE 结合(StoMPP+STE)的两种模式下均显著拉平了深度退化曲线。

研究背景与动机

深度神经网络在算力受限的边缘设备与嵌入式硬件上面临巨大的存储与能耗挑战。二值神经网络(Binary Neural Networks, BNNs)将网络权重和激活值严格约束在 \(\{-1, +1\}\),能够将浮点乘加计算转换为高效的按位异或非运算(XNOR-popcount),从而带来数量级级别的推理加速与能效收益。然而,从头训练全二值网络极具挑战:前向传播中的硬符号量化算子 \(\text{sign}(\cdot)\) 处处导数为零,使得标准的基于反向传播的梯度优化无法直接进行。过去近十年来,学术界普遍采用直通估计器(Straight-Through Estimator, STE)在反向传播中用可导代理梯度(如恒等或截断梯度)替代非导符号函数,但由于前向离散计算与反向连续梯度之间存在固有的不匹配(mismatch),模型性能在网络层数加深时急剧崩溃——例如在 ResNet-50 等深层网络上,常规 STE 训练的 Top-1 准确率往往断崖式下跌,甚至完全陷入近乎随机猜测的灾难性停滞。

现有的大量改进工作主要聚焦于“如何修复 STE 自身的缺陷”,例如设计更加平滑的渐进代理导数、引入梯度形状缩放、学习自适应偏置与尺度因子,或是借助蒸馏和复杂的学习率退火调度。但这些策略依然受制于 STE 的基本架构假设,即在整个训练过程中全局、静态地维持硬符号前向运算与代理反向梯度的脆弱平衡。尽管经典的模型量化领域曾探索过渐进冻结(Progressive Freezing,如 INQ),即在训练过程中逐步将连续参数固定为离散值,但在全二值网络中直接套用全局渐进冻结却完全行不通:一旦激活值参与二值化,深层网络就会迅速崩溃停滞,而同样的调度在仅二值化权重的网络(BWN)中却能顺利收敛。

这一对比揭示了一个深刻且此前未被系统阐明的优化物理机制:在包含二值激活的全二值网络中,一旦某个激活层被硬冻结为符号函数,其反向局部梯度便处处归零,直接在反向传播路径上筑起一道物理上的“梯度阻断壁障(Activation-Induced Gradient Blockade)”,彻底截断了上游更浅层参数接收损失反馈的通路。在仅二值化权重的网络中,激活值始终保持连续可导,因此即使部分权重被冻结,梯度依然能顺畅回传;而全局渐进策略在深浅各层随机冻结激活,导致前向计算图各处的上游梯度通道被随机切断。核心 idea:与其在反向传播中费尽心机拟合不可导的离散算子,不如重塑量化介入的时空顺序——采用由浅入深的逐层渐进冻结范式(StoMPP),确保当前过渡层之后始终保留一段完全连续可导的下游网络作为梯度回传通道,从而在根源上消除激活梯度阻断。

方法详解

整体框架

StoMPP(Stochastic Masked Partial Progressive Binarization,随机掩码局部渐进二值化)为全二值网络的训练建立了一个严谨的逐层离散化脚手架。对于网络中的每一层,算法同时维护实值底层变量(权重与前向激活)以及对应维度的二值掩码 \(M \in \{0, 1\}\),掩码中的 1 表示该分量已进入离散冻结状态,0 表示该分量仍处于连续松弛状态。在整个训练周期内,二值化并非全局铺开,而是按照从输入层到输出层的严格波浪式时序逐层推进。

在任意特定训练阶段,整个网络被自然划分为三个功能明确的区域: 1. 已冻结前缀(Frozen Prefix):由第 \(1\) 层到第 \(\ell-1\) 层组成,所有分量的掩码均为 1,完全执行离散符号运算;由于它们位于当前学习层的前方,其参数已在前期得到充分训练,此时不再需要接收反向梯度更新; 2. 当前过渡层(Transition Layer):即第 \(\ell\) 层,掩码处于动态演进中,利用随机局部冻结与软刷新机制平稳提高二值化比例; 3. 未冻结后缀(Unfrozen Suffix):由第 \(\ell+1\) 层到第 \(N\) 层组成,掩码均为 0,权重和激活完全保持连续可导的代理状态(激活采用截断函数 \(\text{clip}\),权重保持恒等),为损失函数反传至过渡层 \(\ell\) 提供了一条毫无阻碍的顺畅梯度通道。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据输入层"] --> B["已完全冻结前缀层 (1 ~ ℓ-1)<br/>M=1,完全二值化,前向sign算子"]
    B --> C["当前过渡层 ℓ:软刷新随机掩码调度<br/>Bernoulli(p) 局部重采样,平稳递增冻结率"]
    C --> D["未冻结连续后缀层 (ℓ+1 ~ N)<br/>M=0,保持连续clip/identity,提供通畅反向梯度通路"]
    D --> E["网络输出与损失计算"]
    E -.->|反向传播:后缀层无损反传梯度| C
    C -.->|梯度解耦双模更新| F["参数底层实值变量更新"]

关键设计

1. 输入至输出的逐层前向渐进冻结机制:彻底规避激活引起的梯度阻断 传统全局渐进量化在二值激活网络中崩溃的根源,在于激活函数导数 \(\frac{\partial \text{sign}(z)}{\partial z} = 0\) 几乎处处成立。若在深层随机冻结激活值,反向传播经过该硬算子时局部梯度直接降为零,导致更靠近输入端的浅层网络完全失去反向更新信号。为此,StoMPP 将二值化推进顺序严格约束为由输入向输出的单向波浪推进。假设网络总共有 \(N\) 个待量化层,分配给每层 \(E\) 个周期的过渡预算。当网络在第 \(\ell = \lfloor e / E \rfloor\) 层推进时,前序层已固定为离散表示,而后续层 \(\ell+1, \dots, N\) 全部为未冻结状态。由于未冻结激活采用截断函数 \(\text{clip}(z) = \max(-1, \min(1, z))\),在不饱和区间内其局部导数恒为 1,确保了从输出损失到当前第 \(\ell\) 层的整条反向链路始终畅通无阻。这种由浅入深的单向性从结构拓扑上彻底消除了深层激活对浅层更新的屏蔽,使深层网络能够在保持全二值化最终形态的同时稳定扩展深度。

2. 随机掩码局部冻结与软刷新机制:兼顾训练平稳性与动态适应能力 在单个过渡层的离散化过程中,若采用确定性排序(如按权重绝对值大小一次性冻结)或每步全量重采样掩码,容易引发两个极端问题:前者会使模型过早陷入局部劣解且无法应用于无明显绝对值分布的动态激活值;后者则造成前向计算图的高频剧烈抖动。StoMPP 提出了带软刷新(Soft Refresh)的随机掩码策略。在过渡步 \(\tau \in \{1, \dots, T\}\) 内,算法定义了平滑递增的目标冻结比例 \(p(\tau) = (\tau / T)^3\)(默认三次方调度,先慢后快,让模型在低离散度下充分探索)。在每个训练迭代步,系统仅随机抽取总元素数中比例为 \(k = \lfloor n / r \rfloor\) 的一小部分索引(默认刷新率 \(r = 100\),即单步仅扰动约 1% 的参数),并仅对这部分索引重新依 \(\text{Bernoulli}(p(\tau))\) 采样掩码,其余绝大部分索引保持上一时刻的冻结状态不变。这种设计在数学上严格保证了整个张量期望冻结率精确匹配 \(p(\tau)\),同时维持了前向计算的时间连续性,赋予了未冻结与已冻结单元充分协同适应的缓冲期。

3. 解耦梯度代理的双模训练范式:支持纯无 STE 训练与冻结位 STE 增强 为了将连续松弛与离散二值化统一在同一表达式中,StoMPP 构建了前向混合映射: $\(u' = M \odot \text{sign}(u) + (1 - M) \odot \text{Smooth}(u)\)$ 其中激活值的平滑连续代理为 \(\text{Smooth}(u) = \text{clip}(u)\),权重的平滑代理为 \(\text{Smooth}(u) = u\)。在反向传播中,底层实值变量 \(u\) 的梯度被严格形式化为: $\(\nabla_u = \nabla_{u'} \odot \left[ (1 - M) \odot \frac{\partial \text{Smooth}(u)}{\partial u} + M \odot \text{Surr}(u) \right]\)$ 该式将梯度计算显式拆解为未冻结分量的精确导数与已冻结分量的代理导数 \(\text{Surr}(u)\)。这一解耦使 StoMPP 展现出优异的理论与工程灵活性: - 无 STE 模式(STE-free StoMPP):设定 \(\text{Surr}(u) = 0\)。已冻结单元的反向导数严格置零,梯度信号完全且仅由未冻结部分承载。由于在反向路径上不存在任何前后向不匹配的伪梯度,这是学术界首个在大规模深层全二值网络上成功稳定收敛的完全无 STE 训练机制。 - 混合增强模式(StoMPP+STE):设定 \(\text{Surr}(u) = 1\)。仅在已冻结单元上施加 STE 恒等映射,使已固化参数仍能微调吸收全局信息。该模式将逐层渐进拓扑结构与 STE 优势结合,实现了最佳的绝对精度。

损失函数 / 训练策略

StoMPP 在标准交叉熵分类损失下进行端到端优化,无需任何额外的蒸馏损失、量化正则项或复杂的重参数化惩罚。为了严密控制变量并剥离学习率调度对渐进策略的干扰,论文采用极为纯粹的极简训练配置: - 优化器与超参:采用标准 SGD(带动量 0.9),常数学习率 0.1,批大小为 256。 - 权重衰减禁用:全线关闭权重衰减(Weight Decay = 0),以防止 \(L_2\) 正则化与二值权重向 \(\pm 1\) 极化的物理趋势产生冲突。 - 调度分配:默认在 CIFAR-10/100 上训练 200 个 epoch,平均分配到各个量化层;ImageNet 上根据架构分别训练 97 到 131 个 epoch。量化掩码采用三次方单调调度 \(p(\tau) = (\tau/T)^3\) 与软刷新率 \(r = 100\)。首层卷积、最后一层全连接及下采样残差投影层按惯例保持全精度,其余层全部参与二值化。

实验关键数据

主实验

论文在 CIFAR-10、CIFAR-100、ImageNet 视觉基准以及 SST-2(BERT 微调)上展开了系统评测。在完全匹配的极简训练配置下,量化网络在不同深度下的 Top-1 准确率如下表所示:

数据集 模型架构 量化类型 Vanilla STE (%) StoMPP (STE-free) (%) StoMPP + STE (%) 相对 STE 提升
CIFAR-10 ResNet-18 BNN (1W/1A) 77.8 80.9 86.1 +8.3
CIFAR-10 ResNet-34 BNN (1W/1A) 61.5 76.0 82.1 +20.6
CIFAR-10 ResNet-50 BNN (1W/1A) 51.5 69.5 78.6 +27.1
CIFAR-100 ResNet-18 BNN (1W/1A) 49.1 53.8 58.0 +8.9
CIFAR-100 ResNet-34 BNN (1W/1A) 33.7 39.8 47.8 +14.1
CIFAR-100 ResNet-50 BNN (1W/1A) 26.7 40.2 46.5 +19.8
ImageNet ResNet-18 BNN (1W/1A) 41.9 39.5 45.5 +3.6
ImageNet ResNet-34 BNN (1W/1A) 37.0 39.7 48.1 +11.1
ImageNet ResNet-50 BNN (1W/1A) 30.8 34.2 48.5 +17.7
SST-2 BERT-base BNN (FFN 1W/1A) 50.9 (崩溃) 79.0 82.8 +31.9
CIFAR-100 MobileNetV2 BNN (1W/1A) 40.5 42.0 48.0 +7.5
CIFAR-100 ResNet-50 BWN (1W/32A) 64.3 69.0 68.3 +4.0

注:Vanilla STE 在深层网络(如 ResNet-50)以及 Transformer(BERT)上均发生严重的深度性能退化甚至近乎随机猜测(SST-2 上 50.9% 即为完全崩溃)。StoMPP 彻底解决了深度扩展瓶颈。

消融实验

为纯化评估渐进调度本身而非代理梯度的作用,所有消融实验均在无 STE(STE-free StoMPP)设置下于 CIFAR-100 上训练 200 个 epoch 完成。

1. 冻结推进顺序与掩码策略消融(ResNet-18 vs. ResNet-50):

模型类型 掩码时序策略 掩码采样方式 R18 Train (%) R18 Test (%) R50 Train (%) R50 Test (%) 结论与分析
BNN 全局推进 (Global) 随机采样 62.3 53.2 39.5 35.9 深层退化明显
BNN 逐层前向 (Layerwise) 随机采样 83.3 53.8 54.6 40.0 深度扩展性最佳
BNN 反向逐层 (Reverse) 随机采样 28.4 28.4 (崩溃) 9.6 8.6 (极度崩溃) 激活阻断导致浅层无法学习
BWN 全局推进 (Global) 确定性极化 99.5 70.2 96.9 65.9 仅二值权重对顺序不敏感
BWN 全局推进 (Global) 随机采样 99.2 70.3 96.1 67.5 权重无激活阻断效应
BWN 逐层前向 (Layerwise) 随机采样 99.9 69.5 99.9 69.0 保持优异收敛水平
BWN 反向逐层 (Reverse) 随机采样 99.6 69.7 99.6 66.6 无梯度阻断,反向亦可收敛

2. 学习率敏感度与二值架构兼容性消融(CIFAR-100):

实验切片 变量设置 Vanilla STE Top-1 (%) StoMPP (STE-free) Top-1 (%) 观察要点
学习率敏感度 (R18 BNN) \(\text{lr} = 0.001\) 10.4 33.2 小学习率下 StoMPP 依然保持极佳优化效率
学习率敏感度 (R18 BNN) \(\text{lr} = 0.01\) 30.9 51.5 性能大幅领先 STE 达 20.6%
学习率敏感度 (R18 BNN) \(\text{lr} = 0.05\) 46.8 54.2 达到最优区间
学习率敏感度 (R18 BNN) \(\text{lr} = 0.10\) 48.9 53.8 默认基准学习率下保持稳定胜出
Bi-Real Net 架构 (BiReal-18 从头训练) \(\text{lr} = 0.06\), 300 ep 56.0 56.6 无需架构修改即可无缝迁移到二值专用网络
Bi-Real Net 架构 (BiReal-34 从头训练) \(\text{lr} = 0.06\), 300 ep 55.8 59.1 在更深 BiReal-34 上增益扩大至 +3.3%

关键发现

  1. 顺序不对称性证明了激活阻断是深度崩溃的主因:在全二值网络中,反向逐层冻结(从输出向输入)导致 ResNet-50 的测试准确率直接暴跌至 8.6%(近乎随机猜测),而前向逐层冻结能稳定达到 40.0%;反观仅二值化权重的网络(BWN),反向冻结与前向冻结表现几乎一致(66.6% vs 69.0%)。这一强烈的不对称性无可辩驳地证实:是二值激活切断了反向传播,而前向逐层推进是维持深层网络训练梯度的物理生命线。
  2. 独特的“锯齿状(Sawtooth)”动态收敛曲线:观察 StoMPP 的训练轨迹发现,每当一个新层进入冻结阶段时,验证集准确率会出现短暂下挫,随后未冻结参数迅速协同自适应,精度强势反弹回升。这种周期性锯齿形态反映了网络由浅入深消化离散化扰动的内在机制;引入额外跨层跳跃连接(如 Bi-Real Net)能够显著平抑锯齿振幅。
  3. 更快的低 Epoch 收敛效率与鲁棒性:在低轮次预算下(如 50~100 轮),StoMPP 训练准确率上升速度远快于 STE;在调度函数形状(线性、二次、三次、余弦)和软刷新率 \(r \in [10, 1000]\) 的广泛参数扫描中,前向逐层策略均表现出极强的平稳性,无需精细调参。

亮点与洞察

  • 颠覆了“BNN 必须依赖全局连续梯度代理”的刻板认知:学术界长期认为从头训练全二值网络必须依靠 STE 贯穿始终。本文证明,仅仅改变离散化在层间推进的时空拓扑,就能在反向完全不提供代理梯度(\(\text{Surr}(u) = 0\))的条件下训练深层网络并反超常规 STE,兼具坚实的理论优美性与实验说服力。
  • 正交且兼具通用扩展性:StoMPP 本质上是一种优化调度框架,完全独立于特定网络拓扑。论文展示了其在经典卷积(ResNet)、轻量化移动端架构(MobileNetV2)、二值专用残差架构(Bi-Real Net)以及自然语言 Transformer(BERT-base)上的广泛适用性,为未来超大规模 1-bit LLM 的从头预训练提供了极具前景的范式启发。

局限与展望

  • 基线训练配方受限:为了纯化比较二值化规则本身,作者禁用了学习率退火衰减与权重衰减,使得纯 ResNet-18 在 ImageNet 上的绝对分数(45.5%)未对齐那些融合了知识蒸馏、复杂多级训练阶段或专门损失工程的 SOTA 复合方案。
  • 分层时间预算均等化可能并非最优:目前算法对每一层分配完全相等的训练步长,但网络浅层与深层在通道数、感受野及特征抽象级别上存在巨大差异,未来可探索自适应或基于曲率感知的层级过渡分配算法。
  • 梯度阻断机制缺乏微观干预探针:虽然宏观上的正反向对照实验高度吻合理论推导,但尚缺乏对中间层梯度张量模长与奇异值谱的微观动力学数学量化分析。

相关工作与启发

  • vs. 经典 STE 与代理梯度改进(Bi-Real Net, ReActNet, OvSW):经典工作尝试在固定的全二值拓扑下优化反向导数近似,但深层不匹配依然存在;StoMPP 则另辟蹊径,通过时间轴与层级轴的动态连续松弛解耦,直接绕开了梯度不匹配难题。
  • vs. 增量网络量化(INQ, Dong et al.):INQ 仅能处理权重参数的逐批静态冻结,面对激活二值化时会因梯度阻断而瘫痪;StoMPP 首次提出了针对“权重+激活”双二值系统的逐层前向推进与局部软刷新机制,成功克服了全二值网络的深度崩溃障碍。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [揭示激活梯度阻断物理机制,开创了无 STE 逐层离散化新范式]
  • 实验充分度: ⭐⭐⭐⭐⭐ [横跨 CNN/Transformer/轻量网络多任务,包含极具深度说服力的正反向消融与收敛动力学分析]
  • 写作质量: ⭐⭐⭐⭐⭐ [论证逻辑严密,图文对照清晰,实验控制变量严谨]
  • 价值: ⭐⭐⭐⭐⭐ [为深层乃至百亿参数超低比特二值模型的训练稳定性提供了全新的基础设施级思路]