burn-bitnet - BitNet Quantization for Burn

BitNet quantization family for the Burn deep learning framework. 1.58-bit ternary weight quantization + BitNet v2 Hadamard-based activation quantization. Straight-through estimator (STE) gradients throughout.

Papers: BitNet b1.58 (Ma et al., 2024), BitNet v2 (Wang et al., 2025).

Install

cargo add burn-bitnet

Quick start

use burn_bitnet::{weight_quant_ternary, quantize_4bit, quantize_8bit};

// Quantize weights to ternary {-scale, 0, +scale}
let w_q = weight_quant_ternary(weight);

// Quantize activations to 4-bit with Hadamard
let x_q = quantize_4bit(activation);

// 8-bit with Hadamard
let x_q = quantize_8bit(activation);

API

Function	Reference	What	Input
weight_quant_ternary	b1.58 (2024)	W → sign(W−μ)·scale	[M, N]
activation_quant_8bit	b1.58 (2024)	absmax per-token → int8	[B, T, D]
bitnet_v2_quantize	v2 (2025)	Hadamard + absmax/absmean	[B, T, D]
fast_walsh_hadamard	v2 (2025)	FWHT O(n log n)	[M, N]
quantize_4bit	v2 (2025)	bitnet_v2_quantize(x, 4)	[B, T, D]
quantize_8bit	v2 (2025)	bitnet_v2_quantize(x, 8)	[B, T, D]

How it works

Ternary weight quantization (b1.58):
    scale = mean(|W|)
    W_q = sign(W − mean(W)) · scale

Tensor shapes: [in_features, out_features]

BitNet v2 quantization:
    x → FWHT → quantize → FWHT → output

8-bit: absmax per-token → int8 range [-128, 127]
4-bit: absmean per-token → int4 range [-8, 7]

All functions use the straight-through estimator (STE): forward pass uses quantized values, backward gradient flows through the unquantized input.

License

AGPL-3.0. See LICENSE.