BTR: Binary Token Representations for Efficient Retrieval Augmented Language Models

TL;DR

This paper addresses the bottleneck of slow inference in retrieval-augmented language models by introducing Binary Token Representations (BTR), which cache 1-bit token vectors for retrieved passages and apply calibration and compression to maintain accuracy. BTR employs calibrated binarization, offline and runtime token compression, and two training objectives to preserve semantic information and salient query-related details. The approach yields 2–4x faster inference and up to 101x storage reduction across five knowledge-intensive tasks while retaining 90–95% of the original performance, and it also proves effective for encoder-only readers. The work demonstrates that aggressive token-level compression, when properly trained and calibrated, can deliver substantial practical benefits for scalable retrieval-augmented systems with broad applicability and potential for further extension.

Abstract

Retrieval augmentation addresses many critical problems in large language models such as hallucination, staleness, and privacy leaks. However, running retrieval-augmented language models (LMs) is slow and difficult to scale due to processing large amounts of retrieved text. We introduce binary token representations (BTR), which use 1-bit vectors to precompute every token in passages, significantly reducing computation during inference. Despite the potential loss of accuracy, our new calibration techniques and training objectives restore performance. Combined with offline and runtime compression, this only requires 127GB of disk space for encoding 3 billion tokens in Wikipedia. Our experiments show that on five knowledge-intensive NLP tasks, BTR accelerates state-of-the-art inference by up to 4x and reduces storage by over 100x while maintaining over 95% task performance.

Figures (5)

• Figure 1: Retrieval-augmented models use a retrieve-and-read pipeline. The reader can be either an encoder or an encoder-decoder model. BTR creates cacheable binary representations for the passages via decomposition and calibrated binarization to speed up reader inference. BTR further reduces storage by offline compression and improves inference speed by runtime compression.
• Figure 2: BTR reader architecture, where light blue color indicates the model from prior work (based on T5 raffelExploringLimitsTransfer2020 and FiD izacardLeveragingPassageRetrieval2021, and green indicates our methods. We create cacheable calibrated binary token representations for retrieved passages in the reader encoder to speed up inference. Additionally, we compress the precomputed binary token presentations offline to reduce storage costs. We further reduce inference computation by runtime compression for the encoder and decoder.
• Figure 3: Task performance and inference efficiency (throughput and storage) visualization across the baselines and BTR fine-tuned models on three open-domain QA datasets. The area of the circle denotes storage size (we scale the area size of BTR-base and DensePhrase by 10 for clear display). A top-right smaller circle is better.
• Figure 4: Accuracy, storage, and throughput comparisons for different two-stage token compression ratios on the NaturalQuestions dataset. To achieve a good balance between accuracy, storage, and throughput in BTR, we choose a compression ratio of 0.2 for both online and offline token compression. Detailed numbers are in \ref{['tab:compression-ratios']} in \ref{['apdix:ablation']}.
• Figure : Offline Compression for Binary Token Representations