Nemotron-Elastic-12B — один чекпоинт для моделей 6B, 9B и 12B без дополнительного обучения

Nemotron-Elastic-12B от NVIDIA AI предлагает иной подход к семействам моделей: один обученный чекпоинт на 12 миллиардов параметров, который включает вложенные варианты на 9B и 6B. Вместо того чтобы обучать или дистиллировать каждую версию отдельно, команда создаёт эластичную модель, которую можно «вырезать» в меньшие субмодели по потребности, что экономит токены обучения и память при развёртывании.

Гибридная архитектура Mamba-2 и Transformer

Архитектура Nemotron Elastic опирается на дизайн Nemotron-H и сочетает блоки последовательностного SSM Mamba-2 с небольшим набором слоёв глобального внимания. Такое гибридное решение сохраняет возможности для вывода при длинных контекстах и одновременно позволяет структурно уменьшать компоненты сети.

Эластичность через маски и роутер

Эластичность реализована с помощью масок, которые управляют шириной и глубиной модели. Двоичные маски позволяют сокращать размер эмбеддингов, количество Mamba-хедов и каналов, внимание и промежуточные размеры FFN. Глубина уменьшается путём удаления слоёв согласно изученному ранжированию важности, при этом остаточные пути сохраняют поток сигнала.

Модуль роутинга предсказывает дискретные конфигурации для каждого бюджетного режима. Эти предсказания преобразуются с помощью Gumbel Softmax в маски, которые применяются к эмбеддингам, проекциям Mamba, проекциям внимания и матрицам FFN. Команда добавила ряд деталей для корректности SSM: групповую SSM-эластичность, согласованную с группировкой Mamba-хедов, разнородную эластичность MLP с возможностью разных промежуточных размеров по слоям и нормализованную MSE-метрику важности слоёв для решений по глубине. Меньшие варианты формируются как префикс в ранжированных списках компонентов, поэтому 6B и 9B являются истинными вложенными подсетями 12B родителя.

Двухэтапное обучение для задач рассуждения

Обучение ведётся с замороженным учителем — исходной моделью Nemotron-Nano-V2-12B. Студент Nemotron-Elastic-12B оптимизируется одновременно для трёх бюджетов (6B, 9B, 12B) с использованием дистилляции знаний и лосса языкового моделирования. Процесс состоит из двух этапов:

• Этап 1: короткий контекст, длина последовательности 8192, batch size 1536, примерно 65B токенов, равномерная выборка по трём бюджетам.
• Этап 2: расширенный контекст, длина 49152, batch size 512, примерно 45B токенов, неравномерная выборка в пользу полного 12B бюджета (веса 0.5 для 12B, 0.3 для 9B, 0.2 для 6B).

Второй этап особенно важен для задач рассуждения: на AIME 2025 6B улучшился с 56.88 до 68.13 (примерно +19.8% относительного прироста), 9B прибавил 9.7%, а 12B — 4.0% после расширенного контекстного обучения.

Бенчмарки и результаты

Nemotron Elastic оценивали на задачах, требующих интенсивного рассуждения: MATH 500, AIME 2024 и 2025, GPQA, LiveCodeBench v5 и MMLU Pro. Основные средние показатели pass@1 составили примерно 70.61 для 6B, 75.95 для 9B и 77.41 для 12B. 12B эластичная модель в среднем соответствует базовому NanoV2-12B (77.41 против 77.38), 9B близка к NanoV2-9B (75.95 против 75.99), а 6B показывает сильный результат для своего размера, несмотря на отсутствие отдельного обучения.

Экономия токенов и памяти

Ключевая цель Nemotron Elastic — снижение затрат. Получение вариантов 6B и 9B из родителя 12B требует порядка 110B токенов в одном elastic distillation прогоне. Для сравнения, предобучение NanoV2 для 6B и 9B заняло около 40T токенов, а компрессия NanoV2 с Minitron SSM использовала около 750B токенов. Команда отмечает примерно 360-кратное снижение по сравнению с обучением дополнительных моделей с нуля и около 7-кратного снижения по сравнению с базой компрессии.

Память для развёртывания также сокращается: хранение Nemotron Elastic 6B, 9B и 12B вместе требует 24GB BF16 весов, тогда как хранение NanoV2 9B и 12B требует 42GB — экономия примерно 43% при появлении дополнительного варианта 6B.

Практическое значение

Nemotron-Elastic-12B переводит задачу создания семейств моделей рассуждения в плоскость одной эластичной системы. Один чекпоинт даёт несколько конкурентных размеров, упрощая управление парком моделей, снижая затраты на токены и хранение и давая гибкость при развёртывании на устройствах с разными ресурсами.