Проблема масштабирования малых моделей
Языковые модели с количеством параметров от 10 до 100 миллионов привлекательны для локального инференса и быстрых экспериментов, однако большинство из них просто копируют стандартный блок Transformer, не адаптируя его под малый масштаб. Это приводит к неэффективному использованию ресурсов и ухудшению качества на коротких контекстах. Авторы работы предлагают решение — архитектуру Wiola, которая меняет три ключевых компонента каждого слоя.
Три инновационных компонента Wiola
Новизна модели сосредоточена в трех модулях, заменяющих стандартные элементы:
- Spiral Rotary Positional Encoding (Spiral RoPE): Модификация стандартных частот Rotary, где добавляется медленно растущий пер-дименсиональный фактор. Это заставляет фазовые траектории расходиться, улучшая распознавание дальних зависимостей (long-range discrimination) без добавления новых параметров.
- Gated Spiral Attention: Вводит скалярный гейт на голову (per head), зависящий от контента. Гейт вычисляется на основе каузальной кумулятивной статистики потока запросов (query stream). Это обеспечивает мягкий выбор голов (soft head selection) с минимальными вычислительными затратами.
- Butterfly Feed-Forward Block: Заменяет стандартный слой расширения на мультипликативное взаимодействие и внутриблочный bypass-путь. Это улучшает поток градиентов в неглубоких стеках, сохраняя количество параметров на уровне блока с четырехкратным Gated Linear Unit (GLU).
Математическая строгость и эффективность
Авторы не только описывают архитектуру, но и формализуют бюджеты параметров и вычислений. Ключевое преимущество Wiola — доказанное точное и численно верифицированное эквивалентность между полным обучением последовательности (full sequence training) и кэшированным авторегрессионным декодированием. Это означает, что при инференсе не используется никаких аппроксимаций, что критично для стабильности работы на edge-устройствах.
Результаты и доступность
Модель обучена и протестирована на стандартном корпусе Tiny Story. Исследователи опубликовали полностью воспроизводимый протокол обучения и оценки. Референсная реализация доступна как open-source пакет, готовый к публикации весов.
| Компонент | Стандартный Transformer | Wiola 13M | Преимущество |
|---|---|---|---|
| Позиционное кодирование | RoPE (стандартное) | Spiral RoPE | Улучшение дальних зависимостей без доп. параметров |
| Attention | Standard Multi-Head | Gated Spiral Attention | Контент-адаптивный выбор голов, низкая стоимость |
| Feed-Forward | MLP / GLU | Butterfly Block | Лучший градиентный поток, компактность |
Источник: arXiv cs.CL ↗