Variable-Width Transformers: широкие края, узкая середина — как новая архитектура экономит 22% FLOPs
Большинство современных языковых моделей построены на трансформерах, где каждый слой получает одинаковое количество параметров и вычислений. Инженеры масштабируют модель в двух измерениях: ширина (размерность скрытого состояния, d_model) и глубина (число слоев, L). При этом фундаментальное допущение остается неизменным — каждый слой получает одинаковый бюджет. Команда из MIT и MIT-IBM Watson AI Lab задалась вопросом: что если это не оптимально?
Идея: почему слоям не нужна одинаковая ширина
Исследователи из MIT и IBM предложили архитектуру bowser (от формы ⧈ — широкая по краям, узкая в центре). Ключевое наблюдение: параметры слоя масштабируются квадратично с шириной, а вычисления attention — линейно. Это означает, что при фиксированном бюджете параметров можно уменьшить среднюю ширину слоя, сохранив общее число весов.
Формализуем. Параметры доминируют линейными projection-матрицами: QKV-projections, output projection, MLP up- и down-projections. Все они масштабируются как O(d²), где d — ширина слоя. Attention FLOPs (attention scores computation) масштабируются как O(d) при фиксированной длине последовательности. Это означает, что при одинаковом budget параметров форма ⧈ (средние слои уже) дает меньшую average width, а значит — меньше attention FLOPs.
Математика: E[d̄²] < d² при неравномерном распределении. Средний квадрат ширины всегда меньше квадрата средней ширины, если ширины не равны. Практический результат — 22% экономия FLOPs без потери качества.
Авторы ссылаются на исследования, показавшие, что разные слои трансформера выполняют разные функции: ранние слои специализируются на синтаксисе и POS-tagging, поздние — на семантике и генерации. Это не случайность, а устойчивый паттерн в representations learned by Transformers. Логично, что разным функциям нужно разное количество параметров — но до bowser никто не проверял это систематически для всей ширины блока.
Как это работает: parameter-free residual resizing
Главная техническая сложность — изменение размерности между слоями. Если на одном слое ширина 1024, а на следующем 768, нужно как-то «сжать» состояние. Наивное решение — обучить projection-матрицу для каждого перехода — работает, но требует дополнительных параметров и усложняет обучение.
Исследователи использовали parameter-free подход: при сужении состояния координаты просто обрезаются, при расширении — восстанавливаются из наиболее свежей позиции, где они обрабатывались. Skip-connections проходят через все слои без изменений, что позволяет координатам «перепрыгивать» узкие слои напрямую. Это критически важно для gradient flow — без skip connections через всю глубину обучение было бы невозможно.
Формула: если слой j имеет ширину d_j, а предыдущий слой обработал координату i, то элемент восстанавливается из последнего слоя, где d_k >= i. Если ни один предыдущий слой не обрабатывал эту координату — она заполняется нулями.
Авторы провели ablation studies, показавшие, что этот подход превосходит альтернативы: projection-слой с обучаемыми весами и padding нулями. Fixed residual dimension (global residual шире любого слоя) также критичен — иначе narrow-слой создает bottleneck, нарушающий gradient flow и ухудшающий качество.
Эксперименты: от 200M до 3B параметров
Команда обучила модели четырех размеров: 200M, 500M, 1B и 2B параметров (dense), плюс MoE-вариант на 3B (1B active). Для каждого размера — отдельная конфигурация depth/hidden dimension: 200M — 16 слоев, hidden 640; 500M — 24 слоя, hidden 960; 1B — 32 слоя, hidden 1280; 2B — 40 слоев, hidden 1600. Batch size варьировался от 512 до 4096, последовательность — 4096 токенов, данные — DCLM corpus.
Обучение велось до 2.5 Chinchilla-optimal tokens (50× параметров), то есть 2B-модель обучалась на 100B токенов. AdamW optimizer, learning rate с warmup и power decay schedule, SwiGLU activation, без bias terms в linear projections — стандартные современные практики.
bowser сохраняла форму ⧈: широкие early/late layers, узкая середина. Bottleneck position подбиралась: оптимальный bottleneck layer index = 0.75L, bottleneck dimension = 0.3d_base. Для 40-слойной 2B-модели это означает: bottleneck на слое 30 с width 0.3 × 1600 = 480, края занимают остальные 32 слоя.
Результат: при сопоставлении по числу параметров bowser показывает примерно 3% улучшения language modeling loss по сравнению с baseline constant-width. Это звучит скромно, но в контексте scaling laws — заметное улучшение при том же budget.
Главный результат по эффективности: 22% редукция FLOPs на loss-matched scaling curves и 15% снижение KV cache memory и I/O cost. KV cache — ключевой для inference: он масштабируется линейно с width, batch size и sequence length. При том же качестве модели меньший KV cache означает возможность обслуживать больше параллельных запросов на том же hardware.
Авторы подчеркивают: FLOPs считали теоретически, без учёта GPU kernel overhead. Extra kernel launches для variable-width операций добавляют практический overhead. Однако 22% FLOPs-редукция на matmul-heavy операциях, где доминирует собственно умножение, должна сохраняться в значимой степени.
Почему форма ⧈, а не △ или ▽
Исследователи протестировали разные формы: △ (растущая, narrow→wide), ▽ (сужающаяся, wide→narrow), ◇ (растущая затем сужающаяся) и ⧈ (сужающаяся затем растущая). Форма ⧈ устойчиво лучшая на всех масштабах.
Возможная интерпретация: начальные слои собирают информацию из входа — нужно больше емкости для encoding. Средние слои обрабатывают и трансформируют representations — емкость менее критична, можно сузить. Финальные слои производят output — нужна широта для генерации. Это объясняет, почему именно сужающаяся-затем-растущая форма, а не наоборот.
Это противоречит более ранним работам о layerwise-аллокации только FFN-размерностей, где выгоднее было увеличивать средние слои. Bowser перераспределяет всю ширину блока — QKV, attention, FFN — и получает иной оптимальный профиль. Механизм разный: для FFN-only оптимизации middle layers важны для expressiveness, но когда мы перераспределяем весь block width, картина меняется.
Практические ограничения: почему не всё так просто
Архитектура имеет known limitations. Extra kernel launches для variable-width операций: каждый narrow-слой требует отдельного kernel launch, что добавляет overhead на GPU. Текущие GPU-оптимизации — tensor parallelism, pipeline parallelism, activation checkpointing — плохо совместимы с неравномерной шириной: они требуют uniform shape для эффективного partition across devices.
На уровне software-фреймворков: PyTorch, JAX и другие не имеют нативной поддержки per-layer variable width. Нужна либо встроенная поддержка в фреймворках, либо существенная оптимизация custom CUDA kernels. Авторы ожидают, что purpose-built kernels закроют большую часть gap между теоретической и практической эффективностью.
При этом bowser всё ещё matmul-rich: основные операции — те же matrix multiplications, просто с разными размерами. Это означает, что hardware utilization может быть высоким, если kernels оптимизированы.
Что это значит для индустрии
Результаты демонстрируют ранее незамеченную степень свободы в дизайне трансформеров. При текущем масштабировании до десятков миллиардов параметров, 15–22% экономия — значительные числа в терминах GPU-часов и стоимости inference. Для comparison: оптимизация одного параметра batch size дает 5–10% улучшения; 22% — существенно.
Особенно интересна связь с MoE-трендом. Bowser + MoE (показан как 3B MoE variant в статье) комбинирует два измерения оптимизации: routing между experts и перераспределение параметров within layers. Для будущих архитектур это открывает третье измерение — не только how many parameters, но и which layers get how many.
Для inference-практиков 15% редукция KV cache — конкреная выгода. При typical serving setup с sequence length 2048 и batch size 16, KV cache для 2B модели занимает ~3.2GB на запрос. 15% редукция означает ~450MB экономии на каждый запрос, что при высокой concurrency (100+ simultaneous requests) дает существенное снижение memory pressure. Это особенно важно для long-context моделей, где KV cache доминирует memory footprint.
Другое практическое применение — edge deployment. На устройствах с ограниченной памятью (мобильные, embedded) уменьшение KV cache с 15% экономией может означать разницу между возможностью запустить модель и нет. Combined с quantization (INT8/INT4), variable-width дает compounding benefits: параметров меньше, attention FLOPs меньше, KV cache меньше.
Для практиков главный takeaway: при проектировании новой архитектуры стоит экспериментировать с неравномерным распределением ширины. Форма ⧈ — простая отправная точка с измеримым эффектом на моделях от 200M до 3B параметров. Не нужно переписывать фреймворки, чтобы попробовать: можно parameterize layer widths и обучить стандартным tools.
Как bowser сочетается с existing optimizations
Важно понимать: bowser — это orthogonal optimization к большинству существующих техник. Это не замена FlashAttention или KV cache quantization, а дополнение.
С MoE: показано в статье, 3B MoE variant работает. MoE экономит active parameters (только 1B из 3B active per token), bowser экономит average layer width. Комбинация может дать compounding benefits: меньше active FLOPs + меньше KV cache.
С quantization: INT8/INT4 quantization уменьшает memory footprint параметров и activations. Bowser уменьшает average layer width, что означает меньше elements в attention scores и KV cache. Эффекты совместимы и потенциально усиливают друг друга.
Со sparsity: structured sparsity (pruning) или dynamic sparsity — ортогональное измерение. Можно применять bowser shape поверх pruned модели.
Чего bowser НЕ заменяет: архитектурные инновации типа state spaces (Mamba), recurrent mechanisms, или linear attention не заменяются bowser. Это оптимизация того, как распределять parameters внутри transformer block.
Почему это важно для scaling laws
Chinchilla scaling laws утверждают: optimal training tokens ≈ 50 × parameters. Но это для constant-width моделей. Bowser нарушает это предположение: при том же parameter count получается better loss (3% improvement) с тем же compute budget. Это означает, что optimal token-to-parameter ratio может быть ниже для variable-width architectures.
Представим practical implications. Если bowser-2B показывает то же quality, что constant-width 2B, но с 22% меньше FLOPs, то при том же compute budget можно либо (a) обучить на 22% больше tokens и получить better model, либо (b) получить ту же model quality за меньшие деньги. Это имеет смысл только если infrastructure limitations (kernel overhead) будут решены.
Также интересно: bowser дает ~15% меньше KV cache memory. Это критично для inference cost при long context. Standard 4K context моделей с bowser фактически работают как модели с 3.4K context по memory, что может снизить inference cost для приложений с variable-length inputs.
FAQ
Разве не изучали variable depth раньше? Исследования оптимального соотношения width/depth проводились, но все сохраняли constant width within each layer. Bowser — первая работа, систематически исследующая неравномерное распределение всей ширины блока (QKV, attention, FFN) across depth при одинаковом бюджете параметров.
Почему 22% экономия FLOPs, если параметры те же? Параметры доминируют линейными projection-матрицами, которые масштабируются как O(d²). При фиксированном бюджете параметров форма ⧈ требует меньшей average width, так как E[d̄²] < d². Attention FLOPs масштабируются линейно с width — поэтому падают пропорционально.
Можно ли комбинировать bowser с FlashAttention? Концептуально — да, FlashAttention совместим с variable width. Практически — нужна поддержка variable sequence lengths per layer в attention kernel, что не реализовано в стандартных библиотеках. Это добавляет engineering complexity, но не является принципиальным барьером.
Почему bottleneck position = 0.75L? Эмпирически оптимальная позиция для 40-слойной 2B-модели. Слой с index 30 (0.75L) имеет width 0.3d_base. Широкие края занимают ~65% слоев (0–26 и 33–39), узкая середина — ~35%. Это обеспечивает достаточно широкие края для encoding/decoding при минимальной общей average width. Bottleneck position может варьироваться для разных глубин — требует tuning.
Насколько сложно имплементировать bowser? Для research purposes — умеренно. Нужен custom layer width parameterization и modification backward pass для skip connections. Для production — значительно сложнее: текущие фреймворки не поддерживают variable width natively, и нужно либо custom kernels, либо acceptance of overhead.
Были ли tested еще какие-то shapes помимо ⧈? Да, авторы тестировали △ (narrow→wide), ▽ (wide→narrow), ◇ (wide→narrow→wide), и flat (constant width как baseline). На всех масштабах ⧈ была лучшей. Интересно, что △ (противоположная форма) работала хуже baseline, подтверждая, что конкретная форма имеет значение, а не просто неравномерность.
Итог
Архитектура bowser показывает, что uniform width в трансформерах — это convenience, а не оптимальность. При одинаковом числе параметров форма ⧈ дает лучший loss и на 22% меньше вычислений. Главный барьер — инфраструктурный: текущие фреймворки и GPU-ядра оптимизированы под constant width. Когда это изменится, variable-width может стать стандартным инструментом в арсенале дизайнеров архитектур.
Ключевой takeaway: при проектировании больших моделей стоит рассматривать не только total parameter count, но и их распределение по слоям. Края шире, середина уже — простая идея с измеримым эффектом на реальных моделях от 200M до 3B параметров.