WaveletGPT：基于小波的多尺度表征增强大型语言模型训练效率

2024-10-05 09:54:06

斯坦福大学的研究人员首次将小波理论应用于大型语言模型，提出了WaveletGPT，通过在Transformer解码器层中添加多尺度滤波器，加速了模型训练速度，并在文本、音频和音乐等多个领域取得了显著的性能提升。

论文介绍

大型语言模型 (LLM) 已经彻底改变了人工智能，影响了各个科学和工程学科。最初为机器翻译设计的 Transformer 架构已成为 GPT 模型的基础，极大地推进了该领域的发展。然而，当前的 LLM 在其训练方法上面临着挑战，该方法主要侧重于在保持因果关系的同时，根据先前上下文预测下一个 token。这种直接的方法已应用于各种领域，包括机器人技术、蛋白质序列、音频处理和视频分析。随着 LLM 规模的不断扩大，达到数千亿甚至数万亿个参数，人们对人工智能研究可及性的担忧也随之而来，有些人担心人工智能研究可能会局限于行业研究人员。研究人员正在解决的核心问题是如何增强模型能力以匹配更大的架构，或者用更少的训练步骤实现可比的性能，最终解决 LLM 开发中的规模和效率挑战。

研究人员探索了通过操纵中间嵌入来增强 LLM 性能的各种方法。一种方法涉及将手动调整的滤波器应用于潜在空间的离散余弦变换，用于在非因果架构（例如 BERT）中执行命名实体识别和主题建模等任务。但是，这种转换整个上下文长度的方法不适合因果语言建模任务。

两种值得注意的技术，FNet 和 WavSPA，试图改进类 BERT 架构中的注意力块。FNet 用二维 FFT 块取代了注意力机制，但这种操作是非因果的，因为它考虑了未来的 token。WavSPA 在小波空间中计算注意力，利用多分辨率变换来捕获长期依赖关系。但是，它也依赖于非因果操作，检查整个序列长度。

这些现有方法虽然具有创新性，但在其对 GPT 等仅解码器的因果架构的适用性方面存在局限性。它们经常违反对下一 token 预测任务至关重要的因果关系假设，使其不适合直接应用于类 GPT 模型。挑战仍然在于开发能够在保持仅解码器架构的因果性质的同时提高模型性能的技术。

来自斯坦福大学的研究人员提出了将小波融入 LLM 的第一个实例 WaveletGPT，通过将小波融入其架构来增强 LLM。这项技术被认为是同类中的首创，它使用 Haar 小波将多尺度滤波器添加到 Transformer 解码器层的中间嵌入中。这种创新允许每个下一 token 预测在每一层访问多尺度表示，而不是依赖于固定分辨率的表示。

值得注意的是，这种方法在不添加额外参数的情况下将基于 Transformer 的 LLM 的预训练速度提高了 40-60%，鉴于 Transformer 解码器架构在各种模态中的广泛使用，这是一项重大进步。该方法还展示了在相同训练步骤下性能的实质性提高，相当于增加了多层或参数。

基于小波的操作在三种不同模态中均显示出性能提升：语言（text-8）、原始音频（YoutubeMix）和符号音乐（MAESTRO），突出了其对结构化数据集的多功能性。此外，通过使小波核可学习（仅添加一小部分参数），模型实现了更大的性能提升，使其能够从头开始学习中间嵌入上的多尺度滤波器。

所提出的方法在保持因果关系假设的同时将小波融入到基于 Transformer 的大型语言模型中。这种方法可以应用于各种架构，包括非 Transformer 设置。该技术侧重于处理来自每个解码器层的中间嵌入。

对于给定信号 xl(i)（表示第 l 个解码器层沿第 i 个坐标的输出），该方法应用离散小波变换。使用 N+1 层和嵌入维度 E，此过程从解码器块之间的中间嵌入生成 N*E 个长度为 L（上下文长度）的信号。

小波变换，特别是使用 Haar 小波，涉及使信号通过具有不同分辨率的滤波器。Haar 小波是通过缩放和平移操作从母小波派生的方形函数。此过程创建了在不同时间尺度捕获信号信息的子小波。

离散小波变换是通过使信号通过低通和高通滤波器，然后进行下采样来实现的。对于 Haar 小波，这等效于平均和差分运算。该过程通过卷积和下采样生成逼近系数 (yapprox) 和细节系数 (ydetail)。对逼近系数递归执行此操作以获得多尺度表示，从而允许每个下一 token 预测访问中间嵌入的这些多分辨率表示。

这种方法通过关注逼近系数将小波和 LLM 嵌入联系起来，逼近系数捕获了不同级别的结构化数据。对于文本，这种结构的范围从字母到主题模型，而对于符号音乐，它的范围从音符到整个乐曲。该方法使用 Haar 小波，将过程简化为移动平均运算。为了保持因果关系和原始序列长度，该方法计算每个 token 维度内特定内核长度内的先前样本的移动平均值。这将创建输入信号的多尺度表示，从而允许模型跨嵌入维度捕获不同分辨率的信息，而无需更改中间 Transformer 嵌入的结构。

该方法引入了一种独特的方法来合并多尺度表示，而不会增加架构复杂性。它没有计算每个嵌入维度的所有级别的近似信号，而是通过嵌入维度本身的索引对级别进行参数化。这种方法保留了一半的中间嵌入信号不变，而根据它们的索引处理另一半。对于已处理的一半，一个简单的映射函数 f 确定每个坐标的内核大小，范围从级别 I 到 IX 近似值。修改后的信号 xnl(i) 使用因果移动平均滤波器计算，该滤波器的内核大小由 f(i) 确定。此操作维护了 LLM 中至关重要的因果关系假设，并防止信息从未来 token 泄漏。该技术创建了一种结构，其中不同的嵌入维度以不同的速率移动，从而允许模型捕获不同尺度的信息。这种多速率结构使注意力机制能够利用每一层和每个 token 的多尺度特征，从而潜在地增强模型捕获数据中复杂模式的能力。

跨三种模态（文本、符号音乐和音频波形）的结果表明，基于小波的中间操作可以显着提高性能。对于自然语言，验证损失的减少相当于在 text-8 数据集上从 16 层模型扩展到 64 层模型。就训练步骤而言，修改后的架构实现相同损失的速度几乎是原始架构的两倍。对于原始音频，这种加速更加明显，这可能是由于音频信号在短时间尺度上的准平稳性质。与 text-8 和符号音乐相比，原始波形 LLM 设置的收敛速度几乎快了两倍。

比较绝对时钟运行时间，修改后的架构在可学习和不可学习的设置中都显示出计算效率。报告了相对于基线架构完成一个 epoch 所需的时间。事实证明，该方法计算成本低，因为主要操作涉及 Haar 小波的简单平均或学习跨嵌入维度具有可变上下文长度的单个滤波器卷积核。这种效率与性能的提高相结合，突出了基于小波的方法在增强跨不同模态的 LLM 训练方面的有效性，而不会产生显着的计算开销。

本研究介绍了 WaveletGPT，介绍了小波（一种核心信号处理技术）与大型语言模型预训练的集成。通过向中间嵌入引入多尺度结构，性能速度提高了 40-60%，而无需添加任何额外参数。该技术在三种不同的模态中被证明是有效的：原始文本、符号音乐和原始音频。当训练相同的时间时，它表现出显着的性能提升。未来可能的方向包括结合来自小波和多分辨率信号处理的先进概念，以进一步优化大型语言模型。