DiJiang: Efficient Large Language Models through Compact Kernelization

📄 arXiv: 2403.19928v2 📥 PDF

作者: Hanting Chen, Zhicheng Liu, Xutao Wang, Yuchuan Tian, Yunhe Wang

分类: cs.CL, cs.LG

发布日期: 2024-03-29 (更新: 2024-04-01)

🔗 代码/项目: GITHUB

💡 一句话要点

提出DiJiang以降低大语言模型的计算负担

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 频域核化 线性注意力 Transformer 大语言模型 计算效率 加权准蒙特卡洛 离散余弦变换

📋 核心要点

  1. 现有的Transformer模型在计算上负担沉重，尤其是在大规模语言模型中，改进注意力机制通常需要大量重新训练。
  2. DiJiang通过频域核化方法，将预训练的Transformer转变为线性复杂度模型，且训练成本极低，采用加权准蒙特卡洛方法进行高效采样。
  3. 实验结果显示，DiJiang-7B在多个基准测试中与LLaMA2-7B性能相当，但训练成本仅为其约1/50，推理速度显著提升。

📝 摘要（中文）

为了减少Transformer的计算负担，线性注意力机制的研究得到了显著关注。然而，现有的注意力机制改进策略通常需要大量的重新训练，这对于参数众多的大语言模型来说并不实际。本文提出了DiJiang，一种新颖的频域核化方法，能够以较低的训练成本将预训练的标准Transformer转变为线性复杂度模型。通过采用加权的准蒙特卡洛方法进行采样，所提方法在理论上提供了更优的近似效率。实验表明，DiJiang在性能上与原始Transformer相当，但训练成本显著降低，推理速度更快。

🔬 方法详解

问题定义：本文旨在解决大语言模型中Transformer的高计算负担问题，现有方法通常需要大量的重新训练，导致不适用于参数众多的模型。

核心思路：DiJiang通过频域核化技术，将标准Transformer转化为线性复杂度模型，减少训练成本，同时保持性能。采用加权的准蒙特卡洛方法进行高效采样，以提高近似效率。

技术框架：整体架构包括预训练的Transformer模型，通过离散余弦变换（DCT）进行核化，结合加权准蒙特卡洛方法进行采样，最终实现线性复杂度的注意力机制。

关键创新：DiJiang的主要创新在于其频域核化方法，能够在不需要大量重新训练的情况下，将Transformer模型的复杂度降低到线性级别，这与传统方法显著不同。

关键设计：在设计中，采用了离散余弦变换（DCT）作为核化基础，结合加权准蒙特卡洛方法进行高效采样，确保了模型在性能和计算效率上的平衡。具体的参数设置和损失函数设计在实验中进行了详细验证。

📊 实验亮点

实验结果表明，DiJiang-7B在多个基准测试中与LLaMA2-7B的性能相当，但其训练成本仅为后者的约1/50，推理速度显著提升，展示了其在计算效率上的优势。

🎯 应用场景

DiJiang的研究成果具有广泛的应用潜力，尤其是在需要高效推理和低计算成本的自然语言处理任务中。其方法可以被应用于实时翻译、对话系统以及大规模文本生成等领域，未来可能推动大语言模型在资源受限环境中的应用。

📄 摘要（原文）

In an effort to reduce the computational load of Transformers, research on linear attention has gained significant momentum. However, the improvement strategies for attention mechanisms typically necessitate extensive retraining, which is impractical for large language models with a vast array of parameters. In this paper, we present DiJiang, a novel Frequency Domain Kernelization approach that enables the transformation of a pre-trained vanilla Transformer into a linear complexity model with little training costs. By employing a weighted Quasi-Monte Carlo method for sampling, the proposed approach theoretically offers superior approximation efficiency. To further reduce the training computational complexity, our kernelization is based on Discrete Cosine Transform (DCT) operations. Extensive experiments demonstrate that the proposed method achieves comparable performance to the original Transformer, but with significantly reduced training costs and much faster inference speeds. Our DiJiang-7B achieves comparable performance with LLaMA2-7B on various benchmark while requires only about 1/50 training cost. Code is available at https://github.com/YuchuanTian/DiJiang.