用于快速LLM推理的Token合并：背景与Mistral的首次试验

背景

仅解码器模型通过其卓越的生成能力彻底改变了自然语言处理任务。生成新词元的过程是暴力破解式的，因为它基本上是预测一个词元，将其附加到序列中，然后重复此过程，直到达到最大大小或生成了一个词元。当然，这个过程的核心是Transformer架构，随着长度的增加，为了执行自注意力机制，它需要越来越多的计算资源。

我们如何改进这个过程，在不损失预测质量的情况下，使推理更快？

• 我们可以执行更简单的网络操作。这就是量化（8位、4位、三元...）。

• 我们可以拥有一个更小的模型。论文《深层的不合理无效性》研究了从模型中删除层的影响。结果表明，删除模型中最后的注意力模块不会显著影响模型性能，同时能改善延迟。

• “自然而然”出现的问题是：我们需要完整的词元序列来预测下一个词元吗？或者我们可以将其压缩/合并以输入一个缩减的序列？

目标是扭曲大型语言模型的前向传播，以便在调用中合并序列一次或多次。这应该在不需要重新训练或微调的情况下完成。一个简单的模块可以优化这个过程。

为了支持我的观点，“Diffusion世界”已经有一个库，用于在前向传播中平均冗余词元以加速延迟。

提议的技术

因此，我正在寻求将多个词元嵌入向量平均化。简单的平均化不适合，因为它倾向于失去平均向量的幅度。mergekit库中完成的工作引入了一种相对不为人知的平均技术，称为SLERP（球面线性插值）。它通过保留两个向量的球面特性来进行插值。不幸的是，SLERP方法只适用于一对向量，这要求我预处理序列。

合并过程如下：

• 检查序列长度
  • 如果为奇数，则在开头和结尾添加一个“NULL”词元（所有值设为0）
  • 如果为偶数，则在开头、结尾和倒数第二个位置插入一个空词元（所有值设为0）
  • 现在序列长度始终为奇数
• 将序列从（批大小，序列长度，维度）重塑为（批大小，序列长度/2，2，维度）
• 生成一个用于插值的温度数组，形状为（批大小，序列长度/2），所有值设为0.5
  • 对于包含空词元的对，将温度更改为0或1，以完全保留非空词元
• 应用成对SLERP

获得的新序列大小几乎减半，然后合并后将其输入到其他层。空词元背后的直觉是，词元在共享信息（一个汇聚词元）中扮演着核心角色，并且为了语法目的需要完全关注最后一个词元。

这是表示该过程的方案：

一个非常简单的实现，你可以在语言建模头之前合并词元

from transformers import AutoModelForCausalLM, AutoTokenizer
from forward_slerp import merge_tokens

mistral = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")

sentence = "[INST] What is the biggest challenge in life ? [/INST]"

tokens = tokenizer(sentence, return_tensors = "pt")

hidden_state = mistral.model(**tokens)
merged = merge_tokens(hidden_state)

preds = mistral.lm_head(merged)

合并基准测试

我创建了一个Mistral Instruct 7B v0.2模型的合并推理版本。测试在一块H100 GPU上进行。

这个合并过程的评估包括三个部分：合并前向推理与未合并前向推理预测的下一个词元相似性，估算的延迟加速（可以改进），以及Alpaca基准测试。

在前两个部分，我在模型层的不同级别（称为“层截止”）以及不同长度的输入序列上测试了合并过程。这些序列是随机从CNN数据集中抽取的5k文本。

在Alpaca基准测试中，我在第20层应用了合并过程，并解码到第4096个词元。

下一词预测

不出所料（或许你有所期待，我不知道），合并预测和基础预测之间的准确性随着合并应用时机的推进而线性增长。然而，即使在较低级别，合并模型也能通过快速达到与基础模型80%的一致性来获得良好的结果。序列长度似乎没有出现任何模式。

合并推理和未合并推理之间的准确性

当检查前3名和前5名的准确率时，结果甚至更好。非常短的序列结果往往更差，但这可能不确定，因为文本开头是由元数据组成的（例如，“LONDON, England (Reuters) --”）。

合并推理和未合并推理之间的前3名准确率

合并推理和未合并推理之间的前5名准确率

估算延迟加速

由于我尚未实现kv缓存管理，延迟加速可能未达到最佳优化。

延迟加速（y轴：合并推理比未合并推理快多少倍）

目前这些结果对我来说似乎不够稳健，因为我原预计在第一层会有显著更高的加速。然而，长序列从这种方法中受益最大。

AlpacaEval 2.0 基准测试

为了评估模型在推理时合并（层截止设置为20）后的表现好坏，我将合并版本与原始版本在AlpacaEval基准测试（805个提示）上进行了比较。排行榜和注释可在仓库的“code/alpaca”文件夹中找到。

可以提出两点评论：

• 它并没有过度影响输出质量，因为合并模型失去了4%的胜率（未控制长度）和7%的胜率（控制长度）。尽管这并非完全可靠，但它表明模型仍然“活跃”，并且与基础模型相差不远。
• 词元合并增加了模型的冗长性，因为输出的平均长度增加了600个词元。这源于输入到层的序列较短的影响，导致不同的位置编码倾向于延迟eos_token的出现。我未能实现一致的位置编码，但未来会这样做。

该模型仍然优于gemma-7b-it、text_davinci_001或nous-hermes-13b，同时平均了每对词元。它在145个测试模型中排名第88位。

结论与未来工作

正如我提到的，这项工作仅在Mistral 7B模型上进行。这种合并思想在实施方面可能因模型的底层架构而异。我还收到了关于这项技术是否能扩展到“大海捞针”效应不那么明显的更大模型的问题。我将对不同的模型重复这项工作，以查看其效果如何。

此外，由于代码尚未处理kv缓存，合并代码可能未针对实现超快速前向传播进行充分优化。一个主要课题是处理位置编码，以限制“过度生成”效应。

我打算构建此技术的更高级版本，最终在HuggingFace中为任何因果大语言模型构建一个包装器类，以实现更快的推理（类似于accelerate）。

最后，这项工作概述了LLM需要双重架构世界的需求：一个用于训练，一个用于生成。

我很高兴听到您的反馈和意见，并期待我们能进一步推进这项工作。

链接

仓库：https://github.com/samchaineau/llm_slerp_generation HF账户：https://huggingface.co/samchain