BitNet и Mamba: языковую модель уместили в 25 КБ и запустили на BBC Micro на процессоре 1975 года

Автор проекта (его имя в тексте не указано) решил проверить, какую самую сильную языковую модель можно уместить на BBC Micro его отца, домашнем компьютере начала 1980-х на процессоре MOS 6502, выпущенном в 1975 году и известном также по Apple II. Весь код инференса и веса модели нужно было уместить в 25 КБ пользовательской памяти; итоговая конфигурация заняла 22 КБ: 9 КБ код инференса и 13 КБ веса модели.

У 6502 8-битная арифметика и нет инструкции умножения: обычное умножение-с-накоплением 8×8 бит собирается из сдвигов и сложений и стоит 150 тактов. Поэтому веса квантовали методом BitNet, они принимают только три значения: -1, 0 или 1. Такое тернарное умножение сводится к последовательности сложений и вычитаний и занимает всего 30 тактов. Каждый параметр BitNet хранится в 1,58 бита (log2(3)) против 8 бит у int8 и 32 бит у float32; упаковка по 4 параметра в байт позволила уместить в 13 КБ 52 тысячи параметров. Размерность эмбеддинга, 56, а словарь токенов сознательно сведён к 26 буквам латиницы и пробелу: более крупный словарь съел бы слишком большую часть и без того тесного бюджета параметров. Финальный слой предсказания токена (LM head) оставили в int4 для более точного распределения вероятностей, а все остальные матрицы модели, тернарные.

Архитектуру для рекуррентной части выбирали методом проб. Attention отвергли сразу: кэш ключей-значений растёт с каждым сгенерированным токеном и быстро съел бы весь бюджет в 32 КБ оперативной памяти машины. Из архитектур с постоянным объёмом состояния сначала попробовали GRU, но обучение расходилось на каждом прогоне, в режиме BitNet спектральный радиус матриц намного превышает 1, а чтобы удержать его около 1, пришлось бы занулить 98% весов. Единственным способом стабилизировать GRU оказалось хранить главную весовую матрицу в более высокой точности (например, int4), но это резко сокращало возможную размерность модели, поскольку это самая крупная матрица в сети. В итоге выбрали Mamba (модель пространства состояний, SSM): коэффициент затухания в ней вычисляется на этапе инференса для каждого канала отдельно и по построению лежит в диапазоне [0, 128)/128, то есть никогда не превышает 1, взрыв градиентов конструктивно невозможен.

Активации хранятся в 8 битах, а накопление идёт в 16 битах: при модуле каждого слагаемого не больше 128 без переполнения можно сложить до 256 членов. Наивное обрезание 16-битного значения до 8 бит теряет почти весь динамический диапазон, на примере накопления 64 равномерно распределённых значений в диапазоне (-128, 127) стандартное отклонение результата получается около 591, из-за чего 83% значений упёрлись бы в границы обрезания. Вместо этого ввели обучаемый параметр сдвига: перед обрезанием 16-битное значение делят через побитовый сдвиг вправо, и от этого параметра сильно зависит итоговое качество, его изменение на единицу удваивает или уполовинивает амплитуду активаций после слоя. Лучше всего сработало обучение с изменяемым сдвигом в первой половине тренировки и заморозкой во второй.

Для выбора следующего токена жадная генерация (взять максимум логитов), самый простой вариант, но даёт скучный текст. У 6502 нет функции экспоненты, поэтому top-k-сэмплирование через softmax реализовали через таблицу заранее просчитанных экспонент для нужной температуры. А поскольку у 6502 без участия пользователя нет источника случайности, семя генератора псевдослучайных чисел жёстко зафиксировано, модель каждый раз выдаёт один и тот же результат.

Код инференса написан на C и скомпилирован под набор команд 6502 через компилятор CC65; готовый бинарник записывается на BBC Micro через PlayUEF и самодельный кабель-переходник 3,5-мм-джек, кассетный накопитель: ноутбук проигрывает звук из наушникового разъёма, а BBC воспринимает его как кассету. Для проверки корректности реализации на C её сверяли с эталонной моделью на Python в эмуляторе sim65. Попробовать результат может любой желающий: ссылка в статье загружает эмулятор BBC Micro прямо в браузере, подтягивает образ ленты с GitHub и сама вводит нужные команды, установка эмулятора не требуется, а сама генерация занимает несколько минут.

По итогу модель работает, она не особенно умна, но действительно выполняет инференс архитектуры, похожей на современные языковые модели, на железе начала 1980-х. Автор связывает проект с тезисом исследовательницы Сары Хукер (Sarah Hooker) о «лотерее железа» (hardware lottery): передовые архитектуры и алгоритмы развиваются в связке с доступным оборудованием, и при другом наборе аппаратных ограничений оптимальным мог бы оказаться совсем другой выбор архитектуры.

Ключевые факты

Веса квантованы методом BitNet до тернарных значений {-1, 0, 1}, 1,58 бита на параметр против 8 у int8 и 32 у float32; тернарное накопление занимает 30 тактов процессора против 150 у обычного умножения 8×8
Бюджет памяти, 25 КБ, итоговая конфигурация заняла 22 КБ (9 КБ код + 13 КБ веса, 52 тысячи параметров, размерность эмбеддинга 56) при словаре всего в 26 букв и пробел
Attention отвергли из-за роста кэша ключей-значений при 32 КБ ОЗУ; GRU расходился при обучении в каждом прогоне, поэтому выбрали Mamba, её коэффициент затухания по построению не превышает 1
Наивное обрезание 16-битных активаций до 8 бит теряет динамический диапазон (83% значений упирались бы в границы в тестовом примере), вместо него ввели обучаемый параметр битового сдвига
Демо запускается прямо в браузере через эмулятор BBC Micro без установки, генерация занимает несколько минут; из-за отсутствия источника случайности на 6502 вывод модели детерминирован

Почему это важно

Проект наглядно показывает, как далеко можно зайти в сжатии языковой модели, если жёстко привязать архитектурные решения к возможностям конкретного железа: квантование до тернарных весов, отказ от attention в пользу SSM-архитектуры (Mamba), методы, которые обсуждаются и для куда более крупных современных моделей на грани устройства (edge inference). Работа иллюстрирует тезис о «лотерее железа»: то, какая архитектура считается «лучшей», зависит от доступного оборудования, а не является абсолютной истиной.

Кому это важно

Инженерам, которые занимаются квантованием и инференсом моделей на маломощных и встраиваемых устройствах; сообществу ретрокомпьютинга, которому интересно применение современного ML к железу 1970-1980-х; преподавателям, которым нужен компактный наглядный пример устройства языковой модели без абстракций фреймворков.

Как это применить

Демо можно запустить прямо в браузере по ссылке из статьи, эмулятор BBC Micro сам загрузит образ ленты и введёт команды, установка ничего не требует. Инженерам, работающим над инференсом на ограниченной памяти, могут пригодиться конкретные приёмы из статьи: упаковка тернарных весов по 4 на байт, таблица экспонент вместо softmax, обучаемый параметр сдвига вместо наивного обрезания активаций.

Можно ли доверять

Это личный инженерный проект: имя автора в самом тексте не указано, но статья подробно и последовательно объясняет каждое архитектурное решение, приводит код и ссылки на полный набор строительных блоков инференса на GitHub, а результат можно немедленно проверить самому через эмулятор в браузере, это делает написанное легко проверяемым, хотя и без формального рецензирования.

Риски и подводные камни

Это proof-of-concept, а не измеренная по метрикам общая языковая модель: словарь ограничен 27 символами, оценок качества генерации в тексте нет. Вывод модели полностью детерминирован, из-за отсутствия источника случайности на 6502 семя генератора зафиксировано, и результат каждый раз одинаковый, что не отражает поведение обычных языковых моделей со случайной генерацией. Проект остаётся хобби-инженерией, а не рецензируемым исследованием или продуктом.