关于人工智能:辅助生成低延迟文本生成的新方向

大型语言模型现在风行一时，许多公司投入大量资源来扩大它们规模并解锁新性能。然而，作为注意力持续时间一直缩短的人类，咱们并不喜爱大模型迟缓的响应工夫。因为提早对于良好的用户体验至关重要，人们通常应用较小的模型来实现工作，只管它们的品质较低 (例如 代码补全工作)。

为什么文本生成这么慢？是什么阻止你在不破产的状况下部署低提早大型语言模型？在这篇博文中，咱们将从新扫视自回归文本生成的瓶颈，并介绍一种新的解码办法来解决提早问题。你会发现，通过应用咱们的新的辅助生成办法，你能够将硬件中的提早升高多达 10 倍！

了解文本生成提早

文本生成的外围很容易了解。让咱们看看外围局部 (即 ML 模型)，它的输出蕴含一个文本序列，其中包含到目前为止生成的文本，以及其余特定于模型的组件 (例如 Whisper 还有一个音频输出)。该模型承受输出并进行前向传递: 输出被喂入模型并一层一层程序传递，直到预测出下一个 token 的非标准化对数概率 (也称为 logits)。一个 token 可能蕴含整个词、子词，或者是单个字符，这取决于具体模型。如果你想深刻理解文本生成的原理，GPT-2 插图 是一个很好的参考。

模型的前向传递提供了下一个 token 的概率，你能够自在操作 (例如，将不须要的单词或序列的概率设置为 0)。文本生成的步骤就是从这些概率中抉择下一个 token。常见的策略包含抉择最有可能的 token (贪婪解码)，或从它们的散布中抽样 (多项式抽样)。在抉择了下一个 token 之后，咱们将模型前向传递与下一个 token 迭代地连接起来，持续生成文本。这个解释只是解码办法的冰山一角; 请参阅咱们 对于文本生成的博客 以进行深刻摸索。

从下面的形容中能够看出，文本生成的提早瓶颈很显著: 运行大型模型的前向传递很慢，你可能须要顺次执行数百次迭代。但让咱们深入探讨一下: 为什么前向传递速度慢？前向传递通常以矩阵乘法为主，通过查阅相应的 维基百科，你能够看出内存带宽是此操作的限度 (例如，从 GPU RAM 到 GPU 计算外围)。换句话说，_前向传递的瓶颈来自将模型权重加载到设施的计算外围中，而不是来自执行计算自身_。

目前，你能够摸索三个次要路径来充沛了解文本生成，所有这些路径都用于解决模型前向传递的性能问题。首先，对于特定硬件的模型优化。例如，如果你的设施可能与 Flash Attention 兼容，你能够应用它通能够过从新排序操作或 INT8 量化 来减速注意力层，其缩小了模型权重的大小。

其次，如果你有并发文本生成需要，你能够对输出进行批处理，从而实现较小的提早损失并大幅减少吞吐量。你能够将模型对于多个输出并行计算，这意味着你将在大致相同的内存带宽累赘状况下取得了更多 token。批处理的问题在于你须要额定的设施内存 (或在某处卸载内存)。你能够看到像 FlexGen 这样的我的项目以提早为代价来优化吞吐量。

# Example showcasing the impact of batched generation. Measurement device: RTX3090
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2").to("cuda")
inputs = tokenizer(["Hello world"], return_tensors="pt").to("cuda")

def print_tokens_per_second(batch_size):
    new_tokens = 100
    cumulative_time = 0

    # warmup
    model.generate(**inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size)

    for _ in range(10):
        start = time.time()
        model.generate(**inputs, do_sample=True, max_new_tokens=new_tokens, num_return_sequences=batch_size)
        cumulative_time += time.time() - start
    print(f"Tokens per second: {new_tokens * batch_size * 10 / cumulative_time:.1f}")

print_tokens_per_second(1) # Tokens per second: 418.3
print_tokens_per_second(64) # Tokens per second: 16266.2 (~39x more tokens per second)

最初，如果你有多个可用设施，你能够应用 Tensor 并行 调配工作负载并取得更低的提早。应用 Tensor 并行，你能够将内存带宽累赘摊派到多个设施上，但除了在多个设施运行计算的老本之外，你还须要思考设施间的通信瓶颈。该办法的收益在很大水平上取决于模型大小: 对于能够轻松在单个生产级设施上运行的模型，通常成果并不显著。依据这篇 DeepSpeed 博客，你会发现你能够将大小为 17B 的模型散布在 4 个 GPU 上，从而将提早缩小 1.5 倍 (图 7)。

这三种类型的改良能够串联应用，从而产生 高通量解决方案。然而，在利用特定于硬件的优化后，升高提早的办法无限——并且现有的办法很低廉。让咱们接下来解决这个问题！

从新回顾语言模型解码器的正向流传

上文咱们讲到，每个模型前向传递都会产生下一个 token 的概率，但这实际上是一个不残缺的形容。在文本生成期间，典型的迭代包含模型接管最新生成的 token 作为输出，加上所有其余先前输出的缓存外部计算，再返回下一个 token 得概率。缓存用于防止冗余计算，从而实现更快的前向传递，但它不是强制性的 (并且能够设置局部应用)。禁用缓存时，输出蕴含到目前为止生成的整个 token 序列，输入蕴含 _所有地位_的下一个 token 对应的概率分布！如果输出由前 N 个 token 组成，则第 N 个地位的输入对应于其下一个 token 的概率分布，并且该概率分布疏忽了序列中的所有后续 token。在贪婪解码的非凡状况下，如果你将生成的序列作为输出传递并将 argmax 运算符利用于生成的概率，你将取得生成的序列。

from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

inputs = tok(["The"], return_tensors="pt")
generated = model.generate(**inputs, do_sample=False, max_new_tokens=10)
forward_confirmation = model(generated).logits.argmax(-1)

# We exclude the opposing tips from each sequence: the forward pass returns
# the logits for the next token, so it is shifted by one position.
print(generated[:-1].tolist() == forward_confirmation[1:].tolist()) # True

这意味着你能够将模型前向传递用于不同的目标: 除了提供一些 token 来预测下一个标记外，你还能够将序列传递给模型并查看模型是否会生成雷同的序列 (或局部雷同序列)。

(请拜访浏览原文查看动静演示)

让咱们设想，你能够拜访一个神奇的无提早的预测辅助模型，该模型针对任何给定输出生成与你的模型雷同的序列。顺便说一句，这个模型不能间接用，只能辅助你的生成程序。应用上述属性，你能够应用此辅助模型获取候选输入 token，而后应用你的模型进行前向传递以确认它们的正确性。在这个乌托邦式的场景中，文本生成的提早将从 O(n) 缩小到 O(1)，其中生成的 token 数量为 n。对于须要屡次迭代生成的过程，咱们议论的是其数量级。

向事实迈出一步，咱们假如辅助模型失去了它的预测属性。依据你的模型，当初它是一个无提早模型，但它会弄错一些候选 token。因为工作的自回归性质，一旦辅助模型失去一个谬误的 token，所有后续候选 token 都必须有效。然而，你能够应用模型更正谬误 token 并重复反复此过程后再次查问辅助模型。即便辅助模型失败了几个 token，文本生成的提早也会比原始模式小得多。

显然，世界上没有无提早的辅助模型。然而，找到一个近似于模型的文本生成输入的其它模型绝对容易，例如通过相似训练的雷同架构的较小版本模型通常合乎此需要。当模型大小的差别变得显著时，应用较小的模型作为辅助模型的老本在跳过几个前向传递后就显得无关紧要了！当初，你理解了 辅助生成 的外围。

应用辅助模型的贪婪解码

辅助生成是一种均衡行为。你心愿辅助模型疾速生成候选序列，同时尽可能精确。如果辅助模型的品质很差，你将承当应用辅助模型的老本，而收益却很少甚至没有。另一方面，优化候选序列的品质可能意味着应用更慢的辅助模型，从而导致网络加速。尽管咱们无奈为你主动抉择辅助模型，但咱们蕴含了一个额定的要求和一个启发式办法，以确保模型与辅助模型一起破费的工夫放弃在可控范畴内。

首先，咱们要求辅助模型必须具备与你的模型完全相同的分词器。如果没有此要求，则必须增加低廉的 token 解码和从新编码步骤。此外，这些额定的步骤必须在 CPU 上进行，这反过来可能减少了设施间数据传输。可能疾速地应用辅助模型对于辅助生成的益处是至关重要的。

最初，启发式。至此，你可能曾经留神到电影盗梦空间和辅助生成之间的相似之处——毕竟你是在文本生成中运行文本生成。每个候选 token 有一个辅助模型前向流传，咱们晓得前向流传是低廉的。尽管你无奈提前晓得辅助模型将取得的 token 数量，但你能够跟踪此信息并应用它来限度向辅助模型申请的候选 token 数量——输入的某些局部比其它一些局部更容易被预计。

总结一下，这是咱们最后实现的辅助生成的循环 (代码):

1. 应用贪婪解码与辅助模型生成肯定数量的 候选 token。当第一次调用辅助生成时，生成的 候选 token 的数量被初始化为 5。
2. 应用咱们的模型，对 候选 token 进行前向计算，取得每个 token 对应的概率。
3. 应用 token 抉择办法 (应用.argmax() 进行贪婪搜寻或应用 .multinomial() 用于采样办法) 来从概率中选取 next_tokens。
4. 比拟步骤 3 中抉择的 next_tokens 和 候选 token 中雷同的 token 数量。请留神，咱们须要从左到右进行比拟，在第一次不匹配后，后续所有 候选 token都有效。5. 应用步骤 4 失去的匹配数量将 候选 token 宰割。也就是，将输出 tokens 加上刚刚验证失去的正确的 tokens。
5. 调整下一次迭代中生成的 候选 token 的数量 —— 应用启发式办法，如果步骤 3 中所有 token 都匹配，则 候选 token 的长度减少 2，否则缩小 1。

(请拜访浏览原文查看动静演示)

咱们在 🤗 Transformers 中设计了 API，因而应用该办法对你来说是无痛的。你须要做的就是将辅助模型作为 assistant_model 参数传入从而取得提早收益！咱们临时限度了辅助生成的批量大小为 1。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)
outputs = model.generate(**inputs, assistant_model=assistant_model)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
# ['Alice and Bob are sitting in a bar. Alice is drinking a beer and Bob is drinking a']

额定的外部复杂性是否值得？让咱们看一下贪婪解码状况下的提早数 (采样后果在下一节)。思考批量大小为 1，这些后果是间接从 🤗 Transformers 中提取的，没有任何额定的优化，因而你应该可能在你的设置中复现它们。

Space 体验地址: https://hf.co/spaces/joaogante/assisted_generation_benchmarks

通过观察收集到的数据，咱们发现辅助生成能够在不同的设置中显著缩小提早，但这不是灵丹妙药——你应该在利用之前对其进行零碎的评估以清晰应用该办法的代价。对于辅助生成办法，咱们能够得出结论:

1. 🤏 须要拜访至多比你的模型小一个数量级的辅助模型 (差别越大越好) ;
2. 🚀 在存在 INT8 的状况下取得高达 3 倍的减速，否则可能达到 2 倍的减速;
3. 🤯 如果你正在应用不适宜你的模型的 GPU 并且依赖于内存卸载的模型，你能够看到高达 10 倍的减速;
4. 📄 在输出驱动工作中大放异彩，例如主动语音辨认或摘要。

辅助生成的采样办法

贪婪解码实用于以输出为根底的工作 (主动语音辨认、翻译、摘要……) 或事实常识寻求。对于须要大量创造力的开放式工作，例如应用语言模型作为聊天机器人的大多数工作，应该改用采样办法。尽管辅助生成办法是为贪婪解码而设计的，但这并不意味着你不能应用多项式采样进行辅助生成！

从 next token 的概率分布中抽取样本将导致咱们的基于贪婪的辅助生产更频繁地失败，从而升高其提早劣势。然而，咱们能够应用采样中的温度系数来管制下一个标记的概率分布有多尖利。在一种极其状况下，当温度靠近 0 时，采样将近似于贪婪解码，有利于最有可能的 token。在另一个极其，当温度设置为远大于 1 的值时，采样将是凌乱的，从均匀分布中抽取。因而，高温对你的辅助模型更无利，可能保留辅助生成的大部分提早劣势，如下所示。

无妨亲眼看一看，感受一下辅助生成的魅力？

Space 体验地址: https://hf.co/spaces/joaogante/assisted_generation_demo

将来倒退方向

辅助生成表明以后文本生成策略曾经到了可优化的阶段。咱们意识到它目前的难点不在于计算量的问题，因而能够利用简略的启发式办法来充分利用可用的内存带宽，缓解瓶颈。咱们置信，进一步优化辅助模型将使咱们取得更大的提早升高——例如，如果咱们申请辅助模型生成多个间断候选 token，咱们可能可能跳过更多的前向传递。天然地，应用高质量的小模型作为辅助模型对于实现和扩充收益至关重要。

该办法最后在咱们的 🤗 Transformers 库下公布，用于 .generate() 函数，咱们预期将其纳入整个 Hugging Face 宇宙。它的实现也是齐全开源的。因而，如果你正在进行文本生成而没有应用咱们的工具，你能够随时将其作为参考。

最初，辅助生成从新提出了文本生成中的一个关键问题: 模型中所有新 token 都是给定模型以自回归形式计算的后果，同质地前向传递每一个 token。这篇博文提出了这样的想法: 生成的大部分序列也能够由小尺寸的模型同样生成。为此，咱们须要新的模型架构和解码办法——咱们很快乐看到将来会带来什么！

相干工作

在这篇博文最后公布后，我留神到其余作品也摸索了雷同的外围准则 (应用前向传递来验证更长的连续性)。特地地，请看以下作品:

• 分块并行解码, 来自 Google Brain
• 揣测性采样, 来自 DeepMind

Citation

@misc {gante2023assisted,
    author = {{Joao Gante} },
    title = {Assisted Generation: a new direction toward low-latency text generation},
    year = 2023,
    url = {https://huggingface.co/blog/assisted-generation},
    doi = {10.57967/hf/0638},
    publisher = {Hugging Face Blog}
}

致谢

我要感激 Sylvain Gugger、Nicolas Patry 和 Lewis Tunstall 分享了许多贵重的倡议来改良这篇博文。最初，感激 Chunte Lee 设计了精美的封面，你能够在咱们的网页上看到。

原文链接: https://huggingface.co/blog/assisted-generation

作者: Joao Gante

译者: gxy-gxy

排版 / 审校: zhongdongy (阿东)