在数字海洋中畅游：二进制令牌表示（BTR）如何加速检索增强语言模型

作者

在当今人工智能的浪潮中，语言模型（LM）已经成为了我们生活中不可或缺的一部分。它们在信息检索、文本生成和自动问答等领域展现了惊人的能力。然而，这些大型语言模型在实际应用中仍然面临着诸多挑战，例如“幻觉”（hallucinations）、信息陈旧以及隐私泄露等问题。为了应对这些挑战，研究人员们提出了检索增强语言模型（Retrieval-Augmented Language Models）。不过，这些模型在运行时速度较慢，难以扩展，因为它们需要处理大量检索到的文本。为了改变这一现状，来自华盛顿大学的研究团队推出了一种新的技术——二进制令牌表示（BTR）。

BTR的基本原理

BTR的核心思想是通过使用1位向量来预计算每个文本中的令牌表示，从而在推理时显著减少计算量。在BTR的架构中，令牌的二进制表示是通过对隐藏状态的校准二进制化来生成的，这种方法在下游任务（例如问答）中被证明是有效的。虽然这种表示方法可能会导致准确性的下降，但研究团队通过新的校准技术和训练目标有效地恢复了性能。此外，结合离线和运行时压缩，BTR在存储三十亿个维基百科令牌时，仅需127GB的磁盘空间。

文章中的公式

在BTR中，给定一个连续的令牌向量 $h_k = [h_1, h_2, \ldots, h_d]$，我们通过符号函数对其进行哈希，得到二进制表示向量 $b_k = sign(h_k)$。这里，$b_i$ 的值为1，当且仅当 $h_i > 0$，否则为-1。为了更好地保留表示质量，BTR采用了一种校准二进制化的方法，通过保存方差信息来恢复令牌的原始语义。

解决存储和计算的瓶颈

目前，很多现有的检索增强模型在推理时的计算瓶颈主要来自于阅读器（reader）组件。研究发现，在一台高性能的GPU服务器上，读取器的计算占总计算量的60%以上。而BTR通过预计算令牌表示来避免运行时的大量计算，从而实现了推理速度的显著提升。根据实验结果，BTR在五个知识密集型的自然语言处理任务中，将推理速度提升了2到4倍，同时在存储方面也减少了超过100倍，而任务性能仍然保持在95%以上。

训练过程中的创新

为了提高BTR的性能，研究团队在训练过程中引入了两个主要的目标。首先是“段落表示恢复目标”（passage representation recovery objective），该目标可以确保在二进制化之前，令牌表示能够保留段落的语义信息。其次是“查询感知段落令牌蒸馏目标”（query-aware passage token distillation objective），该目标旨在弥补由于预计算段落表示而导致的信息损失。

动态压缩技术的应用

在推理过程中，BTR还利用了动态压缩技术，进一步提高了推理效率。通过对检索到的相关段落进行内部和跨段落的压缩，BTR能够有效地减少计算量。例如，在读取器的上层，查询表示与段落表示进行连接后，BTR会合并相似的令牌，从而减少需要处理的令牌数量。

实验结果的验证

为了验证BTR的有效性，研究团队在五个知识密集型的自然语言处理任务中进行了大量实验，包括自然问题（Natural Questions）、TriviaQA、WebQA、事实核查（FEVER）和多任务语言理解（MMLU）等。实验结果显示，BTR相比于基线模型，不仅提高了推理速度，还在准确性上保持了良好的表现。例如，在Natural Questions任务中，BTR-Atlas base模型的准确率为49.5%，推理速度达到了3.1 QPS。

结论与未来展望

BTR的推出为检索增强语言模型的推理速度和存储效率带来了新的解决方案，使其在处理大量信息时更加高效。未来，研究团队计划将BTR扩展到解码器模型，并探索将二进制令牌表示应用于检索器的可能性，以构建更快、更强大的检索增强语言模型。

参考文献

  1. Cao, Q., Min, S., Wang, Y., & Hajishirzi, H. (2024). BTR: Binary Token Representations for Efficient Retrieval-Augmented Language Models. ICLR 2024.
  2. Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., … & Amodei, D. (2020). Language Models are Few-Shot Learners. NeurIPS 2020.
  3. Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models. arXiv preprint arXiv:2302.13971.
  4. Huang, H., et al. (2022). Privacy Concerns in Large Language Models. arXiv preprint arXiv:2201.04536.
  5. Dettmers, T., et al. (2022). Quantization of Language Models for Inference. NeurIPS 2022.

通过BTR，数字世界的浩瀚信息在我们指尖变得更加触手可及，未来的语言模型将在精准与速度之间找到更加完美的平衡。

发表评论