В современных нейросетях критически важно, сколько физической памяти занимает каждый параметр. В этой работе я попытался уйти от классического float32 в нейросетевом слое к uint8 без квантования. Для этого все вычисления проводились сразу по правилам арифметики остатков в конечном поле Галуа GF(137).
Стоит сказать, что это не замена обычному инференсу и не попытка доказать, что все должны срочно переписать модели на вычеты по модулю 137. Я взял небольшой слой, байтовые веса, нативное ядро, ARM NEON и несколько базовых реализаций для сравнения.
Спойлер: на маленьком контрольном запуске байтовая модель заняла около 1.63 KB вместо 6.50 KB, при этом 1000 прямых проходов в ARM NEON режиме заняли всего 12.1 ms. В лучшем случае на выборке это дает около 4x по памяти и до 4.86x по времени относительно базовой NumPy float32-реализации.
Коротко про идею
Обычный плотный слой можно записать так:
y = activation(x*W + b)
В моём прототипе это было заменено на конечнополевой вариант:
r = (x*W + b) mod 137y = 1, если r >= 42y = 0, если r < 42
То есть матричное умножение остается на месте, но накопленная сумма приводится по модулю 137, а затем проходит через простой порог.
Отдельная операция в ядре выглядит так:
mac137(a, b, c) = (c + a*b) mod 137
Как видите, это не квантизация готовой float32-модели. В моем решении значения изначально живут в байтовом конечнополевом представлении: веса и активации хранятся как uint8, а вычисления идут в арифметике по модулю 137.
Уже предвосхищаю вопрос почему именно 137? Потому что это простое число меньше 256, а значит все значения от 0 до 136 помещаются в uint8. Сама арифметика по модулю 137 образует конечное поле GF(137).
В поле Галуа у любого ненулевого элемента есть обратный ему. Это значит, что для любого a != 0 можно найти такое число a^{-1}, что
a * a^{-1} ≡ 1 (mod 137)
Например, обратный элемент к 2 по модулю 137 — это 69:
2 * 69 = 138 ≡ 1 (mod 137)
То есть деление тоже можно заменить умножением на обратный элемент. В обычной арифметике это звучит естественно, но по модулю так работает не всегда. Например, если взять модуль 256, то у 2 нет обратного элемента: 2*x всегда будет чётным и никогда не даст остаток 1.
Поэтому простой модуль удобнее. И хоть он и даёт маленький дискретный алфавит, но работать с ним удобнее из-за “нормальной” алгебры. В GF(137) можно складывать, умножать и, где нужно, делить на ненулевые элементы.
Нативное ядро
Самое дорогое место в такой схеме не сама идея mod 137, а то, как она исполняется. Если делать всё в Python, быстро становится понятно, что выигрыш по памяти не спасает. Накладные расходы интерпретатора и переходы между массивами съедают почти весь выигрыш.
Поэтому я вынес горячую часть в нативное ядро. Оно делает три вещи в одном проходе:
-
накапливает сумму произведений;
-
приводит результат по модулю 137;
-
применяет пороговую активацию.
Ниже не весь файл, а только маленький фрагмент: редукция по модулю 137 и пороговая активация. В полном ядре эта операция вызывается внутри матричного прохода.
constexpr uint32_t kModulus = 137U;constexpr uint8_t kDefaultThreshold = 42U;// Константа для редукции Барретта:// floor(2^32 / 137). Она позволяет заменить деление на 137// умножением и сдвигом.constexpr uint64_t kBarrettMu = (uint64_t{1} << 32) / kModulus;inline uint8_t mod137_barrett(uint32_t value) { // Приближаем частное value / 137. const uint32_t quotient = static_cast<uint32_t>( (static_cast<uint64_t>(value) * kBarrettMu) >> 32 ); // Вычитаем quotient * 137 и получаем остаток. uint32_t residue = value - quotient * kModulus; // После редукции Барретта остаток может быть немного выше модуля, // поэтому добиваем его одним или несколькими вычитаниями. while (residue >= kModulus) { residue -= kModulus; } return static_cast<uint8_t>(residue);}inline uint8_t threshold_bit(uint8_t residue, uint8_t threshold) { // Простая бинарная активация по остатку конечного поля. return static_cast<uint8_t>(residue >= threshold ? 1U : 0U);}
Зачем здесь редукция Барретта? В горячем цикле не хочется делать обычное деление, если его можно заменить умножением, сдвигом и несколькими вычитаниями. В этом месте компактность нашего модуля 137 тоже оказывается полезной, ведь для маленького фиксированного модуля это удобный приём, потому что и модуль, и константа Барретта нам известны заранее.
Не пугайтесь арифметики остатков, смысл ядра довольно прост. Нам не приходится таскать данные туда-сюда между несколькими слоями абстракции. Умножение, редукция по модулю и пороговая активация выполняются внутри одного нативного цикла, без промежуточных массивов между этапами.
Отмечу, что в полном варианте дополнительно используются:
-
хранение весов и активаций как
uint8_t; -
целочисленное накопление;
-
постоянный пул рабочих потоков в нативном коде;
-
путь ARM NEON для формы слоя с шириной скрытого слоя
H = 32; -
режим повторных вызовов, где измеряется внутренний цикл, а не граница Python/ctypes.
На чём я это проверял
Для первого теста я не стал брать большую нейросеть. Хотелось сначала понять, есть ли вообще смысл в таком представлении? Например сколько оно занимает в памяти, где упирается в Python, а где начинает работать нативное ядро.
Поэтому начал с малого и попытался отличить простые числа от составных. Я взял 500 простых и 500 составных чисел в диапазоне от 1000 до 5000. Затем каждое число превратил в вектор признаков, а именно: брал первые 50 дробных разрядов его записи по основанию числа e.
Подобная идея не дает хорошей классификации сама по себе, зато позволяет быстро собрать фиксированный набор данных и гонять по нему разные реализации одного и того же слоя, в общем всё свелось к удобной воспроизводимости.
Поэтому в этой статье меня интересовала не точность классификации. Главные вопросы были проще:
-
сколько памяти занимает представление;
-
сколько времени занимает 1000 прямых проходов;
-
насколько сильно Python мешает компактной арифметике;
-
что меняется, когда горячий цикл уходит в нативный код.
Что получилось
Самая важная таблица получилась такой:
|
Реализация |
Память |
Время на 1000 прямых проходов |
Комментарий |
|---|---|---|---|
|
NumPy float32 MLP |
6.50391 KB |
58.7784 ms |
базовая строка |
|
GF(137) компактный uint8 |
1.62793 KB |
582.904 ms |
Python съел весь выигрыш |
|
GF(137) нативный объединенный uint8 |
1.62793 KB |
40.1530 ms |
вызов из Python всё ещё заметен |
|
GF(137) NEON нативный цикл |
1.62793 KB |
12.1055 ms |
лучший замер в этой конфигурации |
|
JAX JIT CPU |
6.50391 KB |
42.4758 ms |
сильная float32-база |
Первый вывод был ожидаемый: байтовое представление сразу дает примерно 4x по памяти. Тут никакой магии нет: float32 занимает 4 байта, uint8 1 байт.
Второй вывод оказался полезнее. Просто хранить данные компактно мало. Python-вариант с uint8 получился самым медленным: 582.904 ms на 1000 проходов. То есть компактное представление без нормального исполнения само по себе не спасает.
После выноса горячего цикла в нативный код ситуация изменилась. Путь с вызовом из Python прошел 1000 проходов за 40.1530 ms, а нативный ARM NEON-цикл занял 12.1055 ms.
Если сравнивать именно с моей NumPy float32-реализацией, это около 4x по памяти и до 4.86x по времени. Но стоит учитывать, что это не универсальная оценка скорости GF(137). Это результат конкретной формы слоя, конкретного железа и конкретного режима измерения.
Где результат не впечатлил
Самое неудобное место в таблице это компактный Python-вариант:
GF(137) компактный uint8: 1.62793 KB, 582.904 ms
По памяти всё хорошо, а по времени плохо. Тут важный момент. Байтовое хранение само по себе ничего не гарантирует. Если вычисление остаётся на уровне Python, если используются промежуточные массивы, то накладные расходы быстро съедают весь выигрыш.
Второй неприятный момент появился в контрольных запусках с обычным uint8. Если убрать арифметику GF(137) и оставить более простой пороговый uint8-путь, он может быть быстрее. Но это уже другая функция, не эквивалентная конечнополевому слою.
Поэтому финальный вывод пришлось сузить. Сначала я хотел получить одновременно и скорость, и компактность, но по итогу я не могу честно написать:
GF(137) быстрее стандартного int8-инференса.
Правильнее так:
В этом маленьком воспроизводимом тесте нативное конечнополевое ядро дает компактное хранение и хорошую скорость для конкретной операции. Но стандартная uint8/int8-реализация остаётся обязательной базовой проверкой, а не формальностью.
Мне кажется, это нормальный результат. Не такой эффектный, зато проверяемый.
Как воспроизвести
Активная проверка запускается так:
uv run --extra test pytest
Ожидаемый результат:
5 passed
Ограничения
Главное ограничение: это маленький тест, а не готовая inference-библиотека.
Я проверял конкретную форму слоя, конкретные размеры и конкретное железо. На других размерах матриц результат может измениться. Особенно если сравнивать не с простой NumPy-реализацией, а с хорошо настроенными int8-библиотеками.
Ещё важный момент: это не квантизация обученной модели. Я не брал готовую float32-сеть и не переводил её в int8 с сохранением качества. Слой изначально задан в конечнополевом виде, поэтому сравнение по accuracy здесь не является главным.
Что я считаю честным выводом:
-
байтовое представление даёт ожидаемые 4x по памяти;
-
Python-реализация убивает выигрыш по времени;
-
нативное ядро уже даёт интересные числа;
-
ARM NEON в этом тесте оказался самым быстрым вариантом;
-
обычный uint8/int8 baseline нужно держать рядом всегда.
То есть результат не в том, что GF(137) “лучше int8”. Результат в том, что конечнополевой слой можно сделать компактным, воспроизводимым и достаточно быстрым для маленькой операции. Дальше это уже надо проверять на более нормальных задачах.
Как посмотреть код
Основной репозиторий с кодом и PDF
Отдельный репозиторий статьи про периферийный инференс
ссылка на оригинал статьи https://habr.com/ru/articles/1044172/