Technical Summary: BitNet Text Embeddings (BITEMBED)

Problem Statement

While Large Language Model (LLM)-based text embedders have significantly enhanced retrieval and semantic representation quality, their practical deployment faces two critical bottlenecks:

1. High Inference Latency: LLM backbones are computationally expensive compared to traditional encoder-only models, slowing down the embedding of queries and documents.
2. Substantial Storage and Bandwidth Overhead: High-dimensional, full-precision embeddings impose heavy memory and data movement costs, particularly in large-scale retrieval systems maintaining vectors for millions or billions of documents.

Existing extreme low-bit LLM approaches (e.g., BitNet) focus on general task models but have not been systematically adapted for the specific requirements of text embedding, which demands high-fidelity semantic preservation despite aggressive quantization.

Methodology

The authors propose BITEMBED, a framework that converts pretrained LLM backbones into extreme low-bit embedding encoders. The approach follows a three-stage training pipeline:

1. BitNet-Style Backbone Conversion

The pretrained LLM is converted into a BitNet-style architecture featuring:

• Ternary Weights: Linear weights in transformer blocks are quantized to 1.58-bit values $\{-\Delta_w, 0, \Delta_w\}$.
• Quantized Activations: Token-wise 8-bit quantization is applied to activations entering BitLinear layers.
• Sub-layer Normalization (SubLN): To mitigate sensitivity to activation outliers and scale drift caused by quantization, SubLN modules are inserted inside each transformer block to normalize intermediate states before projection.
• Straight-Through Estimator (STE): Used to approximate gradients for non-differentiable quantization operations during back-propagation.

2. Training Pipeline

To adapt the quantized backbone to representation learning, BITEMBED employs a multi-stage optimization strategy:

• Continual Contrastive Pre-training: The quantized model undergoes pre-training on large-scale text pairs using the InfoNCE loss. This step rebuilds a broad semantic space distorted by extreme quantization before supervised fine-tuning.
• Distillation-Based Supervised Fine-tuning: The low-bit student is fine-tuned on high-quality supervised data (with hard negatives) while being guided by a full-precision (FP16) teacher model. This involves two distillation objectives:
  • Similarity-Distribution Distillation: Minimizes the KL divergence between the student's and teacher's soft similarity distributions over candidate sets, capturing nuanced relative preferences.
  • Attention-Relation Distillation: Aligns the token-level interaction structures (attention maps) of the student with the teacher, specifically from a middle-to-later transformer layer, to preserve structural dependencies.

3. Multi-Precision Output Embedding Training

Beyond quantizing the backbone, BITEMBED trains the output embeddings to be robust across multiple storage precisions (1, 2, 4, 8, and 16 bits).

• Mechanism: During training, the model optimizes an average loss across all supported bitwidths. For lower bitwidths, embeddings are quantized (e.g., scaled binary for 1-bit) and reconstructed for similarity computation.
• Goal: This enables a single model checkpoint to support different memory budgets, allowing practitioners to trade storage cost for performance without retraining.

Key Contributions

1. First Systematic Study: The paper presents the first systematic investigation of extreme low-bit quantization specifically for LLM-based text embedding, targeting both encoder inference efficiency and vector-index storage.
2. Effective Training Framework: It proposes a novel framework combining continual contrastive pre-training, supervised fine-tuning, and dual distillation (similarity distribution and attention relations) to recover representation quality lost during quantization.
3. Multi-Precision Adaptability: It introduces a training strategy for multi-precision output embeddings, demonstrating that a single model can achieve near-full-precision performance while supporting flexible storage configurations.

Experimental Results

The authors evaluated BITEMBED on the MMTEB (eng, v2) benchmark using Qwen3-0.6B and Gemma3-270M backbones.

• Performance Parity: BITEMBED remains largely comparable to full-precision teacher models.
  • On Qwen3-0.6B, BITEMBED achieved an average score of 67.60, only 0.35 points lower than the FP16 teacher (67.95).
  • On Gemma3-270M, it achieved 66.10, closely matching the teacher's 66.71.
• Inference Efficiency: The framework significantly improves token throughput on CPU (8 threads):
  • Qwen3-0.6B: Increased from 364.36 to 830.50 tokens/s (~2.3x).
  • Gemma3-270M: Increased from 1181.28 to 2055.47 tokens/s (~1.7x).
• Multi-Precision Trade-offs:
  • 8-bit and 4-bit embeddings were nearly lossless compared to 16-bit.
  • 1-bit and 2-bit embeddings preserved usable performance (e.g., 64.43 on Qwen3-0.6B) while reducing storage by roughly 16x and 8x, respectively.
• Ablation Findings:
  • Removing SubLN, distillation, or continual pre-training resulted in significant performance drops.
  • Direct fine-tuning of the converted BitNet backbone without these components yielded poor results (58.92 average score), confirming the necessity of the full pipeline.
  • Attention distillation from middle-to-later layers (e.g., layer 18 for Qwen3) proved most effective.
• Comparison with MRL: BITEMBED outperformed Matryoshka Representation Learning (MRL) in most storage scenarios, particularly at extreme compression levels, by maintaining the complete embedding space robust to low-bit reconstruction.

Significance and Claims

The paper claims that BITEMBED demonstrates that extreme low-bit training is a viable and promising direction for efficient semantic representation. By jointly addressing encoding efficiency and vector storage, the framework allows LLM-based embedders to be deployed in resource-constrained scenarios (such as online retrieval services) without sacrificing the semantic quality required for high-performance tasks. The ability to flexibly trade storage for performance via multi-precision embeddings offers a practical solution for scaling text embedding systems.