【KD】Transformer在各个研究领域的轻量化研究进展
目录
引言
Transformer中常用的模型压缩方法
在预训练语言模型中的轻量化Transformer
Q8BERT: Quantized 8Bit BERT
Compressed BERT, known as DistilBERT, is a more compact and efficient version of the original BERT model: lighter, faster, and more affordable in terms of size and weight.
TinyBERT: Distilling BERT for Natural Language Understanding
ALERT is named a compact variant of BERT, designed for the self-supervised learning task of learning language representations.
计算机视觉中的轻量化Transformer
Data-efficient training of image transformation models and an attention-based distillation process.
TinyViT: Fast Pretraining Distillation for Small Vision Transformers
MiniViT: Compressing Vision Transformers with Weight Multiplexing
Dynamic Vision Transformer: An Efficient Vision Transformer Network Employing Dynamic Token Sparsity Mechanism
『多模态中的轻量化Transformer』
MiniVLM: A Smaller and Faster Vision-Language Model
Compressing Visual-linguistic Model via Knowledge Distillation
参考文献
参考
本文系统介绍了Transformer模型中的一些轻量化方法，并从领域角度对主流的轻量化Transformer进行了深入分析。
引言
近年来，Transformer模型在多个领域中被广泛采用，特别是在计算机视觉、自然语言处理以及多模态等前沿领域取得了显著成效。尽管Transformer在多数应用场景中展现出卓越性能，但其庞大的参数规模限制了其在实际应用中的表现，这不仅带来了技术挑战，也给相关研究带来了新的探索方向。笔者系统梳理了Transformer轻量化的主要方法，并从多个领域展示了主流的轻量化Transformer。欢迎各位专家批评指正，共同促进该领域的研究与发展。
预训练模型参数量增长趋势
Transformer中常用的模型压缩方法
量化的基本思想是通过低比特数据类型替代高比特数据类型，同时对存储空间和计算复杂度进行优化。其主要缺点是简单的训练后量化方法在实际应用中会带来较大的性能损失。为此，量化感知训练（Quantized-Aware Training, QAT）被提出，该方法在训练过程中采用全精度优化策略，仅模拟低精度推理过程，从而有效降低了量化带来的性能损失。
剪枝：基于lottery ticket假设，该方法认为模型中关键参数仅占模型参数总量的少数，其余大量参数均为冗余参数，可以通过移除这些冗余参数来实现模型的轻量化。具体而言，剪枝可分为非结构化剪枝和结构化剪枝两种类型。非结构化剪枝通过识别参数矩阵中接近于0或具有接近0趋势的参数并将其置零，从而完成参数稀疏化处理。结构化剪枝则侧重于减少模型中结构化的组件，例如在多头注意力机制中去除不必要的注意力头，在多层Transformer架构中移除某些不必要的层。
非结构化剪枝
在知识蒸馏过程中，通常采用一个大型的教师模型与一个小规模的学生模型之间的互动机制。通过教师模型在监督数据集上生成的"软标签分布"来指导学生模型的学习，这种"软标签"的学习方式能够有效缓解监督数据中标签偏差的问题，从而提升知识迁移的效率和效果。
参数共享：在模型同质的部分间共享参数
在预训练语言模型中的轻量化Transformer
Transformer在自然语言处理领域得到了广泛的应用，其显著的能力推动了预训练领域的发展，并在相关领域带来了革新。然而，随着预训练模型规模的不断扩大，训练与部署预训练模型的成本也在不断增加，因此轻量化研究成为了一个重要方向。由于以BERT为代表的预训练语言模型主要采用Transformer架构，这使得Transformer的轻量化研究成为了当前NLP领域的重要课题。以下是一些经典的预训练语言模型中的轻量化工作，供读者参考。
Q8BERT: Quantized 8Bit BERT
Q8BERT是量化技术在BERT上的直接应用。除了直接应用量化感知训练方式外，Q8BERT还采用指数滑动平均方法来平滑QAT训练过程。经过实验验证，在压缩至原来的4倍参数量的同时，实现了推理速度的4倍提升。实验结果表明，在GLUE和SQUAD数据集上，Q8BERT的性能接近BERT，进一步验证了量化技术在BERT上的有效性。
Q8BERT在下游任务上的表现
DistilBERT, a revised variant of BERT: more compact, more efficient, at a lower cost, and with less weight.
DistillBERT是huggingface推出的精简版Bert模型，它采用了软标签蒸馏技术，在输出层将Bert的层数缩减至原来的一半，并在多个下游任务中展现出良好的性能。
在GLUE上，DisitlBERT只用了60%的参数，保留了97%的性能
TinyBERT: Distilling BERT for Natural Language Understanding
与DistilBERT相似，TinyBERT也采用知识蒸馏在BERT压缩上的应用。与DistillBERT仅在预训练阶段运用知识蒸馏，TinyBERT在预训练和微调两个阶段均采用了两阶段的知识蒸馏策略，从而使小型模型能够有效学习通用与特定任务相关的语义知识。
TinyBERT蒸馏策略
同时，**TinyBERT采用了更加精细化的知识蒸馏机制，对嵌入层的输出结果，Transformer每一层内部的隐藏层输出与注意力机制的输出，以及整个模型的软标签信息均进行了知识蒸馏过程，最终实现了更为精准的知识迁移效果。
更加细粒度的蒸馏方式
最终，TinyBERT通过模型压缩技术实现了BERT模型的精简优化，其压缩后的结构仅包含4层，相较于原模型具有显著的参数量缩减优势，同时其性能不逊于参数量更高的DistilBERT模型。
ALBERT: A Lite BERT for Self-Supervised Learning of Language Representations
ALBERT是Google在2020年ICLR会议上发布的一项重要研究工作。该模型采用了参数分解技术，该技术将词表的嵌入层参数数量显著降低。具体而言，该技术通过将词表规模为V的词汇映射到较小的嵌入维度E，再映射到与隐藏层维度H相等的高维空间，从而实现了嵌入层参数量从O(V×H)优化至O(V×E+E×H)。这种设计不仅有效降低了模型的参数规模，还保持了其在语言理解任务中的性能表现。
此外，ALBERT通过Transformer内跨层的参数压缩实现了参数共享。在跨Tranformer层的完全参数共享下，ALBERT在参数量较低的情况下实现了参数量的大幅压缩，与Bert-base在低参数量条件下取得了相近的效果。同时在相近的参数量下，该方法能够使模型深度增加，隐藏层的维度得到提升，在下游任务中表现更优。
ALBERT的参数量及下游任务表现
计算机视觉中的轻量化Transformer
尽管Transformer在计算机视觉领域的应用相较于NLP领域稍显滞后，但Vision Transformer的问世却使Transformer在视觉模型领域占据了主导地位。
Data-efficient training of image transformers and attention-based distillation
DeiT是Facebook在2021年提出的一项创新工作，其核心方法是通过知识蒸馏技术对视觉Transformer进行降维处理。该研究采用了两种不同的蒸馏方法，成功实现了教师模型与学生模型之间的知识迁移。具体而言，第一种蒸馏方法基于 teacher-student attention 的对齐机制，而第二种方法则采用了 teacher-student feature matching 的策略，确保了知识的有效传递。
Soft Distillation：通过老师模型输出的软标签 进行知识蒸馏。
Hard-label Distillation：基于老师模型预测出的真实标签，将知识传递给学生模型，旨在减少监督学习数据中标签偏差的影响。
DeiT在蒸馏过程中提出了Distillation token的概念，其功能与Class token相似，但Class token被用来基于交叉熵在原始数据上进行训练，Distillation token则被用来模拟老师模型的软分布输出，或者基于老师模型预测的硬标签进行训练。
Disillation token
基于老师模型的蒸馏过程，DeiT在这些条件下的参数量更小、推理效率更快，并且取得了比ViT更好的效果。
TinyViT: Fast Pretraining Distillation for Small Vision Transformers
https://www.ecva.net/papers/eccv_2022/papers_ECCV/papers/136810068.pdf
TinyViT是微软在2022年进行的一次知识蒸馏相关研究，该研究在工程实现上进行了优化改进，使得基于小模型的系统能够在处理更大规模的数据集时依然能够继承和应用大模型的知识。相比之下，DeiT所采用的知识蒸馏方法在计算资源利用方面存在显著劣势，具体表现在蒸馏过程中，教师模型与学生模型需要同时占用GPU内存资源，这不仅限制了batch_size的扩展，还直接影响了学生模型的训练效率。此外，在知识传递过程中，软标签的传播效率也会受到一定的计算资源损耗影响。
为了解决这一问题，TinyViT提出了一种软标签预生成的方法，通过解耦软标签生成与学生成员的训练过程。首先生成并预存软标签，然后利用预存的软标签对学生成员进行训练。由于预存储的软标签向量会导致较大的存储消耗，考虑到这些向量在大多数情况下是稀疏的，作者采用了一种存储稀疏标签的策略，即仅存储概率最高的top-k标签及其对应概率。在对学生成员进行训练时，将这些稀疏标签还原为完整的概率分布，并进行知识蒸馏。整个流程如图所示。
TinyViT的蒸馏流程
在保持相同模型规模的前提下，TinyViT显著提升了知识蒸馏的速度和数据量，并在分类任务方面也表现出色。
MiniViT: Compressing Vision Transformers with Weight Multiplexing
The paper introduces a novel method for compressing vision transformers through weight multiplexing, termed MiniViT. This technique achieves significant compression rates while maintaining model performance, making it a promising approach for efficient vision modeling. The experimental results demonstrate that MiniViT outperforms existing compression methods in terms of both computational efficiency and accuracy, establishing it as a robust solution for transformer-based image compression.
MiniViT是微软在CVPR2022上提出的一种创新方法，该研究团队采用权重复用策略压缩模型参数。与现有的方法不同，MiniViT通过权重复用不仅降低了每层梯度的l2范数的同质化，还减少了最后几层输出特征的相关性。为了解决这一问题，本研究引入了两个关键组件：权重变换模块和权重蒸馏机制。权重变换模块通过在各层之间插入小型自适应结构实现，确保各层输出不至于因参数量相同而产生同质化现象。权重蒸馏机制则通过引入一个辅助的老师模型来引导MiniViT优化输出，从而提升模型性能。整个方法的流程图如下所示。
该压缩方法具有通用性，作者在DeiT和Swin-Transformer架构上进行了实验验证。在更小的模型参数规模下，实验结果表明，采用该压缩方法的模型在ImageNet数据集上的性能表现优异，能够实现与甚至超越现有方法的性能水平。
DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
https://proceedings.neurips.cc/paper/2021/file/747d3443e319a22747fbb873e8b2f9f2-Paper.pdf
该研究由清华大学在Nips2021会议上发布，其研究内容借鉴自模型压缩中的剪枝技术。但其创新点在于仅对Transformer的每一层输入token进行稀疏化处理。研究者认为，在图像处理中，通常会发现存在大量冗余像素，这些像素对模型预测结果的影响较小。因此，通过移除这些冗余部分能够显著提升模型的推理速度。
DynamicViT的主要功能是通过在层之间增加一个轻量预测模块，判断哪些token可以被移除。预测完成后，采用二进制决策掩码完成token的移除。此外，作者调整了训练目标函数，确保每层移除的token数量有限。
DynamicViT的稀疏化模块
基于该稀疏化方法，DynamicViT在原始ViT的基础上实现了显著的加速效果，并在模型性能与推理速度之间实现了良好的平衡。
『多模态中的轻量化Transformer』
多模态模型通常会综合考虑视觉和语言模型的架构设计，因此，其主流架构多为Transformer。然而，在轻量化方面，多模态领域的工作尚显不足，建议读者参考笔者整理的若干具有代表性的多模态轻量化研究工作。
MiniVLM: A Smaller and Faster Vision-Language Model
微软在Oscar模型上实现了轻量化工作，即MiniVLM。其轻量化基于以下观察：在大多数多模态任务中，多模态模型的视觉端通常不需要过于强大的目标检测信息。因此，采用一个相对不精确的目标检测器能够有效地压缩模型的参数量并加快推理速度，同时尽量减少性能损失。
该方法通过融合EfficientNet与Bi-FPN构建了一个轻量化的目标检测器，以实现预期效果。此外，为进一步优化，MiniVLM在多模态Tranformer端实施了压缩措施，将原有的Bert架构替换为更轻量化的MiniLM架构，具体结构如下所示。
轻量化目标检测器
基于允许的精度损失范围内，MiniVLM相较于现有OSCAR模型在推理速度上实现了显著的提升。
MiniVLM的加速效果与下游任务性能
Compressing Visual-linguistic Model via Knowledge Distillation
DistilVLM可视为MiniVLM工作的延伸。不同之处在于，当同时更换目标检测器和Transformer架构时，DistilVLM采用了知识蒸馏以保持性能。其蒸馏策略与TinyBERT一致，亦是经过预训练阶段和微调阶段的两阶段蒸馏：
DistilVLM的蒸馏策略
基于采用了不同的目标检测器，在检测得到的目标区域不同的条件下，后续的知识蒸馏均是无效的。为了解决这一问题，DistilVLM采用了一种视觉token对齐的方法。为了解决这一问题，DistilVLM采用了视觉token对齐的方式，老师模型和学生模型均采用了相同的目标检测器，使得两个模型的检测区域对齐，保证了后续知识蒸馏的有效性。
视觉对齐
最后，在保持与MiniVLM一致的参数设置和推理速度的条件下，DistilVLM实现了性能的显著提升。
参考文献
[1] Q8BERT: Quantized 8Bit BERT
[2] DistilBERT, a compressed variant of BERT; it is more compact, efficient, affordable, and lightweight.
[3] TinyBERT: Distilling BERT for Natural Language Understanding
模型ALBERT：基于轻量版BERT的自监督学习，实现了高效的语义表示。
[5] Data-efficient training of image recognition systems and the application of attention-based distillation methods
[6] TinyViT: Fast Pretraining Distillation for Small Vision Transformers
[7] MiniViT: Compressing Vision Transformers with Weight Multiplexing
[8] DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
[9] MiniVLM: A Smaller and Faster Vision-Language Model
[10] Compressing Visual-linguistic Model via Knowledge Distillation
参考
Fast and Effective！一文速览轻量化Transformer各领域研究进展