Реализации алгоритма TurboQuant на Python и Rust. Источник: arxiv.org/html/2504.19874v1
- Что такое TurboQuant
- Математическая основа
- Алгоритмы
- Теоретические гарантии
- Установка
- Использование
- API
- Результаты демо
- Структура кода
- Реализация на Rust
- Fuzzer
TurboQuant — схема векторного квантования, разработанная для эффективного вывода больших языковых моделей (LLM). Она сжимает вещественные векторы до нескольких бит на координату и при этом:
- минимизирует среднеквадратическую ошибку реконструкции (MSE-режим);
- обеспечивает несмещённую оценку скалярных произведений (режим IP).
Практическое применение — квантование весов и активаций трансформеров, векторные базы данных, ANN-поиск (approximate nearest neighbours).
Пусть вектор x лежит на единичной сфере S^{d-1}. После применения случайной ортогональной матрицы поворота Π координаты вектора y = Π·x становятся почти независимыми и каждая из них имеет маргинальное распределение:
f_X(x) = Γ(d/2) / (√π · Γ((d-1)/2)) · (1 − x²)^((d−3)/2), x ∈ (−1, 1)
Это бета-подобное распределение. При d → ∞ оно сходится к нормальному:
f_X(x) → N(0, 1/d)
Именно эта независимость координат после поворота позволяет применять скалярный квантизатор к каждой координате отдельно — без потери качества, которая была бы при наивном покоординатном квантовании без поворота.
Для каждой координаты решается одномерная задача оптимизации:
min_{c₁ ≤ ... ≤ c_{2^b}} Σᵢ ∫ |x − cᵢ|² · f_X(x) dx
Алгоритм итерации Ллойда:
- Инициализация: центроиды — квантили распределения f_X.
- Границы ячеек: середины между соседними центроидами.
- Обновление: каждый центроид = условное математическое ожидание f_X на своей ячейке.
- Повторять до сходимости (||Δc|| < 10⁻¹⁰).
При d > 50 используется гауссовское приближение N(0, 1/d) — оно точнее работает при больших размерностях и значительно быстрее (аналитические формулы вместо численного интегрирования).
MSE-оптимальный квантизатор вносит мультипликативное смещение в оценку скалярного произведения: E[⟨y, x̃_mse⟩] ≈ (2/π)·⟨y, x⟩. Для задач поиска ближайших соседей и attention-механизмов это недопустимо.
Решение — Quantized Johnson-Lindenstrauss (QJL) преобразование:
Encode: z = sign(S · x) ∈ {−1, +1}^d
Decode: x̃ = (√(π/2) / d) · ‖x‖₂ · Sᵀ · z
где S ∈ ℝ^{d×d} — матрица с элементами N(0, 1).
Свойство несмещённости:
E[⟨y, x̃⟩] = ⟨y, x⟩ для любого фиксированного y
Дисперсия: Var(⟨y, x̃⟩) ≤ (π/2d) · ‖y‖₂²
Инициализация:
- Случайная матрица поворота Π ∈ ℝ^{d×d} через QR-разложение случайной матрицы с элементами N(0,1).
- Вычисление оптимального кодебука {c₁, ..., c_{2^b}} через алгоритм Ллойда–Макса.
Кодирование Quant_MSE(x):
y ← Π · x # поворот
idxⱼ ← argmin_k |yⱼ − cₖ|, j=1..d # ближайший центроид
output: idx (b бит на координату)
Декодирование DeQuant_MSE(idx):
ỹⱼ ← c_{idxⱼ}, j=1..d # восстановление центроидов
x̃ ← Πᵀ · ỹ # обратный поворот
output: x̃
Затраты памяти: b·d бит на вектор.
Комбинирует MSE-квантование с QJL-остатком.
Инициализация:
- Создать TurboQuantMSE с битовой шириной (b−1).
- Случайная проекционная матрица S ∈ ℝ^{d×d} с элементами N(0,1).
Кодирование Quant_Prod(x):
idx ← Quant_MSE(x) # шаг 1: MSE-квантование (b-1 бит)
r ← x − DeQuant_MSE(idx) # шаг 2: остаток
z ← sign(S · r) # шаг 3: QJL на остатке
γ ← ‖r‖₂ # сохраняем норму остатка
output: (idx, z, γ)
Декодирование DeQuant_Prod(idx, z, γ):
x̃_mse ← DeQuant_MSE(idx)
x̃_qjl ← (√(π/2) / d) · γ · Sᵀ · z
output: x̃_mse + x̃_qjl
Затраты памяти: b·d бит + 32 бита (float для γ) на вектор.
| Метод | Дисторсия | Нижняя граница |
|---|---|---|
| TurboQuantMSE | D_mse ≤ (√3π/2) / 4^b | D_mse ≥ 1 / 4^b |
| TurboQuantProd | D_prod ≤ (√3π²·‖y‖²/d) / 4^b | D_prod ≥ (1/d) / 4^b |
TurboQuantMSE достигает результата в ~2.7× от теоретического оптимума. TurboQuantProd — практически оптимален по скалярному произведению.
Несмещённость TurboQuantProd:
E[⟨y, x̃⟩] = ⟨y, x⟩
Зависимости: Python 3.9+, NumPy, SciPy.
pip install numpy scipyФайл turboquant.py не требует установки — достаточно скопировать его в проект.
import numpy as np
from turboquant import TurboQuantMSE
d, b = 256, 4 # размерность и количество бит на координату
# Создать квантизатор (вычисляет матрицу поворота и кодебук)
q = TurboQuantMSE(d=d, b=b, seed=42)
# Один вектор (единичная норма)
x = np.random.randn(d)
x /= np.linalg.norm(x)
idx = q.encode(x) # (d,) uint16 — индексы центроидов
x_hat = q.decode(idx) # (d,) float64 — реконструированный вектор
mse = np.mean((x - x_hat) ** 2)
print(f"MSE = {mse:.6f}")
# Батч векторов
X = np.random.randn(1000, d)
X /= np.linalg.norm(X, axis=1, keepdims=True)
IDX = q.encode(X) # (1000, d) uint16
X_hat = q.decode(IDX) # (1000, d) float64x_raw = np.random.randn(d) * 5.0 # произвольная норма
idx, norm = q.encode_with_norm(x_raw)
x_rec = q.decode_with_norm(idx, norm)from turboquant import TurboQuantProd
q = TurboQuantProd(d=256, b=4, seed=42)
# Кодирование базы (хранимые векторы)
X = np.random.randn(1000, 256)
X /= np.linalg.norm(X, axis=1, keepdims=True)
mse_idx, qjl_signs, res_norms = q.encode(X)
# Оценка скалярного произведения с запросом y
y = np.random.randn(256)
y /= np.linalg.norm(y)
ip_true = X @ y # истинные скалярные произведения
for i in range(len(X)):
ip_est = q.inner_product_estimate(y, mse_idx[i], qjl_signs[i], res_norms[i])
# E[ip_est] == ip_true[i]X_tilde = q.decode(mse_idx, qjl_signs, res_norms) # (1000, 256)| Параметр | Тип | Описание |
|---|---|---|
d |
int | Размерность вектора |
b |
int | Бит на координату (1–16) |
seed |
int или None | Случайное зерно |
| Метод | Описание |
|---|---|
encode(x) |
Вектор(ы) → индексы uint16 |
decode(idx) |
Индексы → реконструированные векторы float64 |
encode_with_norm(x) |
Для ненормированных векторов: возвращает (индексы, нормы) |
decode_with_norm(idx, norms) |
Обратное к encode_with_norm |
mse(x) |
Вычислить среднее MSE на батче |
| Параметр | Тип | Описание |
|---|---|---|
d |
int | Размерность вектора |
b |
int | Бит на координату (≥ 2) |
seed |
int или None | Случайное зерно |
| Метод | Описание |
|---|---|
encode(x) |
Вектор(ы) → (mse_idx, qjl_signs, res_norms) |
decode(mse_idx, qjl_signs, res_norms) |
Реконструировать вектор |
inner_product_estimate(y, mse_idx, qjl_signs, res_norms) |
Несмещённая оценка ⟨y, x⟩ |
bits_per_vector() |
Количество бит на один сжатый вектор |
Вспомогательный класс. Используется внутри TurboQuantProd.
| Метод | Описание |
|---|---|
encode(x) |
x → (знаки ±1, нормы) |
decode(signs, norms) |
Несмещённая реконструкция |
Запуск python3 turboquant.py на d=256, n=1000 случайных единичных векторов:
============================================================
TurboQuant demo | d=256 n=1000
============================================================
--- TurboQuantMSE ---
b= 1 MSE=0.00141 setup=0.01s enc+dec=5.9ms
b= 2 MSE=0.00046 setup=0.01s enc+dec=6.8ms
b= 4 MSE=0.00004 setup=0.21s enc+dec=8.1ms
b= 8 MSE=0.00000 setup=3.29s enc+dec=133.5ms
--- TurboQuantProd (inner product) ---
b= 2 IP bias=+0.00122 RMSE=0.04685 bits/vec=544
b= 4 IP bias=-0.00007 RMSE=0.01389 bits/vec=1056
b= 8 IP bias=+0.00002 RMSE=0.00135 bits/vec=2080
--- Unbiasedness verification (b=4, n=5000) ---
Mean bias over 5000 vectors: -0.000153 (should be ≈ 0)
MSE убывает в ~4 раза при каждом добавлении 1 бита (соответствует теоретическому 1/4^b). Смещение оценки скалярного произведения практически равно нулю.
turboquant.py # библиотека квантования
├── _lloyd_max_gaussian() # Lloyd-Max для N(0, σ²)
├── _lloyd_max_beta_sphere() # Lloyd-Max для маргинала сферы
├── TurboQuantMSE # Алгоритм 1: MSE-квантование
│ ├── __init__ # поворот + кодебук
│ ├── encode / decode # пакетное кодирование/декодирование
│ └── encode_with_norm / # поддержка ненормированных векторов
│ decode_with_norm
├── QJL # Quantized Johnson-Lindenstrauss
│ ├── encode # x → sign(S·x)
│ └── decode # z → (√(π/2)/d)·γ·Sᵀ·z
├── TurboQuantProd # Алгоритм 2: IP-квантование
│ ├── __init__ # TurboQuantMSE(b-1) + QJL
│ ├── encode / decode # (idx, знаки, норма остатка)
│ ├── inner_product_estimate # несмещённая оценка ⟨y, x⟩
│ └── bits_per_vector # размер сжатого вектора в битах
└── _demo() # демо и замеры
fuzzer.py # fuzzer корректности (5 инвариантов)
├── make_unit_batch() # генерация случайных единичных векторов
├── check_roundtrip() # инвариант 1: encode→decode
├── check_mse_monotone() # инвариант 2: MSE убывает с ростом b
├── check_ip_bias() # инвариант 3: несмещённость IP
├── check_ip_variance() # инвариант 4: дисперсия IP в границах теории
├── check_edge_cases() # инвариант 5: граничные случаи
├── run_iteration() # одна случайная итерация
└── main() # CLI: --iters, --seed
Rust-реализация находится в директории rust/ и представляет собой крейт-библиотеку turboquant с исполняемым бинарником для демо.
Зависимости: rand 0.8, rand_distr 0.4 — генерация случайных матриц поворота и QJL.
Требования: Rust 1.85+ (edition 2024), Cargo.
cd rust
cargo build --releaseЗапустить демо:
cargo run --releaseuse turboquant::TurboQuantMse;
let d = 256;
let b = 4;
let q = TurboQuantMse::new(d, b, Some(42));
// Кодирование одного вектора (плоский срез длиной d)
let x: Vec<f64> = /* единичный вектор длиной d */;
let idx: Vec<u16> = q.encode(&x); // индексы центроидов
let x_hat: Vec<f64> = q.decode(&idx); // реконструированный вектор
// Батч: плоский буфер n*d элементов
let x_batch: Vec<f64> = /* n*d элементов */;
let idx_batch = q.encode(&x_batch); // Vec<u16> длиной n*d
let x_hat_batch = q.decode(&idx_batch);use turboquant::TurboQuantProd;
let q = TurboQuantProd::new(256, 4, Some(42));
// Кодирование вектора базы
let qv = q.encode(&x); // QuantizedVec { mse_idx, qjl_signs, res_norm }
// Оценка скалярного произведения с запросом y
let ip_est: f64 = q.inner_product_estimate(&y, &qv);
// E[ip_est] == ⟨y, x⟩
// Декодирование (реконструкция вектора)
let x_tilde: Vec<f64> = q.decode(&qv);
println!("bits per vector: {}", q.bits_per_vector());| Метод | Описание |
|---|---|
TurboQuantMse::new(d, b, seed) |
Создаёт квантизатор: матрица поворота + кодебук Lloyd-Max |
encode(x: &[f64]) -> Vec<u16> |
Плоский буфер векторов → индексы центроидов |
decode(idx: &[u16]) -> Vec<f64> |
Индексы → реконструированные векторы |
Параметры конструктора:
| Параметр | Тип | Описание |
|---|---|---|
d |
usize |
Размерность вектора |
b |
usize |
Бит на координату (1–16) |
seed |
Option<u64> |
Случайное зерно (None — не детерминировано) |
| Метод | Описание |
|---|---|
TurboQuantProd::new(d, b, seed) |
Создаёт квантизатор: TurboQuantMse(b-1) + матрица QJL |
encode(x: &[f64]) -> QuantizedVec |
Вектор → { mse_idx, qjl_signs, res_norm } |
decode(qv: &QuantizedVec) -> Vec<f64> |
Реконструкция вектора |
inner_product_estimate(y, qv) -> f64 |
Несмещённая оценка ⟨y, x⟩ |
bits_per_vector() -> usize |
Размер сжатого вектора в битах |
Вспомогательный тип, используется внутри TurboQuantProd.
| Метод | Описание |
|---|---|
Qjl::new(d, seed) |
Инициализация проекционной матрицы S ∈ ℝ^{d×d} |
encode(x) -> (Vec<i8>, f64) |
x → (знаки ±1, норма ‖x‖₂) |
decode(signs, norm) -> Vec<f64> |
Несмещённая реконструкция |
rust/
├── Cargo.toml
└── src/
├── lib.rs # публичный API: реэкспорт TurboQuantMse, TurboQuantProd, Qjl
├── main.rs # демо и замеры производительности
├── bin/
│ └── fuzzer.rs # fuzzer корректности (5 инвариантов)
├── lloyd.rs # Lloyd-Max квантизатор (гауссовский и сферический варианты)
├── mse.rs # TurboQuantMse: поворот + кодебук, encode/decode
├── qjl.rs # Qjl: матрица S, sign-проекция, несмещённое декодирование
└── prod.rs # TurboQuantProd: MSE + QJL остаток, оценка скалярных произведений
Standalone-фаззер корректности для Python и Rust. Генерирует случайные векторы и параметры, прогоняет 5 инвариантов на каждой итерации.
| # | Инвариант | Критерий |
|---|---|---|
| 1 | Roundtrip decode(encode(x)) не падает, форма верная |
shape совпадает с входом |
| 2 | MSE монотонность больше бит → меньше ошибка | MSE(b+1) ≤ MSE(b) × 1.1 |
| 3 | IP несмещённость оценка скалярного произведения без систематического смещения | |mean(est − true)| < 0.05 |
| 4 | IP дисперсия в пределах теоретической границы QJL (для d ≥ 16) | var ≤ (π/2d)·‖y‖² × 2.0 |
| 5 | Граничные случаи b=1/8, нулевые и почти нулевые векторы, малые d | нет паник, форма верная |
# Базовый запуск (50 итераций, случайный seed)
python fuzzer.py
# С фиксированным seed и количеством итераций
python fuzzer.py --iters 100 --seed 42Пример вывода:
Master seed: 42 iters: 3
[PASS] iter=01 d=4 b=7/6 n=1158
roundtrip=OK
mse_mono=OK
ip_bias=ip_bias=-0.00312
ip_var=skipped(d<16)
edge=OK
...
--- Small-d edge cases ---
[PASS] d=2 OK
[PASS] d=4 OK
Results: 5/5 passed
Exit-код: 0 — все прошли, 1 — есть падения.
# Базовый запуск
cd rust
cargo run --bin fuzzer
# С параметрами
cargo run --bin fuzzer -- --iters 100 --seed 42
# Release-сборка (быстрее)
cargo run --release --bin fuzzer -- --seed 0Пример вывода:
Master seed: 42 iters: 3
[PASS] iter=01 d=64 b=3/4 n=872
roundtrip=OK
mse_mono=MSE(3)=0.00234 MSE(4)=0.00089
ip_bias=ip_bias=+0.00178
ip_var=ip_var=0.00412 threshold=0.04909
edge=OK
...
--- Small-d edge cases ---
[PASS] d=2 OK
[PASS] d=4 OK
Results: 5/5 passed