Which Layer Causes Distribution Deviation? Entropy-Guided Adaptive Pruning for Diffusion and Flow Models¶

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/changlin31/EntPruner
领域: 模型压缩
关键词: 扩散模型剪枝, 流匹配模型, 条件熵偏差, 自适应渐进剪枝, 零样本网络架构搜索

一句话总结¶

针对基于 Transformer 的扩散模型与流匹配模型在下游迁移剪枝易导致分布偏移与模式崩塌的问题，本文提出条件熵偏差（CED）量化各模块对输出生成分布的破坏程度，并结合基于神经正切核（NTK）与梯度信噪比（ZiCo）的零样本代理构建自适应渐进剪枝框架 EntPruner，在 DiT 与 SiT 上实现高达 2.22 倍推理加速且生成质量几无衰减。

研究背景与动机¶

近年来，以扩散模型（Diffusion Models）和流匹配模型（Flow Matching Models）为代表的大规模视觉生成模型展现出惊人的合成质量。随着模型骨干网络从传统的 U-Net 全面转向基于 Transformer 的架构（如 DiT、SiT、PixArt-α），模型凭借极强的可扩展性取得了逼真的画质，但伴随而来的庞大参数量和高昂显存开销，严重阻碍了它们在边缘设备与低延迟交互式场景中的部署。

然而，现存的主流剪枝方法（如 BK-SDM、Diff-Pruning、LD-Pruner）大多将扩散网络视作传统判别式分类网络，依赖权重绝对值大小或梯度范数来评估层重要性。这种做法忽视了生成模型的核心属性——输出分布的多样性（Diversity）与条件保真度（Condition-Fidelity）。在将预训练生成模型迁移至特定下游任务时，不同网络模块在保持分布平衡中的职责显著分化；更严重的是，暴力的一次性剪枝（One-shot Pruning）极易破坏预训练权重已建立的生成流形，导致不可逆的分布漂移甚至模式崩塌（Mode Collapse），使剪枝后的恢复微调极度低效脆弱。

作者深入探究了不同网络层对生成分布的具体影响，发现移除不同 Transformer Block 会诱发截然不同的分布偏移模式：部分层被剔除后熵显著增加，导致输出漂移向随机噪声；而另一些关键层被剔除后熵急剧降低，诱发明显的模式崩塌。核心 idea：利用条件熵偏差（CED）精确度量各 Block 移除后对生成分布的绝对偏移量以判定冗余度，并利用 NTK 条件数与 ZiCo 零样本代理在训练过程中动态确定剪枝时机与剪枝幅度，实现保真度自适应渐进模型压缩。

方法详解¶

整体框架¶

EntPruner 的整体执行流程划分为两个阶段：第一阶段为基于条件熵偏差（Conditional Entropy Deviation, CED）的模块重要性评估与冗余排序；第二阶段为基于多指标零样本神经网络架构搜索（Zero-Shot NAS）的自适应渐进剪枝。整个系统在多阶段训练循环中自主决策当前步剪掉哪些层、剪去多少比例，并无缝继承前序阶段权重，避免了一次性结构破坏。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["预训练扩散/流匹配骨干<br/>(DiT / SiT-XL/2)"] --> B["条件熵偏差评估 (CED)<br/>逐 Block 剔除测量分布熵变"]
    B --> C["熵引导冗余度排序<br/>CED 绝对值低者优先列入剪枝候选"]
    C --> D["零样本自适应候选搜索<br/>NTK 条件数 + ZiCo 梯度打分 + 参数正则"]
    D --> E["多阶段自适应渐进剪枝<br/>分阶段继承权重微调，防止模式崩塌"]
    E --> F["压缩后轻量化模型<br/>参数缩减 30%~50%，加速 1.33x~2.22x"]

关键设计¶

1. 条件熵偏差评估：度量生成模型输出分布的结构偏离

判别模型仅关注决策边界，而生成模型必须维持输出连续概率分布 $p(x)$ 的保真度与覆盖范围。对于去噪网络输出特征分布 $X \sim \mathcal{N}(\mu, \sigma^2)$，其连续微分熵定义为： $$ \mathcal{H}(X) = -\int p(x) \log p(x) dx = \log(\sigma) + \frac{1}{2}\log(2\pi) + \frac{1}{2} $$ 当临时剔除第 $i$ 个 Block 时，作者观察到熵变符号具有直接的物理含义：$\Delta \mathcal{H} > 0$ 表示去噪输出失控、向各向同性高斯随机噪声漂移；$\Delta \mathcal{H} < 0$ 则说明输出特征多样性塌缩、陷入均值崩溃的平凡解。无论熵变方向为何，其偏离均破坏了生成质量。因此，条件熵偏差采用绝对差值进行定义： $$ \text{CED}i = \left| \mathcal{H}_i}) \right| $$ 具有高 CED 绝对值的 Block 是维系生成分布稳定的关键结构（必须予以保护），而低 CED 的 Block 则意味着其被剔除后几乎不引起输出分布扰动，属于下游任务上的高冗余层，被赋予优先剪枝权。}}(X) - \mathcal{H}_{\text{pruned}}(X \mid \text{Drop}{\textit{block

2. 神经正切核与梯度统计驱动的零样本自适应搜索：免验证开销的动态调度

传统渐进式剪枝往往依赖启发式人工设定的层数衰减曲线，或者每步都需要在验证集上耗时评估性能。EntPruner 将每一剪枝阶段 $k$ 下的子网络结构选择重新表述为轻量级零样本 NAS 优化问题： $$ \psi_k^* = \arg\min_{\psi_k \in \Lambda_k} \mathcal{K}(\omega(\psi_k), \Omega) $$ 为了在不产生前向/反向微调开销的前提下预估子网络的可训练性（Trainability）与泛化能力，框架联合引入了流匹配 NTK 条件数与带符号逆方差梯度指标 ZiCo。在流匹配速度场预测 $v_\theta(x_t, t)$ 下，NTK 矩阵 $\hat{\Theta}$ 的特征值主导了梯度优化的最慢收敛模式，采用 NTK 条件数： $$ \mathcal{K}\kappa(\psi) = \frac{\lambda_0}{\lambda_m} $$ 条件数越小，意味着损失曲面越平滑、梯度下降收敛越快且稳定。与此同时，引入改进的带符号 ZiCo 代理衡量层级梯度的绝对均值与方差比率： $$ \mathcal{K} \right) $$ 由于原生 ZiCo 越大表示可训练性越优，作者在目标中取负号转换为最小化问题，促使选出的子网络兼具平滑损失景观与高梯度信噪比。}}(\psi) = \sum_{l=1}^N \log \left( \sum_{\omega \in \omega_l} \frac{\mathbb{E}[|\nabla_\omega \mathcal{L}|]}{\sigma(\nabla_\omega \mathcal{L})

3. 联合排序投票与模型参数正则化：稳健决策机制

由于 NTK 条件数与 ZiCo 梯度的数值量纲和尺度差异巨大，直接线性加权容易造成单一指标主导决策。EntPruner 采用基于相对排名的无参数投票机制，并加入候选子网络总参数量 $\Omega$ 的显式轻量化正则： $$ \psi_k^* = \arg\min_{\psi_k \in \Lambda_k} R(\psi_k), \quad \text{s.t.} \quad R(\psi_k) = R(\mathcal{K}\kappa(\psi_k)) + R(\mathcal{K}(\psi_k)) + \gamma R(\Omega) $$ 其中 }$R(\cdot)$ 代表候选集合内的名次排名序号（1st, 2nd, ...），$\gamma$ 设定为 0.5 作为效率正则因子。各指标下名次越好的候选架构获得越小的累加积分，累计积分最低者当选为当前阶段的最佳子网络 $\psi_k^*$。随后模型直接继承前一阶段对应的已微调参数，进入阶段性适应性微调，完美规避了一次性大比例剪枝引发的破坏性震荡。

实验关键数据¶

主实验¶

论文在 ImageNet 256×256 以及三个典型细粒度下游数据集（CUB-200-2011、Oxford Flowers、ArtBench-10）上对 DiT-XL/2 与 SiT-XL/2 进行了全方位评测。实验涵盖常微分方程（ODE）与随机微分方程（SDE）两种采样求解器，以 50 步采样和 4.0 的无分类器引导（CFG）系数为标准配置。

下表展示了在 SiT-XL/2 架构上不同剪枝算法与微调方案的生成质量（FID、IS）与推理效率对比：

模型与采样设置	方法	剪枝稀疏度	CUB (FID↓ / IS↑)	Flowers (FID↓ / IS↑)	ArtBench (FID↓ / IS↑)	参数量 (M)	推理加速比
SiT w/ ODE	Full Fine-tuning	0%	5.32 / 6.02	11.78 / 3.71	8.80 / 7.32	675.12	1.00×
SiT w/ ODE	LD-Pruner	35%	5.70 / 6.03	12.02 / 3.75	10.78 / 6.63	435.78	1.82×
SiT w/ ODE	EntPruner (Ours)	35%	5.48 / 6.07	11.75 / 3.82	10.03 / 6.88	435.78	1.82×
SiT w/ ODE	LD-Pruner	50%	6.86 / 6.16	12.09 / 3.79	12.81 / 6.35	334.67	2.22×
SiT w/ ODE	EntPruner (Ours)	50%	6.68 / 6.18	11.86 / 3.82	12.65 / 6.41	334.67	2.22×
SiT w/ SDE	Full Fine-tuning	0%	5.17 / 5.87	12.47 / 3.77	13.33 / 6.56	675.12	1.00×
SiT w/ SDE	LD-Pruner	35%	5.24 / 6.10	12.32 / 3.74	16.16 / 6.20	435.78	1.49×
SiT w/ SDE	EntPruner (Ours)	35%	5.22 / 6.11	12.10 / 3.74	15.25 / 6.31	435.78	1.49×
SiT w/ SDE	LD-Pruner	50%	5.98 / 6.19	12.92 / 3.75	18.74 / 5.90	334.67	1.85×
SiT w/ SDE	EntPruner (Ours)	50%	5.83 / 6.15	12.77 / 3.77	18.69 / 5.90	334.67	1.85×

在 ImageNet 256×256 预训练模型直剪场景下（SiT 30% 剪枝比例，ODE 求解器），EntPruner 斩获了 2.69 的 FID，相较于全量 SiT-XL/2（FID 2.15）仅下降 0.54，显著优于 BK-SDM（FID 3.48）、Diff-Pruning（FID 2.79）以及 LD-Pruner（FID 6.81）。在 DiT 架构上剪枝 30% 时，EntPruner 在 Flowers 数据集上实现 11.99 FID（全量微调为 21.05），相对提升幅度达 43.04%，并优于 LoRA、DiffFit、BitFit 等参数高效微调 baseline。

消融实验¶

为了单独验证 CED 熵重要性准则与自适应渐进剪枝机制的独立贡献，作者在 Oxford Flowers 数据集上对 SiT 架构进行了严谨的拆解实验：

| 配置方案 | FID ↓ | IS ↑ | 参数量 (M) | 计算量 MACs (G) | 单图延迟 Latency (s) | 相对加速比 | 说明 | |---|---|---|---|---|---|---| | 全量微调 (Full Fine-tuning) | 11.78 | 3.71 | 675.12 | 228.85 | 0.20 | 1.00× | 原始未剪枝模型基线 | | w/o CED（剪除高 CED 关键层） | 12.06 | 3.81 | 435.78 | 147.13 | 0.09 | 1.82× | 误伤关键分布维持层，FID 恶化 0.31 | | w/o Ada. Pruning（一次性剪枝） | 11.84 | 3.80 | 435.78 | 147.13 | 0.09 | 1.82× | 缺乏训练动态适配，收敛后质量受损 | | EntPruner (完整方法) | 11.75 | 3.82 | 435.78 | 147.13 | 0.09 | 1.82× | 性能全面恢复甚至反超全量微调 |

关键发现¶

生成分布敏感性具有极强层级异质性：SiT-XL/2 在中间层（如 Block 6）呈现极高的正向熵偏离（剔除后输出向白噪声退化），而 DiT-XL/2 则在靠近输出的深层（Block 27-28）呈现严重的负向熵偏离（剔除后陷入低熵模式崩塌）。这证实了生成模型压缩不能采用判别式网络的通用剪枝规则，必须针对架构与分布动态进行定制。
自适应渐进剪枝有效防止模式塌缩：若在训练初始时采用一次性硬剪枝（One-shot），破坏的生成概率流形很难在有限下游微调中完全修补；而分阶段借助 NTK 条件数和 ZiCo 梯度特征渐进式剔除，使模型获得了自适应平滑重参数化的缓冲期。
剪枝兼具正则化效果：在针对特定下游细粒度数据集（如 Flowers）适配时，30%~35% 的剪枝模型 FID 表现反常超越了全量微调（Flowers 上 11.75 vs 11.78），表明适度剪除生成冗余层能有效抑制大模型在小规模数据集上的过拟合。

亮点与洞察¶

提出针对生成模型特性的 CED 评估标尺：创新性地将连续概率分布微分熵引入模块重要性度量，通过熵变绝对值捕捉“噪声发散”和“模式崩塌”两类截然不同但同样致命的分布劣化，突破了传统依赖权重范数的刻板做法。
将扩散模型剪枝重构为零样本 NAS 协同决策问题：精妙地结合了理论收敛界的 NTK 条件数与经验梯度统计的 ZiCo 指标，借助极低计算代价的相对排序投票（Rank-based Voting），彻底消除了剪枝过程中昂贵的搜索与验证开销。
跨扩散与流匹配范式的通用可扩展性：算法对 DDPM 式评分扩散和连续流匹配 ODE/SDE 求解器均保持无缝兼容，为 DiT 类大参数 Transformer 在端侧的落地铺平了道路。

局限与展望¶

分布假设简化：计算微分熵时将输出高维特征简化为局部单变量高斯分布（$\mathcal{N}(\mu, \sigma^2)$），虽然获得了简洁的对数方差解析形式，但在真实多峰或非线性强流形上可能损失高阶统计矩信息。
粒度受限于 Block 级别：目前方法仅聚焦于整个 Transformer Block 的深度结构化丢弃，尚未细化到自注意力机制内部的 Head 级别剪枝或 FFN 通道维度的细粒度结构化稀疏。
跨模态与视频生成的扩展验证留白：尽管理论推导支持更广泛的生成流匹配模型，但当前实验主要覆盖图像级类别条件生成，在更庞大的多模态文生图（如 FLUX、SD3）和长时序视频扩散模型上的适用性仍有待进一步实证。

评分¶

新颖性: ⭐⭐⭐⭐ [首次从输出分布熵变角度审视生成模型剪枝，并构建了基于 NTK-ZiCo 的零样本自适应调度]
实验充分度: ⭐⭐⭐⭐⭐ [跨越 DiT/SiT 两大架构、覆盖 ODE/SDE 采样求解器，并在 ImageNet 及 3 个下游数据集提供详实指标与消融]
写作质量: ⭐⭐⭐⭐⭐ [逻辑层层递进，从分布偏差的直观观察自然过渡到熵公式推导与零样本搜索算法]
价值: ⭐⭐⭐⭐⭐ [为端侧部署大参数量 Diffusion Transformer 提供了即插即用、开箱即用的实用工程与算法解法]