Архитектура "Декодер-Гибрид-Декодер" для эффективного рассуждения с длинной генерацией

Аннотация

Последние достижения в области языкового моделирования продемонстрировали эффективность моделей пространства состояний (State Space Models, SSMs) для эффективного моделирования последовательностей. Хотя гибридные архитектуры, такие как Samba и декодер-декодерная архитектура YOCO, показали многообещающие улучшения производительности по сравнению с трансформерами, предыдущие работы не исследовали потенциал повышения эффективности за счет совместного использования представлений между слоями SSM. В данной статье мы представляем Gated Memory Unit (GMU) — простой, но эффективный механизм для совместного использования памяти между слоями. Мы применяем его для создания SambaY, декодер-гибрид-декодерной архитектуры, которая включает GMU в кросс-декодер для совместного использования состояний чтения памяти из само-декодера на основе Samba. SambaY значительно повышает эффективность декодирования, сохраняет линейную временную сложность предварительного заполнения и улучшает производительность на длинных контекстах, при этом устраняя необходимость в явном позиционном кодировании. В ходе масштабных экспериментов мы демонстрируем, что наша модель показывает значительно меньшую необратимую потерю по сравнению с сильным базовым уровнем YOCO, что указывает на превосходную масштабируемость производительности в условиях крупномасштабных вычислений. Наша крупнейшая модель, улучшенная с помощью Differential Attention, Phi4-mini-Flash-Reasoning, демонстрирует значительно лучшую производительность, чем Phi4-mini-Reasoning, на задачах логического мышления, таких как Math500, AIME24/25 и GPQA Diamond, без какого-либо обучения с подкреплением, при этом обеспечивая до 10-кратное увеличение пропускной способности декодирования на промптах длиной 2K с длиной генерации 32K в рамках фреймворка вывода vLLM. Мы публикуем наш код для обучения на открытых данных по адресу https://github.com/microsoft/ArchScale.

English

Recent advances in language modeling have demonstrated the effectiveness of State Space Models (SSMs) for efficient sequence modeling. While hybrid architectures such as Samba and the decoder-decoder architecture, YOCO, have shown promising performance gains over Transformers, prior works have not investigated the efficiency potential of representation sharing between SSM layers. In this paper, we introduce the Gated Memory Unit (GMU), a simple yet effective mechanism for efficient memory sharing across layers. We apply it to create SambaY, a decoder-hybrid-decoder architecture that incorporates GMUs in the cross-decoder to share memory readout states from a Samba-based self-decoder. SambaY significantly enhances decoding efficiency, preserves linear pre-filling time complexity, and boosts long-context performance, all while eliminating the need for explicit positional encoding. Through extensive scaling experiments, we demonstrate that our model exhibits a significantly lower irreducible loss compared to a strong YOCO baseline, indicating superior performance scalability under large-scale compute regimes. Our largest model enhanced with Differential Attention, Phi4-mini-Flash-Reasoning, achieves significantly better performance than Phi4-mini-Reasoning on reasoning tasks such as Math500, AIME24/25, and GPQA Diamond without any reinforcement learning, while delivering up to 10x higher decoding throughput on 2K-length prompts with 32K generation length under the vLLM inference framework. We release our training codebase on open-source data at https://github.com/microsoft/ArchScale.

Архитектура "Декодер-Гибрид-Декодер" для эффективного рассуждения с длинной генерацией

Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation

Аннотация

Support