ZeroFree-1Bit: 1-Bit Quantization Without Zero (MNIST)

ZeroFree-1Bit — это исследовательский проект, демонстрирующий, что 1-битное квантование весов {-1, +1} без использования нуля может работать не хуже, а иногда и лучше полной FP32-модели.

Главная особенность — квадратичная активация (QuadAct), которая не имеет мёртвых зон (0.19% против 23% у SwiGLU) и позволяет модели обучаться даже при жёстких бинарных весах.

🚀 Ключевые результаты (MNIST)

Модель	Точность	Память на веса	Комментарий
FP32 + ReLU (8 эпох)	98.16%	32 бита	Базовая линия
FP32 + QuadAct (7 эпох)	97.43%	32 бита	Наша активация
ZeroFree-1Bit (sign, 1 бит)	94.69%	1 бит	Честное квантование
ZeroFree-1Bit + Fine-tune (1 эпоха)	97.78%	1 бит	Превзошли FP32 (tanh)

Экономия памяти — 32x (1 бит против 32 бит) при точности, близкой к FP32.

📐 Математическое определение

1. Квадратичная активация (QuadAct)

Для входа z ∈ R и параметра β > 0:

QuadAct(z) = z², если z >= 0 QuadAct(z) = -β · z², если z < 0

Производная (субградиент в нуле):

d/dz QuadAct(z) = 2z, если z > 0 d/dz QuadAct(z) = -2βz, если z < 0

Ключевое свойство: для z < 0 градиент не равен нулю (в отличие от ReLU), что предотвращает «мёртвые зоны». Параметр β = 1/√N, где N — размерность входа, обеспечивает нормировку масштаба.

2. Квантование весов

Обучение (мягкий знак): w_train = tanh(10 · w)

Инференс (жёсткий знак): w_quant = sign(w) = +1, если w >= 0 w_quant = sign(w) = -1, если w < 0

Важно: нуль не используется. Это отличает нас от BitNet b1.58 (где w ∈ {-1, 0, +1}) и экономит 37% памяти.

3. Статистика «мёртвых зон»

Для случайного входа x ~ U(-16, 16) (100000 точек) мы считаем долю точек, где |grad(x)| < 0.01:

DeadZone(act) = (1/N) · Σ I(|d/dx act(x_i)| < 0.01)

Результаты:

QuadAct: 0.19%
SwiGLU: 23.26%

⚡ Экономия вычислений: почему 1 бит быстрее

Для слоя с матрицей весов W ∈ {-1, +1}^{d×d} и входом x ∈ R^d стандартное матричное умножение требует d² операций умножения (FP32):

y_i = Σ_j W_ij · x_j

При W_ij ∈ {-1, +1} умножение заменяется на сложение/вычитание:

y_i = Σ_{j: W_ij=+1} x_j - Σ_{j: W_ij=-1} x_j

Это даёт:

Нет умножений — только сложения и вычитания (на порядок дешевле на CPU/GPU).
Нет нулей — все веса активны, нет «пропусков» (в отличие от BitNet b1.58, где нули требуют ветвлений).

🔍 Мелкий шрифт: пример на ARMv8-A

Приведённая ниже оценка — теоретическая, основанная на тактах инструкций. Реальный бенчмарк зависит от модели, компилятора и микроархитектуры.

Рассмотрим один нейрон y = Σ_{j=1}^{1024} W_j · x_j, где W_j ∈ {-1, +1}, на типичном ARMv8-A (например, Cortex-A76 в Raspberry Pi 5).

Базовый NEON (128-бит): Загрузка 4 входов (float32): ldr q0, [x] (1 такт). Вместо умножения — сложение/вычитание по маске: fadd v2.4s, v2.4s, v0.4s (2 такта, для W=+1) fsub v2.4s, v2.4s, v0.4s (2 такта, для W=-1)

На 1024 элемента: 1024 / 4 = 256 итераций × ~2 такта ≈ 512 тактов.

С расширением DotProd (SDOT) (требует int8-квантования активаций): Загрузка 8 int8 входов и 8 int8 весов: ldr d0, [x] (1 такт). sdot v2.4s, v0.16b, v1.16b (1 такт, 4 умножения-сложения int8).

На 1024 элемента: 1024 / 8 = 128 итераций × ~1 такт ≈ 128 тактов.

Реальный бенчмарк (llama.cpp, ARMv9): на GCP Axion (Neoverse-V2, i8mm/SVE2) включение -march=armv9-a+i8mm+dotprod+sve дало ускорение 2.57x на prompt processing (14.96 → 38.51 tok/s) и 1.32x на генерации (12.03 → 15.96 tok/s) для MoE-модели.

📱 ARMv9+ (оговорка про SME2)

На новых ARMv9-чипах (Apple M4, iPhone 16, серверные Neoverse V2+) появляется SME2 (Scalable Matrix Extension 2). Это отдельный матричный сопроцессор с ZA-хранилищем, который специально заточен под нейросети.

Для нашего подхода это идеальная платформа:

SME2 имеет dedicated lookup table instructions для бинарных нейросетей (compressed neural network support). Это буквально создано для весов {-1, +1}.
Arm заявляет до 5x ускорения AI-нагрузок на SME2 по сравнению с предыдущим поколением CPU.
Интеграция через KleidiAI уже работает в llama.cpp, PyTorch, ONNX Runtime без изменения кода.

Оговорка: SME2 доступен только на ARMv9.3-A и новее. На старых ARMv8-A (Raspberry Pi 4/5, большинство мобильных чипов до 2024) его нет. Поэтому наш пример выше — это ARMv8-A с NEON/DotProd, а SME2 — это потенциал на будущее, который мы планируем задействовать, когда платформа станет доступнее.

🔬 Почему это работает?

Квадратичная активация даёт ненулевые градиенты даже для отрицательных входов. У SwiGLU 23% мёртвых зон, у нас — 0.19%.
Веса строго {-1, +1} без нуля. Это на 37% меньше памяти, чем BitNet b1.58 (которому нужен ноль).
BatchNorm + tanh(10x) при обучении стабилизируют процесс, а калибровка + 1 эпоха fine-tune возвращают точность после честного квантования.

📦 Установка и запуск

pip install torch torchvision matplotlib python zero_free_1bit_mnist.py

Скрипт обучает три модели (FP32+ReLU, FP32+QuadAct, ZeroFree-1Bit) и строит график results.png.

📊 Графики

results.png — сравнение точности и памяти.
loss_curves.png — кривая loss по эпохам (если сгенерирован).

💡 Для систем с ограниченными ресурсами: логиты

Мы понимаем, что не все системы могут позволить себе вычисление полного softmax (особенно на edge-устройствах или при работе с большими словарями). Поэтому в модель добавлен режим возврата логитов — сырых выходов последнего слоя до применения softmax.

Это даёт несколько преимуществ:

Гибкость: вы можете применить свою собственную функцию нормализации (например, ZeroFree-Softmax из нашего исследования) или вовсе обойтись без неё, если задача позволяет.
Экономия вычислений: softmax требует двух проходов (поиск максимума и суммы экспонент). Логиты позволяют пропустить этот шаг.
Совместимость: многие фреймворки (ONNX, TensorRT, TFLite) ожидают на входе именно логиты, а не вероятности.

Чтобы получить логиты, просто вызовите модель без финального softmax:

logits = model(x) # вместо probs = torch.softmax(model(x), dim=-1)

📜 Лицензия

MIT. Используйте, модифицируйте, распространяйте.

🙏 Благодарности

Сообществу Hugging Face за transformers, trl, datasets.

Downloads last month: -; Downloads are not tracked for this model. How to track