精读笔记
Problem Setting
论文标题：Post-Training Pruning for Diffusion Transformers（arXiv preprint / 2026）。这篇论文实际处理的是 DiT 在无需 retraining 的条件下如何被 aggressive pruning，而不是一般意义上的模型压缩。关键矛盾是：post-training pruning 要依赖少量 calibration activation 做一次性重要性估计，但 DiT 的权重尺度、token 数、层结构异质性使得 LLM pruning 的 saliency 近似不再可靠。以前方法卡在两个地方：一是 Wanda/SparseGPT 风格指标在 DiT 上过度依赖权重项，二是固定 granularity 假设把层间/通道间的真实重要性结构抹平。困难不在实现剪枝，而在高稀疏率下 ranking error 会被 diffusion iterative denoising 放大，最终表现为生成质量和语义一致性快速崩溃。
Motivation
已有路线不够的原因不是 pruning 理论本身失效，而是它们默认的统计 regime 错了。LLM 中长序列带来较大的 activation norm 和更稳定的 token 统计，Wanda 的 |W|·||X|| 排序在那个 regime 下可以工作；DiT 在 latent space 中 token 少得多，activation 统计弱，权重尺度又更大，于是 importance ranking 更容易被 weight magnitude 主导。作者的核心观察有两个：DiT 权重显著大于 LLM 权重；importance map 在二维权重空间中存在不同聚类形态。真正缺的是一个 DiT-aware 的 saliency calibration，而不是更复杂的 pruning optimizer。这个动机是成立的，因为 post-training 方法没有训练反馈，只能依赖先验设计来减少 ranking bias。
Core Idea
论文真正的核心思想是：不要把 DiT 当成 LLM Transformer 来剪，而要重新平衡 weight 与 activation 在 saliency 估计中的角色。STW 把权重项平方，本质上是在把 pruning perturbation 重新解释为二阶能量，而不是一阶幅值贡献；这和 OBD/OBS 的原始精神更一致，也更适合 DiT 中较大权重导致的一阶近似失真。
另一个核心思想是 pruning granularity 本身也是 inductive bias。统一 per-layer 或统一 per-channel 都假设重要性分布的空间结构一致，但 DiT 中不同模块的重要性聚类形态不同。CAG 引入的不是新信息源，而是改变了比较域：对于紧凑分布用 layer-level allocation，对于通道异质明显的层用 channel-level allocation。它的本质区别在于从“全模型统一剪枝规则”转向“根据 importance geometry 选择剪枝坐标系”。
Method
方法层面最关键的是 STW 与 CAG，而不是完整 pipeline。STW 解决的是 saliency ranking 的尺度偏置：原始 Wanda 指标使用 |W|·||X||，在 DiT 中由于权重大、activation token 统计弱，容易让权重幅值支配排序；改成 |W|^2·||X|| 后，更接近二阶 loss perturbation 的能量形式，也把 DiT 权重尺度压到更接近 LLM pruning 有效的范围。
CAG 解决的是 sparsity allocation 的比较范围问题。论文用输出通道平均重要性及其归一化方差来刻画通道异质性：低异质层采用 per-layer pruning，高异质层采用 per-channel pruning。机制上的变化是避免在 adaLN 等强通道结构层里进行错误的全层竞争，也避免在 MLP/self-attention 等较均匀层里过度细粒度分配。
校准数据只用于收集 activation norm，且需要覆盖 diffusion timestep。这里的必要性很强：DiT 的 activation sensitivity 随 timestep 变化，单一 early/late timestep 会给出偏置统计。这个部分虽然不是主要创新，但对 post-training pruning 的稳定性很关键。
Key Insight / Why It Works
最可能真正有效的是 STW 对 saliency 排序的修正。DiT pruning 的失败很可能不是因为缺少更高级的优化，而是因为重要性估计的尺度项错配；一旦 ranking 偏了，高稀疏率下删掉的就是 diffusion denoising 中反复使用的关键权重，误差会跨 step 积累。|W|^2·||X|| 看起来简单，但它把指标从 LLM heuristic 拉回到二阶 perturbation energy，这个 inductive bias 对 DiT 的大权重 regime 更合理。
CAG 的作用更像辅助但很重要：它不是发现新的 saliency，而是减少 allocation noise。尤其 adaLN 这类 conditioning/normalization 相关层可能具有强通道语义或调制结构，per-layer 统一排序会破坏某些通道的整体功能。CAG 保留了这种 latent channel structure，因此在高稀疏率下收益明显。
这篇的增益不太像 scaling、retrieval、test-time compute 或 data coverage，而更像 better inductive bias + representation/statistics alignment。它利用的是 DiT 参数-激活统计和 layer geometry，而不是更多数据或更大计算。需要警惕的是，energy-based 解释有一定事后合理化成分：文中没有严格证明 |W|^2·||X|| 是 DiT 的最优 Hessian surrogate，更多是经验统计 + ablation 支撑。CAG 的聚类解释也偏 heuristic，增益来源可能部分来自给 adaLN 等特殊层更合适的手工 granularity，而不一定来自通用 clustering principle。
Relation To Prior Work
它最接近 Wanda/SparseGPT/OBD-OBS 这条 post-training pruning 谱系，而不是 diffusion-specific retraining pruning。和 SparseGPT 的差异在于它不做重建补偿；和 Wanda 的差异在于 saliency 指标的权重指数与 granularity 自适应。看似新的 energy-based formulation，本质上是把 OBD 的二阶权重项重新注入 Wanda 这类高效 activation-aware metric 中。
和 Diff-Pruning 这类 U-Net diffusion pruning 相比，它的实质不同在于不依赖 retraining，也不主要建模 timestep dynamics，而是针对 DiT Transformer 权重/activation 统计做一次性校准。真正新增的信息是：DiT 的 pruning failure 可以被解释为 saliency scale mismatch 与 granularity mismatch，而不是简单的“diffusion model 更难剪”。
创新程度应当谨慎评价：STW 形式非常简单，严格说是已有二阶 pruning 思想的 DiT-specific 重组；但它抓住了一个被 LLM pruning 迁移时忽略的统计错配，这比堆复杂模块更有价值。CAG 也是 heuristic，但它明确把 granularity 作为可适配对象，这在 DiT pruning 里是实质性的工程-方法结合点。
Dataset / Evaluation
评估覆盖 DiT-XL/2、PixArt-Σ、FLUX.1-dev，包含 ImageNet、COCO、MJHQ，不是单模型单 benchmark 结果，因此足以支持“在多种 DiT backbone 上比 Wanda/magnitude 更稳”的主张。尤其 FLUX.1-dev 这种大模型上的结果增强了方法的外推可信度。
但 evaluation 主要验证的是离线质量保持，而不是部署闭环。参数量和 TFLOPs 下降被报告为压缩/加速收益，但 unstructured sparsity 是否能在真实 GPU inference 中端到端加速，文中未充分说明。2:4 structured pruning 有补充，但主结果仍偏质量指标。
benchmark 对核心 claim 的支持是中等偏强：ablation 分离了 STW 与 CAG，calibration size/timestep 也有分析，说明不是单纯偶然调参。但仍缺少更细的层级归因，例如哪些层贡献最大、CAG 是否主要只在 adaLN 有效、不同 timestep 下被剪权重是否一致。真实世界部署、不同 prompt distribution、视频/多模态 DiT 泛化都没有被充分覆盖。
Limitation
最大隐含前提是 DiT 的权重尺度和 activation 统计稳定，且这种统计差异正是 pruning failure 的主因。如果未来 DiT 架构使用不同 normalization、weight scaling、tokenization、flow matching schedule 或更长 latent sequence，STW 的指数选择未必成立。
第二个上限是它没有恢复机制。没有 retraining、distillation 或 reconstruction compensation 意味着一旦 saliency ranking 错了就无法修正，因此在 70%/80% 极端稀疏下明显退化是预期结果。这不是一个能突破稀疏上限的方法，而是在 40%–50% 区间内做更稳的 one-shot selection。
第三，CAG 可能把问题从“如何全局分配 sparsity”转移为“如何选择 granularity”。文中用归一化方差解释聚类异质性，但阈值和层类型之间的关系没有充分理论化。增益来源不清：可能主要来自对 adaLN 的特殊处理，而不是普适 clustering-aware principle。
第四，速度收益存在 deployment gap。FLOPs/Params 下降不等于真实 diffusion serving 加速；稀疏 kernel、batch size、resolution、memory bandwidth、scheduler 开销都会影响实际收益。文中未充分说明端到端 latency、吞吐、显存峰值和 sparse kernel 兼容性。
Takeaway
- 1. DiT pruning 不能直接继承 LLM pruning 的 saliency 假设；同为 Transformer，权重/activation/token 统计不同会改变 importance ranking 的有效性。
- 2. 对 post-training compression 来说，最重要的可能不是更复杂的优化，而是找到目标模型族的正确统计 regime。
- STW 的价值就在于用一个很小的指标修正对齐了 DiT 的二阶扰动结构。
- 3. Pruning granularity 应被视为 inductive bias，而不是 implementation choice。
一句话总结
这篇论文是把 LLM post-training pruning 迁移到 DiT 时的一次统计校准：真正贡献不在复杂算法，而在指出并修正 DiT saliency 的尺度错配与 granularity 错配。