Что такое AttnRes и зачем перестраивать Gemma 4 31B
Проект AttnRes - это попытка переосмыслить базовый механизм передачи информации внутри трансформера. Автор проекта, работающий в одиночку, взял модель Gemma 4 31B и заменил в ней стандартный residual stream на attention-based routing. Цель - не получить идеальную модель, а доказать жизнеспособность альтернативной архитектуры на бета-версии. Код и методология открыты под лицензией Apache 2.0.
Мотивация проста: стандартный residual stream передаёт информацию между слоями равномерно, суммируя выход каждого слоя с его входом. Это работает, но не оставляет модели возможности выбирать, какую именно информацию из предыдущих слоёв стоит пропустить дальше, а какую - приглушить. AttnRes даёт такую возможность через механизм внимания, который вычисляет веса для агрегации признаков с разных слоёв.
Схожий подход уже применялся в Kimi K3 от Moonshot - одной из самых заметных моделей последнего времени. Разбор архитектуры Kimi K3 показывает, что attention-based routing там использовался для повышения эффективности обработки длинных контекстов. Это подтверждает: идея не маргинальная, а перспективная.
Residual stream vs. Attention-based routing: ключевое отличие
Стандартный residual stream работает по формуле: выход слоя складывается с его входом, и сумма передаётся следующему слою. Каждый слой добавляет свою добавку к общему представлению, но не решает, насколько эта добавка важна в контексте других слоёв. Это линейное накопление информации.
Attention-based routing заменяет простое суммирование на взвешенную агрегацию. Вместо того чтобы тупо складывать выход слоя с входом, модель вычисляет коэффициенты внимания между текущим представлением и выходами всех предыдущих слоёв. Затем она смешивает их с этими весами. Результат: слой может «вслушаться» в конкретный предыдущий слой, если его признаки релевантны текущей задаче, и проигнорировать шум от остальных.
На практике это означает более гибкую обработку признаков. Модель учится маршрутизировать информацию динамически, в зависимости от контекста. Для длинных последовательностей это критично: не все токены одинаково полезны для предсказания следующего, и AttnRes позволяет фокусироваться на значимых.
Дистилляция с weaning-графиком: как обучить модель без бюджета Google
Главная проблема проекта - ресурсы. Один человек, без доступа к кластерам Google, пытается обучить модель с 31 миллиардом параметров. Решение - дистилляция знаний с weaning-графиком. Исходная Gemma 4 31B выступает учителем, а модифицированная версия с AttnRes - учеником. Задача ученика: воспроизвести поведение учителя, попутно осваивая новую архитектуру маршрутизации.
Дистилляция как метод уже стала индустриальным стандартом. Тот же Google запустил Gemini Distillation Service - сервис, позволяющий разработчикам дистиллировать большие модели в компактные без собственной инфраструктуры. AttnRes идёт другим путём: не сжатие, а архитектурная перестройка при сохранении размера. Это сложнее вычислительно, но открывает дорогу к новым свойствам модели.
Weaning-график: плавное отключение residual stream
Резко вырезать residual stream и включить AttnRes нельзя - модель потеряет связность представлений и качество упадёт катастрофически. Weaning-график решает эту проблему постепенным переходом.
На старте обучения вклад стандартного residual составляет 100%, а AttnRes - 0%. С каждым шагом коэффициент residual уменьшается, а коэффициент AttnRes растёт. Типичный график - линейное затухание residual за N шагов до нуля. К концу обучения модель полностью переходит на attention-based routing.
Почему это работает: на ранних этапах ученик видит стабильные представления от residual и учится предсказывать поведение учителя. По мере роста вклада AttnRes модель адаптирует внутренние представления под новую схему маршрутизации, не теряя накопленных знаний. Резкое переключение вызвало бы шок - градиенты стали бы нестабильными, и обучение развалилось бы.
Диверсификация данных: 20 языков и синтетические примеры
Дистилляция требует данных. В идеале - тот же датасет, на котором обучался учитель. Но доступа к нему нет. Автор AttnRes решает проблему генерацией синтетических примеров с помощью самой Gemma 4 31B.
Модель берёт тексты и переводит их на 20 языков, создавая мультиязычное покрытие. Это сохраняет manifold знаний - сложную поверхность в пространстве параметров, которую исходная модель выучила за время обучения. Одноязычная дистилляция рискует «схлопнуть» этот manifold, потеряв редкие паттерны и знания о низкоресурсных языках. Мультиязычная аугментация распределяет давление дистилляции по более широкой области, снижая риск катастрофического забывания.
Результат: ученик видит разнообразные примеры, которые активируют разные участки manifold знаний учителя. Это дешевле, чем собирать реальный мультиязычный датасет, и быстрее, чем ждать доступа к оригинальным данным.
Top-K 12 логгитов: дистилляция распределения вероятностей вместо one-hot
Стандартная дистилляция языковых моделей часто использует one-hot целевой токен или label smoothing. Проблема в том, что one-hot отбрасывает всю информацию о распределении вероятностей, которую выучил учитель. Модель видит только «правильный» токен и не знает, что учитель считал ещё 5-10 токенов почти такими же вероятными. Label smoothing немного размазывает вероятность, но не решает проблему структурно - хвосты распределения всё равно теряются.
AttnRes использует top-K 12 логгитов. Для каждого токена учитель выдаёт полное распределение вероятностей по всему словарю. Ученик получает не один целевой токен, а 12 наиболее вероятных токенов с их логитами. Функция потерь заставляет ученика воспроизводить относительные соотношения между этими логитами.
Модель рассматривается как предсказатель распределения, а не просто следующего токена. Это сохраняет информацию о неопределённости: если учитель колебался между «кошка» и «кот», ученик тоже научится присваивать обоим вариантам высокую вероятность. Редкие, но важные токены из длинного хвоста распределения не исчезают полностью - они остаются в топ-12 и продолжают влиять на обучение.
Практический эффект: модель после дистилляции лучше справляется с задачами, где важна нюансировка - генерация креативных текстов, перевод с учётом стиля, решение задач с неоднозначным ответом.
Ограничения и вызовы: B300, квантование и длинные хвосты
Проект откровенно говорит о своих ограничениях. Это бета-версия, и многие проблемы ещё предстоит решить. Три главных вызова - вычислительные требования, квантование и длинные хвосты распределений.
Почему нужно ждать B300: вычислительные требования AttnRes
Attention-based routing добавляет вычислительную нагрузку. Каждый слой теперь вычисляет внимание не только к предыдущему слою, а ко всем предыдущим слоям. Для модели с 31B параметров и несколькими десятками слоёв это означает квадратичный рост вычислений по глубине модели.
Текущие GPU, включая H100, справляются с инференсом стандартной Gemma 4 31B, но AttnRes требует больше памяти и операций. Автор проекта рассчитывает на следующее поколение ускорителей - B300, которые должны обеспечить достаточный запас по памяти и пропускной способности. Без них обучение и инференс остаются на грани практической применимости для одиночного разработчика.
Проблемы с квантованием тоже связаны с нестандартной архитектурой. Методы вроде EdgeRazor, которые сжимают модели до 1,88 бит на параметр через смешанное квантование и дистилляцию, пока не тестировались на архитектурах с attention-based routing. Стандартные инструменты квантования заточены под классические трансформеры с residual stream. Для AttnRes потребуется адаптация или разработка новых методов.
Длинные хвосты распределений - третья проблема. Top-K 12 логгитов сохраняет больше информации, чем one-hot, но всё равно отсекает хвост за пределами топ-12. Для частотных токенов это не критично. Для редких - может привести к потере способности модели генерировать специфическую лексику, имена, термины. Автор признаёт этот риск и рассматривает его как компромисс между вычислительной эффективностью и полнотой знаний.
Moonshot Kimi K3 и другие подтверждения перспективности attention-based routing
AttnRes - не первый проект с идеей attention-based routing. Kimi K3 от Moonshot использовал похожий механизм под названием Attention Residuals. В Kimi K3 этот подход позволил улучшить обработку длинных контекстов и повысить стабильность обучения при масштабировании модели.
Детальный разбор архитектуры Kimi K3 показывает, что Attention Residuals там сочетались с другими инновациями: Stable LatentMoE с 896 экспертами и механизмом Quantile Balancing для равномерной загрузки экспертов, Gated MLA для компрессии KV-кэша, и KimiDeltaAttention для линейного внимания. AttnRes фокусируется именно на замене residual stream, но общее направление мысли совпадает.
Успех Kimi K3 - модель показала результаты на уровне лидеров индустрии - подтверждает, что attention-based routing имеет практический потенциал. Открытость кода AttnRes под Apache 2.0 позволяет сообществу развивать идею параллельно с проприетарными разработками Moonshot.
Как начать экспериментировать: открытый код и следующие шаги
Код проекта доступен на GitHub под лицензией Apache 2.0 - это разрешает коммерческое использование и модификацию. Минимальные требования для запуска пока высоки: нужен GPU с большим объёмом памяти, желательно из линейки H100 или ожидаемых B300. Для экспериментов на меньших масштабах автор рекомендует начинать с модели поменьше - например, Gemma 4 9B, если она появится в открытом доступе.
Структура кода разделена на модули: замена residual stream на attention-based routing, weaning-график с настраиваемым расписанием, дистилляция с top-K логгитами, генерация синтетических данных. Каждый модуль можно тестировать независимо.
Направления для контрибуции, которые автор считает приоритетными:
- Адаптация методов квантования под архитектуру с attention-based routing
- Оптимизация вычислений - снижение квадратичной зависимости от глубины
- Тестирование на других базовых моделях помимо Gemma 4 31B
- Эксперименты с разными схемами weaning-графика: косинусное затухание, ступенчатое, с тёплым стартом
Проект не обещает готового продукта. Это исследовательский код, который доказывает концепцию. Для ML-инженеров и исследователей, которые хотят понять, куда движутся архитектуры трансформеров за пределами стандартных рецептов, AttnRes - ценный полигон для экспериментов.