随着Transformer架构在自然语言处理领域的广泛应用与显著成果，其对计算资源的需求也日益凸显。Transformer在语音识别、机器翻译、文档生成与摘要等任务上取得了卓越表现，这进一步推动了大规模语言模型（LLMs）时代的到来。然而，随着模型参数量的指数级增长，尤其是其中注意力机制所涉及的全序列交叉相关计算，训练成本和推理时的能耗显著增加，严重限制了模型的大规模应用‍

为解决这一问题，华为诺亚方舟实验室联合北京大学提出了一种名为“DiJiang”的创新方法，这是一种基于频域核化的策略，旨在将预训练的标准Transformer转化为具有线性复杂度的模型，从而极大地降低训练成本，还大幅提升了推理速度。

该方法的核心在于运用加权准蒙特卡洛采样技术，在理论上保证了更高效的近似效果，证明了逼近原始注意力机制的准确性。为了进一步减少训练过程中的计算复杂性，DiJiang采用离散余弦变换（DCT）操作作为核化基础

上图展示了DiJiang方案示意图，其中Transformer注意力机制中查询与键的计算被高效地映射到了频率域，利用快速离散余弦变换（DCT）完成这一映射过程。通过这种方法，有效地消除了softmax运算步骤，从而大幅度减少了Transformer模型的计算复杂度

实验结果与分析

研究人员在各种规模的语言模型上进行了广泛实验,验证了DiJiang架构在无需全量重训练的情况下,能够成功转化预训练的Transformer模型,并保持其原有的高水平表现

首先在Pythia模型上进行评估。在6个公共数据集上,DiJiang模型平均而言的表现几乎与原始Pythia模型持平,但训练成本仅为原先的1/16。与此同时,推理速度也显著提升,最高可达10倍。还在OPT、TinyLLaMA等其他模型上进行了微调,结果同样令人振奋

在LLaMA2-7B这一极具挑战的大规模预训练模型上,DiJiang-7B取得了与原模型相当的表现,但所需训练成本却只有后者的大约1/50。无疑,这一切都凸显了DiJiang架构在节省算力资源和减小碳足迹方面的巨大潜能

对比分析进一步揭示了DiJiang的优势所在。将其与现有的线性Transformer模型如Linformer、Performer等进行了对比。结果显示,这些方法在微调场景下表现不佳,而DiJiang得益于高效的注意力近似能够达到接近原始模型的性能。此外,它的收敛速度也明显快于其他方法,有望进一步节约训练时间和计算开销

研究人员通过可视化注意力图直观地展示了DiJiang对于注意力机制的高度保真逼近。这一视觉证据再次印证了DiJiang的卓越性能,突出了它对注意力计算关键信息的准确捕捉和建模能力,从而在产出整体性能时无须受损

上图展示了由不同方法生成的注意力分布图。显而易见，原始Transformer的注意力图（图（a）部分）富含信息，为其强大的功能奠定了基础。相比之下，像Performer（图（b）所示）这样的其他线性注意力方法产生的注意力图未能充分捕捉到词元之间的关系，导致生成的图与原始Transformer所呈现的图存在较大差异，即使经过微调，最终也会导致模型准确性下降。相反，DiJiang方法（图（c）所示）通过使用加权准蒙特卡洛方案，精确地逼近了原始注意力机制。这让模型能够有效地模拟不同词元之间的关系，尽管推理效率显著提高，但得到的结果几乎与原始Transformer相同

结语

DiJiang为提高预训练大规模语言模型的计算效率提供了一种全新的高效策略。它突破了过去的做法,避免了全量重训练的高昂消耗,却能够使模型以可控的代价实现线性化升级

研究人员还探索了一系列其他方法，如快速傅里叶变换混合tokens（FNet）、softmax替代方案（SoftMax-Free Transformer）及离散余弦变换图像建模（Discrete Cosin Transformer）等，这些均展示了学术界在提升Transformer效率方面的持续努力和多元创新思路

paper：https://arxiv.org/pdf/2403.19928.pdf

遥遥领先？