BitNet Inference Framework

BitNet-Embeddings is a family of multilingual text embedding models designed for high efficiency in storage and inference. They use a decoder-only architecture with last-token pooling and L2 normalization to generate dense embeddings. These models are suitable for tasks such as text retrieval, clustering, semantic similarity, classification, bitext mining, and reranking.

Key characteristics:

• Architecture: Transformer-based with BitLinear layers (BitNet framework).
• Quantization: Native 1.58-bit weights and 8-bit activations (W1.58A8).
• Context Length: Up to 32,768 tokens.
• Pooling: Last-token (EOS) pooling followed by L2 normalization.

bitnet.cpp is the official inference framework designed for 1-bit Large Language Models (LLMs), such as BitNet b1.58. It provides optimized kernels for fast and lossless inference on both CPU and GPU platforms.

Key benefits include:

• High Performance: Achieves significant speedups on both ARM (1.37x to 5.07x) and x86 (2.37x to 6.17x) CPUs.
• Energy Efficiency: Reduces energy consumption by up to 82.2% on x86 and 70.0% on ARM.
• Scalability: Capable of running large models (e.g., 100B parameters) on a single CPU at human-readable speeds (5-7 tokens per second).

TL2 tiling strategy is more complex because it requires BK % 6 == 0. To accommodate this, the dimension K is split into threeK and twoK. Computation for (M, threeK) is performed using TL2, while computation for (M, two_K) is performed using TL1.

To use codegen_tl2.py, ensure the following constraints are met:

• M % BM == 0
• K % BK % 32 == 0
• BM % bm == 0
• bm must be exactly 32

The BitNet CPU inference engine implements two parallelization strategies in its kernels to improve throughput:

1. Weight Parallel: Processes multiple weight rows or columns in a single kernel call. This is effective for reducing kernel launch overhead.
2. Activation Parallel: Built on top of weight parallel, this strategy amortizes the cost of I2_S weight unpacking across multiple activation elements.

Recommendation: For the I2_S quantization format, Activation Parallel is recommended because it provides significant benefits during the unpack operation. The current kernel defaults to activation parallel.

TL1 tiling strategy cuts weights into M / BM weights of shape (BM, K), then further cuts those into K / BK weights of shape (BM, BK). Finally, (BM, BK) weights are cut into (bm, compute_num / bm) compute blocks.

To use codegen_tl1.py, ensure the following constraints are met:

• M % BM == 0
• K % BK == 0
• BM % bm == 0
• bm must be chosen from the set {32, 64}

The I2_S format is used for efficient CPU inference of BitNet models by packing ternary weights ${-1, 0, +1}$ into a 2-bit representation.

Quantization and Encoding

1. Quantization: scale = 1/mean(|w|), then q = round(w * scale).clamp(-1, 1).
2. Encoding:
   • -1 $\rightarrow$ 0
   • 0 $\rightarrow$ 1
   • +1 $\rightarrow$ 2

Data Layout

• Weights are processed in blocks of 128 values.
• Each block is packed into 32 bytes.
• Each byte stores 4 values using the bitwise pattern: byte = (c0 << 6) | (c1 << 4) | (c2 << 2) | c3.
• A float32 scale is appended to the end of the packed data buffer.

Tensor Type Assignment in I2_S mode

• 2D linear weights: I2_S ternary packed.
• Embedding weights: float16.
• Norm weights (1D): float16.

BitNet-Embeddings models use a specific quantization and architectural setup to optimize for inference:

Quantization (W1.58A8)

• Weights: Quantized to ternary values {-1, 0, +1} using absmean quantization.
• Activations: Quantized to 8-bit integers using absmax quantization (per-token).
• Note: These models are trained from scratch with this scheme, not post-training quantized.

Architectural Details

• Position Embeddings: Uses Rotary Position Embeddings (RoPE).
• Normalization: Employs SubLN (sub-layer normalization) for training stability.
• Linear Layers: No bias terms are used in linear or normalization layers.
• Output: The output embeddings can be further quantized to 8, 4, 2, or 1 bit to trade off storage cost against retrieval performance.

Microsoft has released several specialized 1-bit models optimized for different tasks:

1. BitNet-b1.58-2B-4T (LLM)

• Type: Ternary (1.58-bit) Large Language Model.
• Size: 2.4B parameters.
• Training: 4 trillion tokens.
• Features: Chat-ready, supports GPU inference, and offers massive CPU speedups.

2. BitNet-embedding-0.6B (Embedding Model)

• Type: 1-bit Embedding Model.
• Size: 0.6B parameters.
• Features: Competitive embedding quality with 2 bits per weight; supports optimized I2_S conversion on x86 CPUs.

3. BitNet-embedding-270M (Lightweight Embedding Model)

• Type: 1-bit Embedding Model.
• Size: 270M parameters.
• Features: Designed for resource-constrained edge deployment with a minimal memory footprint.

The BitNet-Embeddings family includes two primary 1.58-bit models. Choose based on your required embedding dimension and performance needs:

Note: There are also bf16 versions (e.g., harrier-oss-v1-270m) available for comparison.

Use Cases

• Direct Use: Efficient information retrieval for RAG, web search, enterprise search, and question answering. Text clustering, classification, and bitext mining.

Limitations

• No Text Generation: BitNet-Embeddings does not generate human-readable text; it only produces dense embedding vectors.
• Language/Domain Sensitivity: Performance may be limited in low-resource languages or niche domains (legal, medical, scientific) not well-represented in training data.
• High-Risk Applications: Not recommended for commercial or real-world applications without further testing and development.