AC/DC on a Budget - Alternating Sparse Phases¶
会议: NeurIPS2026
arXiv: 2610.27758
论文: https://arxiv.org/abs/2610.27758
领域: 模型压缩
关键词: 稀疏训练、交替阶段、动态重布线、符号翻转、计算预算
一句话总结¶
本文深入揭示了交替稀疏-稠密训练(AC/DC)中稠密阶段的核心收益源自被掩码参数接收到的大梯度所驱动的掩码重布线与权重符号翻转,进而提出了零初始化、缩短稠密窗口以及仅更新高梯度非活动权重的 ARC(交替重布线压缩)方法,在显著降低训练计算量(FLOPs)的同时匹配甚至超越了基线性能。
研究背景与动机¶
在深度学习模型参数量与算力开销急剧增长的背景下,神经网络稀疏化已成为缓解硬件瓶颈的关键手段。然而,绝大多数高性能稀疏化方法(如迭代幅度剪枝 IMP、连续重参数化或基于稠密梯度的动态稀疏训练 DST)仍极度依赖稠密模型的优化过程来探索并确定具备优良泛化能力的稀疏子结构;相比之下,完全从头开始训练稀疏掩码的初始化剪枝(Pruning at Initialization)性能往往显著落后。这一现象表明,稠密训练阶段所提供的结构对齐与支持集扩展(support expansion)能力至今难以被简单替代。
交替压缩/解压缩训练(AC/DC)通过显式交替执行固定支持集的稀疏训练阶段与全参数更新的稠密训练阶段,取得了优异的稀疏模型精度。但这种交替模式在每次解压缩时均需承担全量稠密训练的昂贵 FLOPs 开销。从马鞍面动力学(saddle-to-saddle dynamics)的理论视角来看,神经网络训练可被视为在奇异模态构成的马鞍区域之间逐级跃迁的过程;固定稀疏阶段主要负责巩固当前已学到的模态,而稠密阶段则充当了“重布线窗口”,赋予模型访问非当前支持集(off-support)方向的能力,从而逃离当前的稀疏鞍点。
进一步的经验诊断证实,在稠密阶段中重新引入的被掩码权重虽然数值接近零,却会接收到不成比例的极大梯度,这些大梯度驱动了至关重要的掩码重布线以及最终保留参数的符号翻转修正。既然稠密阶段的核心价值集中于这些高梯度非活动坐标的探索与符号对齐,那么对所有参数进行全量密集更新便存在巨大的计算冗余。核心 idea:将全量稠密解压缩替换为仅更新当前稀疏支持集与少量最高梯度非支持集坐标的稀疏重布线阶段(ARC),并结合非活动参数零初始化与稠密周期截断,以极低 FLOPs 保留支持集跃迁收益。
方法详解¶
整体框架¶
ARC(Alternating Rewiring Compression,交替重布线压缩)旨在保留 AC/DC 交替探索与巩固机制的前提下,消除高昂的全稠密训练开销。算法在初始稠密预热(Dense Warmup)后,进入稀疏阶段与重布线阶段的周期性交替循环,并在最后通过稀疏冷却(Sparse Cooldown)阶段收敛至目标稀疏度模型。
在每一个重布线阶段的起点,算法不展开全量稠密解压缩,而是仅对当前稀疏支持集 \(S_t\) 与通过少量前向/反向探针筛选出的大梯度非支持集坐标 \(G_t\) 所构成的扩展集合 \(A_t = S_t \cup G_t\) 进行更新。重布线阶段结束后,算法再次执行幅度剪枝,挑选出高质量权重进入下一个稀疏阶段进行表征巩固。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["密集预热<br/>Dense Warmup"] --> B["幅度剪枝<br/>生成初始稀疏支持集 St"]
B --> C["稀疏巩固阶段<br/>仅更新 St 内活跃权重"]
C --> D["梯度探测与重布线选择<br/>基于探针批次挑选高梯度非支持集 Gt"]
D --> E["非活动参数零初始化<br/>Zero Init 重置 Gt 坐标"]
E --> F["稀疏重布线阶段<br/>在扩展支持集 At = St ∪ Gt 上训练"]
F --> G{"是否达到<br/>预定循环周期?"}
G -- "否,更新 St" --> C
G -- "是" --> H["稀疏微调与冷却<br/>Sparse Cooldown 最终收敛"]
关键设计¶
1. 梯度引导的稀疏重布线:以紧凑扩展集替代全量稠密更新
标准 AC/DC 在解压缩阶段需要更新所有参数,算力开销巨大;然而经验诊断表明,离开当前稀疏鞍点所需的核心方向几乎全由那些位于掩码外部但梯度能量极高的少数坐标所主导。为此,ARC 彻底舍弃了全稠密解压缩,改为在稀疏重布线阶段仅更新扩展支持集 \(A_t = S_t \cup G_t\)。其中 \(S_t = \operatorname{supp}(\theta_t)\) 为继承自上一阶段的活跃稀疏坐标,而新增的探索坐标 \(G_t\) 定义为在被掩码集合 \(S_t^c\) 中梯度绝对值最大的前 \(q\) 个参数: $\(G_t = \operatorname{Top}_q\left(\left|(g_t)_{S_t^c}\right|\right)\)$ 这些梯度通过在阶段转换点对少量随机抽样的微调探针批次(例如 5 到 100 个 mini-batches)累积梯度获得,且在整个重布线周期内固定支持集 \(A_t\)。这既赋予了高潜力非活动坐标重新生长并竞争进入下一轮掩码的机会,又将阶段更新控制在极低稀疏度(例如仅增加 10%~30% 密度),将探测开销摊销到整个训练阶段。
2. 零初始化策略:破除过时状态阻碍并促进关键符号校正
在传统 AC/DC 中,被重新引入的掩码外部参数通常沿用其上一次离开稠密阶段时的历史陈旧值(Stale init),这导致它们保留了旧的符号和数值偏置。然而,参数符号翻转(sign-flip)实验表明,最终存活权重在训练后期的符号纠正对于稀疏网络泛化至关重要,而陈旧值会阻碍反向传播方向对参数符号的有效修正。本文提出的零初始化(Zero init)在进入重布线或稠密阶段前,将候选非活动参数硬重置为零: $\(\theta_i \leftarrow 0, \quad \forall i \in S_t^c \text{ (在 ARC 中仅作用于 } i \in G_t \text{)}\)$ 重置为零彻底消除了陈旧参数的残留记忆,使最新的梯度方向能够无阻碍地直接决定该参数的初生正负符号与增长轨迹,从而在训练中后期持续释放关键的符号校正能力。
3. 截断稠密窗口:缩短重布线时长以最大化计算收益
标准交替训练通常采用对称的周期配置(例如 10 个稠密轮次交替 10 个稀疏轮次)。然而,激活 Gram 矩阵特征空间相似度诊断揭示,稠密阶段在经历早期平滑且大幅度的表征漂移后,很快就会建立起主要的子空间变换,过长的稠密轮次并不会带来额外的重布线增益。基于这一洞察,本文提出了非对称的截断调度策略:在保持总周期长度 \(D + S\) 不变的前提下,大幅削减重布线轮次至 \(D' < D\)(例如从 10 轮缩减至 3 轮),同时将节省出的轮次全部转为计算开销极低的稀疏巩固阶段 \(S' = D + S - D'\)。该策略在完全不损失模型最终精度的前提下,直接将理论训练 FLOPs 压低至极具竞争力的水平。
损失函数 / 训练策略¶
ARC 继承了基于压缩梯度的随机迭代硬阈值(Stochastic IHT)收敛保证。在目标函数满足受限平滑性(Restricted Smoothness)与集中 PL 条件(Concentrated Polyak–Łojasiewicz Condition)的前提下,通过压缩梯度 \(h_t = P_{A_t} g_t\) 更新参数并辅以硬阈值投影 \(T_k(\cdot)\),可保证线性收敛至误差界 \(O(\beta \sigma_{\rm cg}^2 / \lambda_{\rm cg})\)。在调度实现上,模型在经过稠密预热后,按照预定的周期交替执行稀疏训练与扩展集 \(A_t\) 训练,学习率遵循连续衰减且在阶段切换时不重置,各层稀疏权重在阶段末尾通过全局幅度剪枝恢复至目标稀疏度 \(k\)。
实验关键数据¶
主实验¶
在 CIFAR-100(ResNet18)、ImageNet(DeiT-Small、ResNet50)以及自然语言建模 OpenWebText(GPT-2 124M)上,ARC 与 AC/DC 基线、RigL 以及 Top-KAST 进行了广泛对比。
| 数据集 / 模型 | 稀疏度 | 指标 | ARC (本文) | AC/DC Baseline | RigL |
|---|---|---|---|---|---|
| CIFAR-100 / ResNet18 | 90% | Top-1 准确率 (%) | 75.21 ± 0.26 (ARC 0.3+Zero) | 74.70 ± 0.40 | 72.73 ± 0.41 |
| CIFAR-100 / ResNet18 | 95% | Top-1 准确率 (%) | 75.17 ± 0.98 (ARC 0.3+Zero) | 73.77 ± 0.17 | 71.55 ± 0.26 |
| ImageNet / DeiT-Small | 70% | Top-1 Acc / 训练 GFLOPs | 76.50% / 13.34G (ARC 0.2) | 75.89% / 19.01G | 73.89% / 10.32G |
| ImageNet / DeiT-Small | 50% | Top-1 Acc / 训练 GFLOPs | 77.82% / 18.11G (ARC 0.2) | 77.43% / 21.52G | 76.25% / 15.30G |
| ImageNet / ResNet50 | 95% | Top-1 Acc / 相对训练FLOPs | 73.09% / 0.24× (ARC-Zero) | 73.14% / 0.53× | 69.70% / 0.13× |
| OpenWebText / GPT-2 | 80% | 验证集困惑度 (PPL ↓) | 21.25 (ARC-Zero) | 23.03 | 23.07 |
| OpenWebText / GPT-2 | 80% | 零样本 7 任务平均准确率 (%) | 37.29% (ARC-Zero) | 36.27% | 35.67% |
消融实验¶
针对 CIFAR-100 上 90% 稀疏度的 ResNet18,系统评估了稠密阶段长度截断、零初始化(Zero init)与不同重布线密度预算对准确率与理论算力的影响:
| 配置 | Top-1 准确率 (%) | 相对 FLOPs 状态 | 说明 |
|---|---|---|---|
| AC/DC Baseline (10-D, 10-S, Stale) | 74.70 ± 0.40 | 1.00× (基线) | 标准对称阶段交替 |
| AC/DC + Zero init (10-D, 10-S) | 75.07 ± 0.38 | 1.00× | 零初始化提升显著,促进后期符号修正 |
| AC/DC (3-D, 17-S, Stale) | ~74.65 | ~0.72× | 稠密阶段截断至 3 轮,精度基本无损 |
| ARC (0.1) | 74.45 ± 0.14 | 0.52× | 仅激活 10% 高梯度非支持集坐标 |
| ARC (0.1) + Zero init | 74.74 ± 0.36 | 0.52× | 达到标准 AC/DC 精度,FLOPs 减少近半 |
| ARC (0.3) + Zero init | 75.21 ± 0.26 | 0.65× | 超越 AC/DC 基准精度 +0.51% |
| ARC (0.1) + Zero init (5-D, 15-S) | 74.44 ± 0.06 | 0.42× | 三项改进联合,大幅削减计算量且性能坚挺 |
| 随机选择非支持集坐标 (Random 0.1) | ~72.80 | 0.52× | 证明梯度引导选择对重布线的不可或缺性 |
关键发现¶
- 重布线窗口的饱和效应:稠密或重布线阶段的有效探索高度集中在前几个 epoch 内,将每轮稠密轮数从 10 轮截断至 3 轮几乎不造成精度下降,为计算量压缩释放了巨大空间。
- 零初始化的符号解耦作用:Zero init 破除了历史数值偏置,使得在训练中后期(甚至最后的周期中)仍能产生关键的参数符号反转纠正,在各类架构中普遍带来 0.3%~1.0% 的绝对精度提升。
- 重布线梯度的不可替代性:若使用随机选取的坐标集替代梯度挑选的 \(G_t\) 集合,模型精度会出现断崖式下跌(在 CIFAR-100 上跌幅近 2%),证实了高梯度非支持集坐标承载了逃离鞍点所必需的能量方向。
- 大语言模型上的迁移增益:在 GPT-2 预训练中,ARC-Zero 将验证集困惑度从 AC/DC 的 23.03 显著降至 21.25(相对改善 7.7%),并在 7 项下游零任务评测中有 6 项超越 AC/DC,平均分达 37.29%。
亮点与洞察¶
- 诊断视角严密深刻:将稀疏交替训练机制形式化映射至马鞍面动力学(Saddle-to-Saddle Dynamics),通过梯度-权重相平面散点图、最终掩码符号翻转时序追踪以及激活 Gram 矩阵特征空间相似度三大度量,透彻拆解了稠密阶段“为什么有效”。
- 计算-精度帕累托前沿重塑:打破了过去“要么承受全稠密高开销(AC/DC),要么忍受纯稀疏低精度(RigL)”的二元困境,以极小梯度的探针采样精准锁定了高回报参数子集。
- 软硬件协同的扩展潜力:不仅在非结构化稀疏上表现卓越,在附录中针对 NVIDIA GPU 的 2:4 半结构化稀疏实验中,ARC 仅用 75% 的通道密度重布线便取得了与 AC/DC 全稠密相当的精度(76.06% vs 76.41%)。
局限与展望¶
- 缺乏硬件端原生加速支持:目前实验主要基于 PyTorch 密集张量掩码(Masked Dense Tensors)实现,主要展示的是理论 FLOPs 的压缩,在缺乏专用稀疏计算算子支持的现有通用硬件上无法直接转化为训练墙钟时间(Wall-clock time)的线性加速。
- 探针计算的额外超参:引入了挑选集合 \(G_t\) 所需的探针 mini-batches 数量以及密度预算 \(q\) 两个超参数,虽然对超参相对鲁棒,但在超大规模分布式集群训练时需要统筹多卡通信同步。
- 结构化扩展仍待深入:目前对 2:4 稀疏的研究属于通道级探索,未来需进一步探索细粒度块稀疏(Block Sparsity)与推断加速硬件的端到端软硬件结合。
相关工作与启发¶
- vs AC/DC (Peste et al., 2021): AC/DC 在每个解压缩阶段对全部模型参数进行稠密优化;ARC 揭示了稠密优化的关键仅在于高梯度非活动权重,用 \(A_t = S_t \cup G_t\) 的稀疏重布线替代全稠密,以约一半的 FLOPs 达到甚至超越其精度。
- vs RigL (Evci et al., 2020): RigL 是一种纯稀疏到稀疏(sparse-to-sparse)的动态稀疏训练方法,虽然 FLOPs 极低,但由于完全缺乏稠密预热和集中的支持集扩展窗口,高稀疏度下性能普遍较弱;ARC 则在保持极具竞争力的 FLOPs 预算下大幅拉开与 RigL 的精度差距。
- vs Top-KAST (Jayakumar et al., 2020): Top-KAST 前向与后向分别使用基于权重幅度的前 \(k\) 个参数;而 ARC 仅在阶段转换处通过梯度能量探索并锁定外部坐标,并在后续训练中双向均参与优化,具备更强的逃离稀疏鞍点能力与更优的收敛保证。
评分¶
- 新颖性: ⭐⭐⭐⭐☆ 理论与实证结合紧密,精准定位交替稀疏训练的核心动力学机制并提出高效的稀疏重布线替代方案。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 CIFAR、ImageNet(CNN 与 ViT)、GPT-2 预训练及 2:4 半结构化稀疏,消融与动力学诊断极其扎实。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑严谨流畅,从鞍点理论动力学推导到经验度量设计再到算法落地环环相扣。
- 价值: ⭐⭐⭐⭐☆ 为大规模模型(尤其是大语言模型)的绿色低成本稀疏预训练提供了极具启发性的工程与算法范式。