精读笔记

Problem Setting

这篇论文不是在做一般的 diffusion acceleration，而是在处理一个更尖锐的问题：当图像生成/编辑模型已经大到无法在消费级硬件上运行时，能否通过结构化剪枝把它压到可部署规模，同时不把编辑、参考图遵循、文本渲染等高层能力一起剪掉。

真正困难点有三层。第一，目标压缩比非常激进，HunyuanImage-3.0 从 80B 到 20B，本质上不是小幅 slimming，而是重构了大部分 MoE 容量。第二，模型架构不是以往剪枝论文最熟悉的纯 DiT，而是 MoE-based decoder-only backbone，参数集中在专家层，剪枝后会直接改变路由后的表征分布。第三，图像编辑比普通 T2I 更依赖条件保持和局部变换，任何中间层错位都会被后续层放大成 identity drift、instruction drift 或细节退化。

以前方法卡住的地方在于，它们多半假设剪枝后只需要在固定局部块上做 teacher feature matching，或者在输出噪声/velocity 上做 KD。但高压缩比下，student 自己前几层已经产生了偏移；如果后续监督仍建立在 teacher trajectory 上，就会出现典型 off-policy mismatch：训练时看到的是 teacher state，推理时滚动的是 student state。TMP 试图解决的关键矛盾正是：局部监督易优化但不等于端到端一致，端到端监督更真实但在大剪枝后太难优化。

Motivation

已有路线不够的核心原因不是剪枝指标不够聪明，而是恢复训练没有处理“分布漂移后的表示对齐”。importance metric、NAS、mask learning、cheap replacement、time-step-aware pruning 都在回答“删哪里”或“何时用小模型”，但对于极高压缩后每层输入分布被 student 自身污染的问题，处理得不充分。

作者的关键观察是：剪枝后的 collapse 很大程度来自 representation misalignment 的级联传播。局部 feature supervision 可以缓解，但如果一直停留在固定局部块，它只是在 teacher-provided clean input 上训练 student 的某个片段；这和真实推理时 student 自己生成上游特征的情形不同。因此缺口不是一个更复杂的 pruning score，而是一个能从局部可学性逐渐过渡到全局一致性的训练组织方式。

从这个角度看，TMP 的动机比较合理：高压缩模型需要 curriculum，因为一开始直接端到端对齐太不稳定；同时需要 on-policy 成分，因为最终部署时只能在 student 自己的 hidden-state manifold 上运行。

Core Idea

TMP 的核心思想是把剪枝后的恢复过程重新建模为“逐步扩张作用域的表示对齐问题”。它不是只让每一层或每个 block 模仿 teacher，也不是直接让整个 student 预测 teacher output，而是先把网络切成很多短区间，在局部区间内做 feature matching；随后按树结构合并相邻区间，让监督目标从短程局部一致逐渐变成长程表示一致，最终接近端到端 token representation matching。

这个设计引入的 inductive bias 是：深层生成模型的剪枝恢复应该遵循一个 coarse-to-fine / local-to-global 的优化路径。早期局部监督降低 credit assignment 难度，后期合并监督迫使跨层误差被联合吸收，而不是每个小块各自对齐 teacher 后在组合时仍然不兼容。

另一个本质变化是 mixed-policy。off-policy 蒸馏使用 teacher hidden states，提供干净、稳定、低方差的目标；on-policy 蒸馏使用 student rollout 到某层后的 hidden states，再让 teacher 的对应局部层作用在这些 student states 上构造目标。这等于让监督落在 student 实际会访问的状态分布附近。相比 prior 的纯 teacher-forcing local KD，这更像是在修正 student manifold 上的局部动力学，而不是只拟合 teacher manifold 上的函数值。

Method

第一，结构化剪枝本身是相对朴素的启发式：HunyuanImage 上做 MoE expert pruning 加 MLP width pruning，Z-Image 上做 MLP width pruning。它解决的是容量与显存的硬约束，核心变化是直接减少常驻参数；但这部分不是论文最有技术含量的地方，更像是为后续恢复训练制造一个极高压缩的 student。

第二，off-policy feature distillation 解决的是剪枝后局部模块缺少稳定学习信号的问题。它把 teacher 在某区间入口的 hidden states 输入 student 对应区间，让 student 匹配 teacher 区间出口。必要性在于降低优化难度，避免一开始就让严重受损的 student 承担完整端到端轨迹。

第三，on-policy feature distillation 解决的是 teacher-forcing 与 student inference 不一致的问题。它先让 student rollout 到某层，得到带偏差的 student hidden states，再用 teacher 的局部层在这个状态上生成 stop-gradient target，要求 student 对应区间追上这个 target。核心变化是监督分布从 teacher trajectory 扩展到 student trajectory，这比普通局部 KD 更贴近部署时的误差传播路径。

第四，tree-structured interval merging 解决的是局部目标之间不兼容的问题。随着训练推进，相邻区间被合并，多个短程 loss 被替换成更长跨度的统一 loss。这相当于 curriculum：先学局部映射，再学跨区间组合，最后逼近全局表示一致。它的价值在于让 optimization path 更平滑，而不是一次性要求剪枝模型重建完整 teacher dynamics。

第五，prediction alignment / velocity distillation 是最终输出空间校准。它解决的是 feature 对齐后仍可能存在生成目标偏差的问题，尤其是 flow matching velocity prediction 和 Hunyuan 的 autoregressive text/token head 对齐。但从机制上看，这更像收尾，不太可能是高压缩可行性的主要来源。

Key Insight / Why It Works

我认为这篇最重要的 insight 是：高压缩结构化剪枝的主要瓶颈不是参数重要性估计，而是剪枝后 hidden-state dynamics 的再对齐。传统剪枝方法容易把问题看成静态子网络选择；TMP 实际上把它看成“在 student 自身轨迹上重建 teacher 的局部动力系统”。这比单纯匹配最终噪声预测更合理，因为图像生成/编辑模型的高层能力很大程度编码在中间 token 表示的逐层变换里。

最可能的核心贡献是 mixed-policy feature distillation 与 tree merging 的组合。单独 off-policy KD 给稳定 teacher anchor，但会训练-推理错配；单独 on-policy KD 更真实，但早期 student states 太差，target 可能噪声大且优化不稳。混合策略本质上是在 bias-variance 和 stability-realism 之间折中。tree merging 则提供了一个 curriculum，使模型不会一开始就被长程误差压垮，同时也不会永远停留在局部 imitation。

这不是 retrieval，不是 test-time compute，也不是纯 memory reuse；更准确地说，它是 representation alignment + curriculum learning。对 HunyuanImage 的部署结果还混入了工程因素：FP8 quantization、CPU/GPU offloading、按阶段释放显存，这些主要解决“能否跑起来”，不是“剪枝后能力为何保留”。论文里把 20B 跑上 4090 作为亮点是有实际价值的，但这部分属于 engineering deployment，不应和 TMP 的算法贡献混为一谈。

哪些可能只是辅助？velocity distillation 大概率是必要但非决定性的后处理；magnitude/score-based pruning 也不是新机制。真正让 75% 压缩后还能恢复的，应该是分阶段 feature-space 对齐。如果没有强 ablation，不能排除增益主要来自大规模恢复数据、teacher 冗余容量、MoE expert redundancy，而不是 tree/mixed-policy 本身。文中未充分说明训练数据规模、采样分布、合并 schedule 敏感性，因此增益来源不清。

另一个需要警惕的点是“能力保留”的归因。图像编辑样例和私有人评可以说明没有明显崩，但不能证明复杂 reasoning 被保留。所谓 multi-image reasoning 很可能仍主要来自原 teacher 的表征冗余和数据覆盖，剪枝模型只是没有破坏这些常见模式；若测试分布真正超出恢复数据覆盖，能力下降可能更明显。

Relation To Prior Work

TMP 最接近的技术谱系是 diffusion/DiT structured pruning 后的 local distillation，例如 FastFLUX、PPCL 这类 block-wise 或 contiguous-layer supervision。它与这些工作的共同点是都承认剪枝后需要中间层监督，而不仅是输出 KD。

真正不同点在于两个方面。第一，TMP 不把局部区间当成固定训练单元，而是用树状合并让监督区域随训练变大。这使 local KD 从一种静态辅助 loss 变成 curriculum objective。这个变化是实质性的，因为它直接处理了“局部最优组合后不等于全局一致”的问题。

第二，TMP 把 on-policy distillation 从语言模型输出层概念迁移到图像生成 backbone 的中间表示层。这个迁移比表面上更重要：扩散/flow 模型的最终输出只是一步 velocity/noise，而中间 representation 才承载多条件融合、空间语义和编辑指令执行。让 teacher 在 student hidden states 上构造局部 target，本质上是让 student 学会回到 teacher dynamics 的稳定流形。

哪些看似新其实是重组？expert pruning 用 gating score、width pruning 用 magnitude、输出 velocity KD 都是已有思想。树状 curriculum 也不是全新机器学习概念。但把这些组合到 MoE 图像编辑模型的极高压缩场景，并明确处理 off-policy/on-policy representation mismatch，是这篇相对实质的新增信息。

与 time-step-sensitive pruning、dynamic routing、supernet 方法相比，TMP 更偏离 runtime adaptive compute，而是追求一个静态小模型。这对低显存部署更友好，因为不需要保留原模型或 supernet 的最坏情况显存。

Dataset / Evaluation

评估覆盖了两个层面：HunyuanImage-3.0 的图像编辑压缩，以及 Z-Image-Turbo 的 T2I 压缩。这个设计能支撑“方法不只适用于单一 DiT/T2I”的基本 claim，因为它跨了 MoE 与 DiT、编辑与生成、80B 与 6B 两个规模区间。

但最核心的 HunyuanImage claim 主要依赖私有 image editing benchmark 和 human preference，公开可复核性弱。-2.5% 偏好下降听起来很小，但没有充分说明评测集分布、prompt 难度、参考图复杂度、评审协议和失败案例类型。对于图像编辑模型，平均偏好尤其容易掩盖 identity preservation、局部编辑准确性、多对象绑定等细分能力退化。

Z-Image 部分使用公开 benchmark，更能说明 4B pruned model 在通用 T2I 指标上没有明显崩。但这些 benchmark 是否能验证 TMP 的核心 claim——高压缩后 representation dynamics 被更好恢复——并不充分。它们主要测结果质量，不测中间表征对齐、误差传播、on-policy 分布下的鲁棒性。

真实部署方面，单 4090、1024×1024、8 steps 是有意义的工程验证，说明模型不仅是参数量变小，而是能在消费级 GPU 上跑。但这里的成功依赖 FP8 与动态 offload，延迟、吞吐、CPU 内存带宽、不同显卡上的稳定性没有展开。它验证了 feasibility，不等于验证了高效 production serving。

Limitation

第一，方法强依赖强 teacher 和足够相近的 student architecture。TMP 的 feature matching 假设 teacher/student 层区间之间存在可对齐语义；如果剪枝改变了深层信息组织方式，或跨架构差异更大，这种逐区间 matching 可能会变成错误约束。

第二，泛化性仍未被彻底证明。MoE Hunyuan 与 DiT Z-Image 两个案例说明方法有一定通用性，但都属于现代 transformer-based flow/diffusion 图像模型，且 student 是从 teacher 直接结构化裁剪而来。跨模态、跨任务、跨更激进结构变换是否成立，文中未充分说明。

第三，增益归因不清。缺少足够细的 ablation 来区分：tree merging、on-policy、off-policy、velocity KD、剪枝启发式、恢复数据规模分别贡献多少。尤其在 80B MoE 模型中，大量专家本身可能高度冗余，剪掉低 gating-score experts 后能力下降有限并不意外；TMP 可能是在恢复这个冗余系统，而不是发现了普适剪枝原理。

第四，数据依赖可能很重。on-policy distillation 需要 student rollout 覆盖部署时会遇到的 hidden states；如果恢复数据没有覆盖长尾编辑指令、多参考图组合、复杂文字渲染或反事实场景，student 会在未覆盖状态空间中漂移。论文没有给出恢复数据分布，核心能力可能主要来自数据覆盖。

第五，所谓 reasoning / multi-image capability 保留需要谨慎。图像编辑中的“推理”经常更像训练分布中的模式组合与 reference-following retrieval，而不是显式规划。TMP 并没有引入新的长期状态建模或规划机制，只是尽量不破坏 teacher 原有表征。因此它能保留多少真正组合泛化，不能从当前实验直接推出。

第六，工程部署把问题部分转移到了系统侧。20B FP8 + offloading 能在 24GB 上运行，但这可能牺牲速度、增加 CPU-GPU 传输依赖，也可能对驱动、kernel、显存碎片管理敏感。论文强调“可运行”，但没有充分讨论实际交互延迟和批量服务成本。

Takeaway

• 最值得记住的第一点：大模型剪枝的关键正在从“找不重要参数”转向“剪枝后 hidden-state dynamics 如何恢复”。
• 对于生成模型，尤其是多条件图像编辑模型，中间表示对齐比最终输出 KD 更接近能力保留的本质。
• 第二点：on-policy distillation 不只适用于 LLM token output，也可以迁移到视觉生成模型的中间层。
• 只要 student rollout 会改变未来监督分布，就存在 off-policy mismatch；在 hidden space 里构造 on-policy target 是一个可迁移 insight。

一句话总结

TMP 是一篇把结构化剪枝从静态参数裁剪推进到 mixed-policy、curriculum-style 表征动力学恢复的工作，其主要贡献在于让极高压缩比的大型图像生成/编辑模型在工程部署上变得可行，但其算法增益与数据/冗余/系统优化之间的归因仍需更严格验证。