摘要:在大型语言模型(LLMs)迅速发展的今天,一项来自EPFL(瑞士联邦理工学院洛桑分校)、东北大学和微软联合研究团队的创新成果正在改变我们对模型效率的认识。这篇名为"zip2zip: 通过令牌压缩实现大语言模型的推理时自适应词汇表"的研究由Saibo Geng、
想象一下,如果你只会说单个词,而不能组合词语成句,交流会变得多么繁琐。这正是当前大语言模型面临的一个问题。它们使用固定的"词汇表"——一组预先定义的单词或子词(称为"tokens")。这就像给模型一本固定的字典,不管它在讨论医学、编程还是法语文学,都只能使用这本字典里的词汇。当模型需要表达特定领域的术语或非英语语言时,常常需要用多个基本tokens拼凑而成,这就像用积木一块块搭建复杂的结构,既耗时又低效。
研究团队提出的zip2zip框架就像是给模型提供了一个动态扩展的字典,能够在对话过程中即时学习和创建新的"复合词"(hypertokens)。这些复合词由多个基本Tokens组成,但作为整体处理,大大提高了处理效率。想象一个人在交谈中发现自己经常重复使用某些词组,于是决定创造一个简短的缩写来代替——这正是zip2zip的工作原理。
zip2zip的核心是基于Lempel-Ziv-Welch(LZW)压缩算法,这也是常见的zip文件压缩技术所使用的算法。它能在模型推理(即模型生成回答)的过程中,动态识别并合并经常一起出现的tokens,形成新的超级tokens(hypertokens)。这种方法显著减少了处理相同内容所需的tokens数量,从而提高了效率。
研究表明,使用zip2zip可以将输入和输出序列长度减少20-60%,大幅提高推理速度。更令人惊喜的是,研究团队证明现有的LLM可以通过参数高效微调在短短10个GPU小时内适应这一框架,无需重新训练整个模型。
这项技术对于提高大语言模型的实用性具有深远意义。它能让模型在相同的计算资源下处理更长的文本,降低运行成本,并提供更快的响应速度。特别是对于专业领域(如医学、法律或编程)或非英语语言的应用,zip2zip提供了显著的性能提升。
让我们深入了解这项技术的工作原理、实现方式以及实验结果,看看它如何改变大语言模型的运行方式。
一、理解问题:为什么当前的分词方式不够高效?
在深入了解zip2zip之前,我们需要理解为什么当前大型语言模型的分词方式存在效率问题。当你使用ChatGPT或其他大语言模型时,它实际上并不是直接处理我们的文字,而是先将文字切分成一个个"tokens"。这些tokens可能是完整的单词、常用词组或单词的一部分。
举个简单的例子,当模型处理"machine learning"这个词组时,它可能将其分割为["machine", "learning"]两个tokens;但当处理更专业的词汇如"transcription factor"时,可能会被分割为["trans", "cri", "ption", "factor"]四个tokens。这就是为什么有时你会发现大语言模型在处理特定领域的专业术语或非英语文本时显得不够高效。
研究者们指出,当前大多数LLM使用的是在通用文本语料库上训练的静态词汇表,通常采用字节对编码(Byte Pair Encoding,BPE)等算法生成。虽然这种全局优化的词汇表在平均情况下表现尚可,但在处理特定领域或语言时却不尽如人意。
比如,当处理生物医学文本时,专业术语如"transcription factor"或"mitochondrial DNA"可能被分割成多个tokens,导致表示同一内容需要的tokens数量增加2-3倍。这不仅增加了内存和计算需求,还增加了用户成本。
研究团队发现,目前业界尝试解决这个问题的方法主要是在特定领域或语言适应时扩展词汇表。例如,为中文、韩文或特定专业领域(如心理健康、法律)定制专门的词汇表。然而,这种方法需要为每个目标领域或语言重复进行,并且需要维护多个分词器。
另一种趋势是直接增加词汇表的大小——从32K增加到128K甚至200K tokens。但研究表明,简单扩大词汇表的收益递减,超过一定规模后可能反而会降低模型性能。
这些局限性表明,我们需要一种自适应的分词机制——一种能够在推理时动态调整词汇表以适应输入文本的机制,而无需重新训练模型或维护多个分词器。这正是zip2zip所要解决的核心问题。
二、zip2zip的核心理念:动态词汇表如何工作?
zip2zip的革命性理念在于它能让语言模型在对话进行时实时调整词汇表,就像我们人类在交谈中会即兴创造简称或缩写一样。这种动态适应能力无需事先重新训练模型,也不需要为每种语言或每个专业领域准备单独的词汇表。
为了实现这一目标,研究团队将每个LLM与一个"超级词汇表"(hyper-vocabulary)关联起来,用于补充模型的原始静态词汇表。原始词汇表中的tokens被称为"基础tokens",而超级词汇表中的每个条目都是一个"超级token",代表基础tokens的合并序列。
在每次对话开始时,超级词汇表初始为空集,然后在解码过程中通过识别和合并上下文窗口中重复出现的token子序列而逐步填充。这就像是模型在交谈过程中不断学习和创建新的"词组",使后续的交流更加高效。
zip2zip的持续词汇扩展过程是其核心特性。随着解码的进行,zip2zip会持续合并共同出现的tokens作为新的超级tokens添加到超级词汇表中,并在新生成的tokens上反复应用合并操作。这种持续扩展允许模型以更紧凑的方式表示较长的、重复出现的基础tokens序列。
超级tokens被视为模型中的一等公民,在整个解码过程中与基础tokens可互换使用。重要的是,这个过程完全发生在推理过程中,无需修改底层分词器或要求模型重新训练。
研究团队选择Lempel-Ziv-Welch(LZW)压缩算法作为实现词汇扩展的核心技术。LZW是一种基于字典的无损压缩方法,可以增量构建变长序列的码本。在zip2zip中,码本初始化为基础token词汇表,并在遇到重复的token模式时通过添加新的超级tokens进行扩展。
为了控制动态扩展词汇表的增长,研究者们设置了一个最大合并大小M,限制单个超级token能代表的基础tokens数量。LZW算法特别适合zip2zip,主要因为它具有以下特性:
1. 它是流式的——在步骤t创建的超级tokens可以在步骤t+1立即重用;相比之下,像BPE这样的方法需要访问完整序列并离线操作。
2. 它是自包含的——压缩后的token序列可以完美地重建输入的基础tokens,无需传输或保存码本。
3. 它是无歧义的——当基础tokens和超级tokens都可用时,LZW算法会明确决定使用哪一个,无需人为干预。
这些特性使得zip2zip能够在不改变现有LLM库和接口的情况下,实现动态、自适应的分词机制,为语言模型带来显著的效率提升。
三、技术实现:超级嵌入和架构设计
实现动态词汇表面临一个关键挑战:超级tokens没有在原始模型的嵌入层(和投影层)中预定义的固定嵌入向量,因为它们不是原始词汇表的一部分。zip2zip通过引入一个巧妙的机制来解决这个问题。
研究团队引入了一个"超级编码器"(hyper-encoder),这是一个神经网络,接收组成超级token的基础tokens的嵌入作为输入,并输出相应的超级token嵌入。具体来说,对于一个由M个基础tokens y1:M组成的序列,超级编码器fφ产生超级token嵌入h = fφ(y1:M)。对于由少于M个基础tokens组成的超级tokens,研究者们会将输入序列填充到长度M。
由于基础tokens的嵌入映射保持不变,超级编码器fφ本质上是将拼接的基础token嵌入从(M × d)维空间映射到d维的超级token嵌入向量,执行非线性降维。这就像是将多个单词的含义压缩成一个新词的含义,保留原始信息但形式更紧凑。
对于输出投影层,如果底层transformer绑定了嵌入和投影矩阵,可以重用相同的超级编码器来计算用于投影的表示。否则,会训练一个单独的超级编码器来生成超级token的投影向量。
zip2zip的整体架构(如图2所示)包含以下几个关键步骤:
1. 输入文本首先被分词为基础tokens(步骤1)。
2. 然后通过在线LZW压缩模块将token序列压缩成超级tokens流(步骤2)。
3. 由于超级tokens不是模型原始嵌入层的一部分,它们的嵌入向量在推理过程中由超级编码器即时计算(步骤3-4)。
4. 嵌入后,基础token嵌入和超级tokens嵌入都通过基础模型的标准transformer层处理,产生上下文化的隐藏状态(步骤5-6)。
5. 在输出投影层,超级token投影向量(在绑定情况下与超级token嵌入向量相同,否则由单独的超级编码器计算)被附加到原始投影矩阵中(步骤7)。
6. 这允许模型在基础词汇表和超级词汇表的并集V∪Vh上计算联合softmax(步骤8)。
7. 结果概率分布覆盖V∪Vh,采样的token可以是基础token或超级token(步骤9)。
8. 在下一个循环中,新生成的token(步骤10)——无论是基础还是超级——都被附加到输入序列中,过程重复(回到步骤1)。
9. 在生成结束时,超级token序列通过LZW解码函数解压缩为基础tokens序列(步骤11-12)。
整个过程以完全自回归的方式工作,新生成的tokens也会被合并成供未来步骤使用的超级tokens。
研究团队特别强调了两个重要点:
1. 一致的词汇表更新:扩展的词汇表(包含新创建的超级tokens)必须在输入嵌入层和输出投影层之间以一致的方式更新,维持超级token集的一致视图。未能一致更新两侧可能导致两类错误:无法解码的超级tokens,或模型尝试解码不存在的超级token。
2. 超级嵌入缓存:虽然超级token嵌入是即时计算的,但它们是上下文无关的,因此可以在推理步骤之间缓存。类似于transformer的KV缓存,这使得增量更新成为可能:在每个步骤只需处理新创建的超级tokens。由于码本随着上下文中的tokens数量n线性增长,总缓存大小也线性增长。因此,超级token嵌入的计算成本在每一步保持恒定——即每步计算一个token嵌入。
这种精心设计的架构使zip2zip能够在不修改基础模型架构的情况下,实现动态词汇表扩展,显著提高语言模型的效率。
四、训练方法:如何教会模型使用动态词汇表
要让语言模型能够有效使用动态扩展的词汇表,研究团队设计了一种特殊的训练方法。传统的语言模型训练使用因果语言建模(CLM)目标,即最小化数据序列在模型下的负对数概率:
min(θ) E[y~D][-log πθ(y)]
其中πθ(y)表示模型πθ下token序列y的概率。
在zip2zip中,研究团队的目标是在压缩序列上优化相同的CLM损失。设C为在线压缩算法(如LZW),φ为超级编码器的参数。给定序列y~D,设z = C(y)为其压缩形式。训练目标变为:
min(θ,φ) E[y~D][-log πθ,φ(C(y))] = min(θ,φ) E[z~C(D)][-log πθ,φ(z)]
这里,πθ,φ(z)表示压缩序列z的概率,由基础模型权重θ和超级编码器参数φ参数化。
为了构建压缩数据集C(D),研究团队首先使用标准分词器对语料库进行分词,然后应用LZW压缩算法。这个预处理步骤在训练前执行一次,可以通过批处理进行高效并行化。
虽然超级tokens引入了额外的词汇动态,但训练仍然是完全可并行的。研究团队利用语言模型中使用的标准因果掩码机制,允许模型并行预测每个位置的下一个token——无论是基础token还是超级token。为了在推理过程中消除对序列码本更新的需求,研究团队通过对整个输入序列应用LZW压缩来预先计算固定的码本。这个预先计算的码本在训练过程中一致使用,以条件化token预测,确保效率和与标准训练管道的兼容性。
研究团队还引入了一个辅助重建目标,鼓励超级token嵌入保留有关其底层基础token序列的足够信息。具体来说,模型被训练为从超级token嵌入重建原始基础token嵌入。他们使用结合了标准下一个token预测损失和辅助重建损失的组合损失来联合优化语言模型和超级编码器:
min(θ,φ,ψ) E[y~D][-log πθ,φ(C(y))] + λ E[y1:M][Δ(y1:M, fψ(fφ(y1:M)))]
其中fφ是超级编码器,fψ是旨在从其超级嵌入重建相应基础tokens的解码器,Δ是基础tokens y1:M与重建基础tokens fψ(fφ(y1:M))之间的重建损失函数,如交叉熵损失。超参数λ控制语言模型预测误差与自编码器重建误差之间的权衡。
这种联合优化目标鼓励超级编码器学习一个紧凑的d维流形,嵌入在更高维的(M × d)空间的基础token嵌入中,而语言模型πθ,φ学习在给定前面上下文的情况下预测下一个(超级)token。重建损失可以视为一种自编码形式,其中超级token充当压缩的潜在表示,重建鼓励语义内容的保存和压缩的无损性。
适配预训练语言模型。提出的目标(方程2、3)自然地与预训练语言模型集成。在这种设置中,基础模型可以保持冻结,只训练超级编码器以适应压缩的token序列。也可以使用LoRA等参数高效方法来适配基础模型的选定组件,实现使用最少计算资源的有效适配。
五、效率优势:zip2zip如何加速推理过程
zip2zip通过增加平均token长度来提高效率,从而减少表示相同文本所需的tokens数量。这种压缩适用于输入(如提示)和输出(如完成),导致上下文长度更短。因此,模型执行更少的计算——无论是在注意力机制中还是前馈层中——更重要的是,在推理过程中需要更少的自回归解码步骤。
由于大型语言模型的延迟主要由顺序解码的成本驱动,将输出tokens数量减少n%会导致解码延迟大约减少n%,这一点在实验部分得到了经验证明。
不过,需要注意的是,zip2zip引入了一些额外的计算开销,主要来自超级编码器在推理时即时计算超级嵌入的过程。为了解决这一问题,研究团队实现了超级嵌入缓存并使用自定义Triton内核优化计算。正如实验结果所示,这些优化措施使得zip2zip能够在保持高质量输出的同时,显著提高推理效率。
六、实验结果:zip2zip在实际应用中的表现
为了评估zip2zip的有效性,研究团队对Phi-3模型(3B和14B)进行了适配,并从四个维度进行了评估:(1)token效率,(2)语言建模困惑度,(3)下游任务性能,(4)推理效率。
### 训练设置
研究团队没有更新完整的模型权重,而是采用了参数高效微调技术LoRA。此外,他们训练了超级嵌入和超级投影模块。他们将最大合并大小设置为M = 3,并使用两层transformer编码器作为超级编码器。
微调过程非常高效,4B参数模型只需约10个H100-GPU小时,14B参数模型需要40个H100-GPU小时,仅使用0.1亿训练tokens。有趣的是,重建损失在训练期间几乎收敛到零,表明模型几乎可以完美地从超级token表示中恢复原始基础token序列,这突显了所学压缩的高度信息保存性。
### 样本输出和超级token模式
研究团队展示了几个例子,提供了zip2zip模型如何生成文本的直观理解。从示例中可以看出,模型成功地在输出中生成了基础tokens和超级tokens的混合。超级token比例在Python代码生成示例中高达40%,在生物医学文本生成示例中为20%。许多超级tokens对应于有意义的语义单位或特定领域术语,如"torch"、"Attention"、"MultiHead"(代码生成);"mRNA"、"transcribed"、"synthesis"(生物医学);"Eiffel"、"de la"、"Gustave"(法语)。
### Token效率
研究团队定义token效率η为每个token表示的平均字节数,即η = Bytes(x) / Tokens(x),其中Bytes(x)是文本x的UTF-8编码中的字节数。他们使用四种LLM的分词器评估了token效率——Llama-3(128K词汇量)、Qwen-2(150K)、Phi-4(200K)和Gemma-3(256K)——每种都有不同的基础词汇量大小。token效率在五个代表性领域进行测量:代码、数学、聊天、多语言和网络。
结果显示,应用LZW zip2zip一致地提高了所有分词器和领域的token效率。在代码和数学等结构化领域的收益特别显著——比基础分词器高出50%。有趣的是,词汇量更大的模型并不总是实现更好的token效率,这表明简单地增加词汇量并不足以提高效率。
### 困惑度
研究团队在四个语料库上评估了zip2zip模型的困惑度:Wikitext、the Pile以及Paloma的两个子集:mC4(C4的多语言子集)和dC4(又称C4-100D,跨越100个领域的C4子集)。
结果表明,zip2zip模型的字节级困惑度有适度增加,表明语言建模性能有所下降。这是可以预期的,因为压缩通常会导致预测性和模型性能之间的权衡。不过,困惑度的增加相对较小,特别是考虑到zip2zip带来的显著效率提升。
### NLP基准测试
研究团队在七个广泛使用的NLP基准上评估了zip2zip的性能:ARC-[Challenge, Easy]、HellaSwag、LAMBADA、OpenbookQA、PIQA、Winogrande和GSM8K。
结果表明,使用zip2zip微调的模型在大多数任务上的表现与基线相似。然而,在GSM8K上,主要涉及数值计算的任务,模型表现出显著下降。由于这类任务对分词敏感,模型偶尔会生成格式错误或重复的数字。虽然token级操作对LLM已知具有挑战性,但自适应分词似乎加剧了这个问题。
为了验证zip2zip在非英语语言上的有效性,研究团队在机器翻译任务上评估了模型,包括WMT14和WMT16。结果表明,使用zip2zip时,BLEU、CHRF和TER指标在所有测试的语言对上都有小幅度下降。然而,下降相对较小,表明即使在压缩表示中,模型仍然保持强大的多语言能力。
### 推理效率
zip2zip通过降低需要生成的tokens数量来减少解码时间。然而,它引入了额外的计算量,主要来自超级编码器在推理时即时计算超级嵌入的过程。为了解决这一开销,研究团队实现了超级嵌入缓存并使用自定义Triton内核优化计算。
实验结果显示,zip2zip在所有四种测试设置中都实现了显著的速度提升,无论是在预填充阶段还是解码阶段。最显著的收益观察到在使用Phi-3.5-4B模型的512+256设置中。在数据中心级GPU(如NVIDIA H100)上的改进更为显著,而在消费级硬件(如Apple M1)上则相对适中。
研究团队还专注于高效的LZW分词实现。zip2zip在推理过程中引入了额外的LZW压缩步骤,以及生成结束时的解压缩步骤。为了最小化开销,他们实现了一个基于Rust的zip2zip分词器,其性能优于Python版本,并与HuggingFace的快速BPE分词器延迟相当。
七、与相关工作的比较
zip2zip与现有研究的关系主要体现在以下几个方面:
### 词汇扩展
先前的工作探索了在特定领域或语言适应期间扩展分词器词汇表以提高分词效率。例如,有研究通过添加新tokens来适配LLaMA模型以支持中文、韩文和专门领域如心理健康和法律。还有研究者探索了如何通过更好地选择要添加的tokens子集来有效扩展词汇表。
与这些方法不同,zip2zip是首个实现推理时动态词汇扩展的方法,基于输入上下文构建新tokens,无需预先重新训练或修改分词器。
### 提示压缩
提示压缩方法包括GistTokens、Selective Context、LLMLingua、Summary Vectors、In-context Autoencoder等,这些方法减少了输入token长度,但不影响输出tokens数量,而输出生成通常主导整体生成时间。相比之下,zip2zip同时压缩输入和输出token序列。
### 潜在tokens表示
潜在token表示或patches的概念主要在计算机视觉中探索,如Token Merging和Token Pooling,旨在减少序列长度同时保留语义内容。最近,Byte Latent Transformer (BLT)将这一概念扩展到语言建模,完全放弃tokens直接在字节级别操作。
BLT和zip2zip都采用输入的层次建模方法,但在三个关键方面有所不同:(1)目标:BLT旨在替代分词器,而zip2zip寻求扩展和改进它;(2)算法:BLT使用基于熵的分割,而zip2zip应用基于LZW的token压缩;(3)训练:BLT需要从头开始训练,而zip2zip支持预训练模型的持续适应。
此外,Lester等人提出通过直接在用算术编码压缩的文本上训练LLM来提高语言模型效率。虽然两种方法都利用压缩来提高效率,但zip2zip强调动态词汇扩展以实现现有模型的升级训练,而Lester等人的方法需要从头开始训练。
八、讨论与局限性
尽管zip2zip取得了令人印象深刻的成果,研究团队也坦率地讨论了该方法的几个局限性和未来研究方向:
### 超越LZW
虽然研究团队采用LZW进行超级tokens的动态构建,但zip2zip广泛兼容任何在线压缩算法。未来的工作可能探索提供压缩效率和模型性能之间不同权衡的替代方案。
### 码本管理策略
LZW算法使码本随着上下文窗口中的tokens数量线性增长。实证结果显示,在生成过程中只有约25%的超级tokens被重用,留下了优化的巨大空间。两个潜在的改进是:(1)修剪或选择性保留策略以减少未使用的条目,以及(2)码本预填充,如果能提前预测可能的tokens,这可能会有所帮助。
### 压缩-质量权衡
压缩和建模之间存在内在权衡:随着token空间被更积极地压缩,冗余减少——但可预测性也随之降低——使模型更难预测下一个(超级)token。在极端情况下,最优压缩方案如算术编码产生的序列在统计上与随机噪声无法区分,使语言模型难以学习。
研究者们观察到这种效应表现为更高压缩级别下困惑度增加(表3),这可能会通过降低生成质量而削弱压缩的好处(尽管在表4和表5中的任务中相对较小)。在压缩和模型性能之间找到正确平衡仍然是未来研究的重要方向。
九、结论
zip2zip框架为大型语言模型带来了一种革命性的动态词汇适应方法。通过整合基于LZW的token压缩与动态超级token嵌入机制,zip2zip实现了序列长度和解码步骤的显著减少,从而提高了推理效率,同时只需对架构进行最小修改。
实验表明,zip2zip在保持各种任务上的强大性能的同时,实现了显著的推理效率提升。这些发现突显了将动态分词集成到LLMs中的潜力,为LLM效率研究开辟了新方向。
最重要的是,zip2zip提供了一种实用的解决方案,可以提高大型语言模型在各种应用场景中的效率和响应速度,特别是在处理特定领域的内容或非英语语言时。随着大型语言模型继续在各行各业的应用扩展,这种提高效率的方法将变得越来越重要,使得先进的AI技术更加经济实惠和广泛可用。
来源:至顶网一点号