WhiteDLG's

知识蒸馏与高效模型设计 · 四篇论文解读

2026-3-4 Deep Learning / Compression 256 Views

自蒸馏概述

# self-distillation (原文标题,保留为段落风格)

—— 知识在模型内部迁移:深层向浅层传递

1. Multi-Stage Vision Token Dropping: Towards Efficient Multimodal Large Language Model

关注的问题

多模态大语言模型(如能够理解图片和文字的LLaVA)在处理输入时,会将图片分割成许多小的视觉单元(称为视觉token)。这些视觉token数量巨大,而且包含大量重复或冗余的信息(spatial and temporal redundancy),导致模型在推理(理解并生成答案)时需要处理海量数据,计算量大,速度慢。

现有方法的局限性

之前的一些解决方案试图通过丢弃不重要的token来减少计算量。但它们判断一个token是否重要的依据比较片面,通常只依赖于单个阶段(比如只看图片编码阶段,或者只看文本预填充阶段)的信息。

架构

MustDrop架构图

根据你提供的论文内容,我仔细通读了全文。这篇论文的核心方法论中,并没有明确提及或设计“自蒸馏”机制。

论文提出的 MustDrop 是一种 “训练无关”(Training-free)的多阶段令牌压缩方法。它不涉及任何需要额外训练的“蒸馏”(Distillation)过程,而是完全依靠推理过程中的启发式规则和注意力分数来动态决定丢弃哪些视觉令牌。

既然论文本身没有设计自蒸馏,我将为你详细拆解其实际的核心机制(多阶段丢弃策略),然后再结合框架图进行解释。

一、论文实际的核心设计方法:多阶段生命周期评估

MustDrop 的核心思想是:一个视觉令牌是否重要,应该在其整个“生命周期”中动态判断。 这个生命周期分为三个主要阶段,每个阶段都有针对性的策略:

1. 视觉编码阶段 (Vision Encoding Stage) - 处理视觉自身冗余

  • 目标: 在不依赖文本的情况下,减少视觉令牌的空间冗余,并锁定绝对重要的令牌。
  • 方法:
    • 局部空间合并 (LSM): 将图像令牌视为2D网格,通过滑动窗口扫描相邻令牌。如果窗口内的令牌相似度极高(超过阈值),就将它们合并成一个代表令牌。这能去掉大量因空间连续性产生的重复信息(如图2a中物体的连续区域)。
    • 关键集保留 (Key Retaining): 利用ViT中[CLS]令牌对其他令牌的注意力分数,找出那些对理解图像整体至关重要的令牌。这些令牌会被放入一个“关键集”,在任何后续阶段(预填充、解码)都绝对不允许被丢弃

2. 预填充阶段 (Prefilling Stage) - 处理跨模态冗余

  • 目标: 利用用户输入的文本指令,筛掉与问题无关的视觉令牌。
  • 方法: 双注意力过滤 (Dual-Attention Filtering)。这是一个非常精细的筛选机制,包含两个维度:
    • 全局注意力过滤: 计算一个视觉令牌与所有文本令牌的总注意力得分。如果某个视觉令牌与整体文本的相关性极低,就把它放入待删除候选集 \( D \)。
    • 个体注意力过滤: 这是为了防止误删。有些视觉令牌可能与整体文本相关性不高,但却与某个特定单词(例如问题中的“number”、“color”)高度相关。MustDrop会检查每个视觉令牌与任何一个文本令牌的最大注意力值。如果最大值依然很低,才最终判定为不重要并删除。

3. 解码阶段 (Decoding Stage) - 处理缓存冗余

  • 目标: 进一步压缩KV缓存,减少显存占用和计算量。
  • 方法: 输出感知缓存策略 (Output-Aware Cache Policy)
    • 观察: 作者发现,在生成回答(解码)时,到了较深的层(如Layer 2之后),图像令牌几乎不再受到生成令牌的关注(注意力稀疏)。
    • 操作: 既然浅层之后视觉信息已无大用,MustDrop在KV缓存中只保留经过前两轮筛选后最终存活下来的极少数视觉令牌。这意味着在解码阶段的大部分层中,视觉部分的KV缓存被大幅缩减。

二、框架图 (Figure 2) 详细解释

现在结合上图,我们来逐一解析图2中的四个子图是如何对应上述方法的。

(a) 视觉编码阶段

  • 图释: 图中展示了“Window scan”(窗口扫描)和“Spatial Merging”(空间合并)。
  • 对应方法: 这就是局部空间合并的直观体现。原始密集的网格令牌经过窗口扫描后,相似相邻的令牌被合并(如图中箭头所示),令牌数量显著减少(图注中提到 Retained ~76% Tokens)。
  • 同时发生的事: 图中用灰色框出的“Key Set”就是关键集。它基于与[CLS]的高注意力分数被选出,这些令牌会被标记为“Not To Prune”(不可修剪),无论后续文本是什么。

(b) 预填充阶段概览

  • 图释: 展示了文本令牌(Language Tokens)和视觉令牌(Vision Tokens)一起输入LLM的多层网络。
  • 对应方法: 这个阶段的核心是利用文本语义来指导修剪。图中显示了随着层数加深(Layer 1 到 Layer L),视觉令牌的数量在动态变化(绿色的 Vision Tokens 在不断减少)。具体如何减少,由图(c)决定。

(c) 预填充阶段中的双注意力过滤

  • 图释: 这是一个放大的细节图,显示了“Dual-Attention Filtering”的机制。
  • 对应方法:
    • 左侧分支(Global): 对应全局注意力过滤。计算视觉令牌 \( V_j \) 与所有文本令牌(\( T_1, T_2, ..., T_N \))的注意力之和。
    • 右侧分支(Individual): 对应个体注意力过滤。寻找视觉令牌 \( V_j \) 与单个文本令牌(如 \( T_i \))之间的最大注意力值。
    • 决策: 结合这两个分数,综合判断该视觉令牌是否应该被剪掉(To Prune)还是保留(Not To Prune)。这种双路设计避免了只看整体而忽略细节,或者只看细节而忽略全局的问题。

(d) 解码阶段的输出感知缓存

  • 图释: 左侧是原始的KV Cache(所有层都存了很多视觉令牌),右侧是 MustDrop 的 Output-Aware KV Cache。
  • 对应方法: 这就是输出感知缓存策略
    • 图中显示,在 MustDrop 的设计下,从较浅的层(如 Layer 1, Layer 2)开始,KV Cache 中只保留最终层(Final Layer)还会用到的极少数视觉令牌(Retained ~8% Tokens)。
    • 例如,图中在 Layer 1 原本需要存储 309 个视觉 KV,经过策略调整后,只存储最终会存活的那 65 个。这极大地节省了从浅层到深层的存储空间和访存开销。

总结: 图2非常清晰地展示了 MustDrop 的流水线作业
Step (a) 在视觉编码器里做初步清洗(合并相似物 + 保护关键物);
Step (b) & (c) 在 LLM 预填充时结合文本做精细筛选(双注意力把关);
Step (d) 在 LLM 解码生成时做最后的缓存瘦身(只留最终幸存者)。

总结:视觉编码阶段、预填充阶段、解码阶段不断筛选视觉信息token。

2. Self-Distillation: Towards Efficient and Compact Neural Networks

自蒸馏示意图

总结:一样的逻辑,知识是在同一个模型内部进行迁移的,从模型的深层部分转移到浅层部分。

损失函数设计

  • 监督学习信号(交叉熵损失),保证分类正确性;
  • 知识蒸馏信号(KL 散度损失),实现从深层到浅层的知识迁移;
  • 特征模仿信号(L2 损失),进一步提升特征表示的一致性。

3. Wavelet Knowledge Distillation: Towards Efficient Image-to-Image Translation

关注的问题

GANs在图像处理(比如把草图变成逼真照片)方面很厉害,但模型太大、太慢,不适合在实际应用(如手机APP)中部署。

实现

这篇论文《Wavelet Knowledge Distillation: Towards Efficient Image-to-Image Translation》的核心设计方法可以概括为:从频率视角分析GAN的生成缺陷,并提出一种专注于高频信息的知识蒸馏方法,以在压缩模型的同时保持生成质量。以下是其设计思路和核心方法的详细解释:

1. 研究动机与问题分析

1.1 问题背景

GAN在图像翻译任务中表现优异,但模型参数量大、计算成本高,难以部署在资源受限的设备上。

1.2 频率视角的实验发现

作者首先通过离散小波变换(DWT)将生成图像和真实图像分解为不同频带(低频+高频),并计算它们之间的归一化L1距离,得到以下关键观察:

  • 所有GAN在低频信息(如轮廓、大块颜色)上生成效果都很好;
  • 但在高频信息(如纹理、边缘、细节)上,尤其是小模型表现明显较差;
  • 这说明:GAN压缩的关键瓶颈在于高频信息的生成能力

2. 核心方法:小波知识蒸馏(Wavelet Knowledge Distillation, WKD)

2.1 基本思想

传统知识蒸馏方法直接让学生模仿教师生成的完整图像,忽略了高频信息的重要性。WKD则通过DWT将图像分解后,仅在高频部分进行蒸馏,让学生专注于学习细节信息。

2.2 技术实现步骤

2.2.1 图像分解

  • 使用离散小波变换将教师生成的图像 \( f_t(x) \) 和学生生成的图像 \( f_s(x) \) 分别分解为四个子带:
    • LL:低频近似
    • LH、HL、HH:三个方向的高频细节(水平、垂直、对角线)

2.2.2 蒸馏目标

  • 仅对高频子带(LH、HL、HH)计算L1损失,忽略低频部分:
\[ \mathcal{L}_{\mathrm{WKD}} = \frac{1}{n} \sum_{i=1}^{n} \| \Psi^H(f_t(x_i)) - \Psi^H(f_s(x_i)) \|_1 \]

其中 \( \Psi^H \) 表示提取高频分量的操作。

2.2.3 总损失函数

  • 在原有GAN训练损失(如对抗损失、循环一致性损失等)基础上,加入WKD损失:
\[ \mathcal{L}_{\mathrm{overall}} = \mathcal{L}_{\mathrm{origin}} + \alpha \cdot \mathcal{L}_{\mathrm{WKD}} \]

其中 \( \alpha \) 为平衡超参数。

3. 实验设计与验证

3.1 模型与数据集

  • 模型:Pix2Pix、CycleGAN、Pix2PixHD
  • 数据集:Horse→Zebra、Edges→Shoes、Cityscapes 等
  • 教师模型:原始大模型(ngf=64)
  • 学生模型:通道数减少的轻量版(ngf=32/24/16)

3.2 评价指标

  • FID(越低越好):衡量生成图像质量
  • mIoU(越高越好):用于Cityscapes语义分割任务

3.3 主要实验结果

3.3.1 定量结果

  • WKD在所有数据集和模型上均显著优于传统KD方法
  • 在CycleGAN上实现 7.08倍压缩 + 6.80倍加速,FID几乎无下降;
  • WKD可与其他特征蒸馏方法联合使用,进一步提升效果。

3.3.2 定性结果

  • 学生模型生成图像在纹理、边缘、细节上更接近教师;
  • 避免了传统KD方法中常见的“模糊”、“条纹混乱”等问题。

4. 消融实验与讨论

4.1 频带选择的影响

  • 仅蒸馏低频 → 性能严重下降;
  • 仅蒸馏高频 → 最佳效果;
  • 同时蒸馏高低频 → 效果略差于仅高频;
  • 验证了“高频是瓶颈”的结论。

4.2 判别器与生成器的关系

  • 若只压缩生成器,不压缩判别器 → 判别器过强,对抗失衡,生成器难以学习;
  • 若将判别器也适当压缩(如4.01倍)→ 生成器性能最佳;
  • 表明:判别器的压缩有助于维持对抗平衡,提升生成器效果

4.3 知识蒸馏范式的比较

  • TSKD(传统教师-学生)效果最好;
  • TAKD(教师助理)、DML(互学习)、SD(自蒸馏)在图像翻译任务中效果有限;
  • 说明:图像翻译任务对高质量教师依赖更强,不同于分类任务。

5. 方法总结

组件说明
出发点GAN生成图像高频信息不足,尤其是小模型
核心工具离散小波变换(DWT)
蒸馏对象仅高频子带(LH、HL、HH)
损失函数L1 loss on high-frequency components
优势提升细节生成能力,压缩率高,性能损失小
扩展性可与特征蒸馏方法结合,兼容多种GAN架构

6. 贡献总结

  1. 首次从频率角度系统分析GAN生成质量,指出高频信息是压缩瓶颈;
  2. 提出小波知识蒸馏(WKD),专注于高频信息的学习;
  3. 实验验证WKD在多个数据集和模型上的有效性,实现高压缩比、低性能损失;
  4. 揭示判别器压缩对生成器性能的积极作用
  5. 比较多种KD范式在图像翻译任务中的表现,为后续研究提供参考。

什么是小波?

小波是一种数学工具,专门用来分析图像的频率。

在这篇论文里,作者用“小波变换”把图像拆解开来:

  • 低频部分:图像的轮廓、大色块(粗略的构图)。
  • 高频部分:图像的纹理、边缘、细节(比如皮毛的纹路、鞋带、高光)。

4. AUTOENCODERS AS CROSS-MODAL TEACHERS: CAN PRETRAINED 2D IMAGE TRANSFORMERS HELP 3D REPRESENTATION LEARNING?

关注的问题

要让AI模型理解3D世界(比如点云数据),通常需要大量人工标注的3D数据,这非常昂贵且困难。

相比之下,2D图像和文本数据更容易获取,并且已经有许多非常强大、成熟的预训练模型(如基于Transformer的模型)。

架构

ACT框架图

这篇论文的核心贡献是提出了一个名为 ACT(Autoencoders as Cross-Modal Teachers) 的自监督学习框架。你提到的“自蒸馏”设计,在论文的语境中,实际上是一种跨模态、两阶段的知识蒸馏方法。它并非传统意义上单个模型的自我蒸馏,而是利用在其它模态(2D图像或语言)上预训练的强大家模型作为“教师”,通过蒸馏的方式指导一个3D“学生”模型的学习。

下面我将先详细解释这种独特的两阶段“自蒸馏”设计方法,然后再结合你提供的框架图进行解析。

第一部分:ACT的两阶段“自蒸馏”(跨模态知识蒸馏)设计方法

ACT 框架通过两个核心阶段,将2D/语言模型的知识迁移到3D模型上。

阶段一:训练“跨模态教师”模型 (3D Autoencoding with Pretrained Foundational Transformer)

这个阶段的目标是将一个在2D图像或语言上预训练的Transformer模型,转化成一个能够理解3D点云的“教师”模型

  • 输入:3D点云数据(如ShapeNet中的物体)。
  • 过程
    1. 冻结与提示:他们冻结(Freeze)了预训练Transformer(如ViT或BERT)的大部分参数,以保留其从海量数据中学到的通用知识(如纹理、轮廓、语义等)。同时,引入了少量的可学习提示(Prompt)嵌入。这些提示作为可调参数,引导冻结的Transformer去适应和处理3D数据。
    2. 3D自编码器训练:他们构建了一个以这个冻结Transformer为核心的3D自编码器。编码器部分是“冻结Transformer+提示”,解码器则负责从编码后的特征重建原始3D点云。
    3. 自监督学习:整个自编码器通过一个叫做离散变分自编码器(dVAE) 的方式进行训练。训练目标是让模型能够尽可能地重建输入的点云。这个过程完全是自监督的,不需要任何人工标注。
  • 结果:经过这个阶段的训练,这个原本不懂3D的2D/语言Transformer,学会了如何将3D点云编码成富含语义信息的特征。它已经成功转变为一个强大的 “跨模态3D教师”

核心思想:通过自编码重建这个“通用”任务,利用提示调优,将冻结教师模型中的“暗知识”(dark knowledge)引导到3D特征空间中,从而创造出一个优质的教师。

阶段二:知识蒸馏训练“3D学生”模型 (Masked Point Modeling as Cross-Modal Knowledge Distillation)

这个阶段的目标是训练一个纯粹的3D Transformer学生模型,让它学会教师模型对3D数据的理解方式

  • 输入:同样是3D点云数据。
  • 过程
    1. 构建教师-学生体系:将第一阶段训练好的、参数固定不变的“跨模态3D教师”模型作为监督信号。同时,初始化一个待训练的、结构相对简单的3D Transformer学生模型。
    2. 掩码点建模:这是训练学生模型的核心任务。他们会对输入点云的局部块进行随机掩码(Mask),即遮挡掉一部分信息。然后,让学生模型去预测被遮挡部分的内容。但学生预测的目标,并不是原始的点云坐标,而是教师模型在看到一个完整(未被掩码)点云时,为这些被遮挡部分所编码出的“潜在特征(latent features)”
    3. 特征级蒸馏:学生模型的优化目标是最小化它预测的特征与教师提供的“答案”特征之间的差距(如余弦相似度损失)。这个过程就是知识蒸馏,学生被强制去模仿教师的内部“思考”过程和特征表示。
  • 结果:通过这种方式训练出来的3D Transformer学生模型,不仅在结构上适用于3D任务,更重要的是,它继承了来自2D/语言教师的强大语义理解能力和几何表示能力。

核心思想:将掩码建模视为一种知识蒸馏。教师模型提供的特征作为比原始点云坐标更高级、更富含语义的“软标签”,从而引导学生模型学习到更强大的3D表示。

第二部分:ACT框架图(Figure 1)解析

现在,我们对照你提供的图来梳理整个流程。

(a) 在大规模数据上预训练Transformer

  • 这部分展示了ACT所利用的“知识源头”。
  • 图中列举了在大规模数据(如ImageNet图像、自由形式的语言)上预训练的Transformer模型,例如ViT(视觉Transformer)、BERT(语言模型)和CLIP(视觉-语言模型)。
  • 这些模型已经掌握了丰富的2D视觉概念或语言语义,构成了ACT的“知识库”。

(b) 第一阶段:使用预训练Transformer进行3D自编码

  • 这部分对应我们上面讲的阶段一。
  • 输入:一个“可爱的蓝眼睛猫坐在枕头上的点云”。
  • 处理流程
    1. 局部几何嵌入:首先通过最远点采样(FPS)和分组(Grouping)将点云切分成多个局部块(类似于图像的patch)。然后,一个轻量的点网络(如DGCNN)将这些局部块编码成一系列Token(嵌入向量)
    2. 提示调优:这些Token加上可学习的提示(Prompt),一起输入到一个冻结的预训练Transformer(如ViT)中。这个冻结的Transformer是知识的核心载体,通过提示来适应3D任务。
    3. 3D自编码:Transformer的输出再经过处理,最终通过一个解码器(Decoder) 尝试重建原始的输入点云(图中的“reconstruction loss”)。这个自编码过程让预训练Transformer学会了将3D信息映射到它自己的特征空间。
  • 输出:训练完成后,这个由“冻结Transformer+提示”构成的编码器(Teacher Encoder) 就成为了一个强大的“跨模态教师”。

(c) 第二阶段:将掩码建模作为跨模态知识蒸馏

  • 这部分对应我们上面讲的阶段二。
  • 输入:同样是点云数据。
  • 处理流程
    1. 构建两个分支
      • 教师分支(顶部):就是阶段一训练好的、参数固定的“跨模态教师”。它对完整的点云进行编码,输出高质量的潜在特征(Latent Features) 作为“知识”。
      • 学生分支(底部):是一个待训练的3D Transformer学生模型。它对输入点云进行同样的分块和嵌入。但关键一步是,它对一部分Token进行了掩码(Mask),然后用一个标准的Transformer(Point Cloud Transformer)来处理这些被破坏的序列。
    2. 特征级蒸馏:学生模型需要预测被掩码位置的特征。它的预测目标,正是教师模型在对应位置输出的潜在特征。损失函数(如图中的 \(\mathcal{L}\))计算学生预测和教师特征之间的差距,并通过反向传播更新学生模型的参数。
  • 输出:训练完成后,这个3D Transformer学生模型就具备了强大的几何理解能力和泛化性能,可以用于各种下游任务(如分类、分割)。

关于掩码

在ACT框架的第二阶段(Masked Point Modeling as Cross-Modal Knowledge Distillation),掩码(Mask)的对象是输入点云经过分块和嵌入后得到的“Token”(标记)

为了让你更直观地理解,以下是具体的掩码内容和实现细节:

1. 掩码的具体对象:局部几何块 Token

在输入3D Transformer学生模型之前,原始点云会经过一个预处理步骤:

  1. 采样与分组:对于一个包含 \(N\) 个点的点云,首先使用最远点采样(FPS)选出 \(N_s\) 个种子点(例如 64 个)。然后以每个种子点为中心,利用K近邻(KNN)算法寻找其周围的 \(K\) 个邻居点,形成一个局部几何块(一个 patch,包含 \(K\) 个点,例如 32 个点)。这样,整个点云就被划分成了 \(N_s\) 个局部块。
  2. 嵌入为Token:通过一个轻量级的PointNet或DGCNN网络(称为局部几何嵌入模块),将这些局部块(\(K\times 3\) 的坐标)编码成一个特征向量 \(\mathbf{x}_i\)。这个特征向量就是输入给Transformer的一个 “Token”
  3. 掩码操作:在学生分支中,会对这一系列Token \(\mathbf{X} = \{\mathbf{x}_1, \mathbf{x}_2, ..., \mathbf{x}_{N_s}\}\) 进行掩码。

结论:被掩码的并不是单个的3D点(point),而是描述局部几何形状的“块”所对应的Token

2. 掩码的策略与比例

论文中对掩码的方式进行了详细的实验和讨论(见论文第8页,Figure 2(a)):

  • 随机掩码(Random Masking)
    • 这是主要采用的策略。对于输入的 \(N_s\) 个Token,随机选择其中的一部分(按照一定的比例)进行掩码。
    • 最佳比例:论文实验发现,较高的掩码比例(Masking Ratio)效果更好。例如,在ScanObjectNN数据集上的实验显示,当掩码比例在 60% - 80% 时,模型在下游任务中的表现最佳。
    • 极端情况:当掩码比例为 0% 时,相当于普通的全Token知识蒸馏,效果反而不佳。这说明掩码任务本身是关键,它能迫使学生模型学习更鲁棒的几何关系。
  • 块掩码(Block Masking)
    • 论文也尝试了块掩码(类似于图像MAE中的策略),即连续地掩码一大片相邻的Token。
    • 实验发现,块掩码对较低的掩码比例更敏感,且整体效果不如高比例的随机掩码。

3. 被掩码部分如何处理?

在输入给学生模型之前,被选中的掩码位置会被替换为一个共享的、可学习的掩码嵌入向量(learnable corruption embedding \(e^{[M]}\)),如论文公式(6)所述。

\[ \mathbf{Z}^{\mathcal{M}} = \mathbb{1}(\mathcal{M})\odot \mathbf{e}^{[M]} + \mathbb{1}(1 - \mathcal{M})\odot \mathbf{T} \]

这意味着,学生模型看到的输入是一部分真实的Token(可见的几何块)和一部分占位符(掩码Token)。

4. 学生模型需要预测什么?(而不是重建原始点)

这是ACT最核心的设计。学生模型的任务是预测被掩码位置的“特征”,而不是预测被掩码的原始点坐标。

  • 预测目标:学生模型的预测目标是教师模型(第一阶段训练好的跨模态教师)在看到完整点云时,为这些掩码位置编码出的潜在特征(latent features)。
  • 损失函数:通过最小化学生预测的特征与教师提供的特征之间的余弦相似度损失(\(\mathcal{L}_{\mathrm{cos}}\)),学生模型被强迫去模仿教师的特征表示空间,从而继承了来自2D/语言模型的语义知识。

总结

掩码的内容:是代表局部几何块的 Token(数量通常为 64 个)。
掩码的比例:通常是 60% 到 80% 的 Token 被随机掩码掉。
掩码的目的:让学生模型在被破坏的输入下,去预测由强大的跨模态教师模型提供的、富含语义的特征表示,而不是去复原具体的点云坐标。

  Tags: Knowledge Distillation, Model Compression, Vision Transformer, 3D Vision
Comments Section (Giscus Loading...)