Предсказание нескольких токенов (MTP, когда модель дополнительно предсказывает несколько последующих токенов) уже стало стандартным рецептом предобучения: оно улучшает результаты на downstream-задачах и может служить черновиком для спекулятивного декодирования. Авторы точно разложили градиент перекрестной энтропии MTP на две части: член само-дистилляции KL, который притягивает голову MTP к распределению следующего слова самой модели (полезный, улучшает спекулятивное декодирование), и остаточный член с очень низким отношением сигнал/шум, который мешает обучению и усугубляется с ростом масштаба. Soft-MTP использует «сквозную маршрутизацию градиентов» (straight-through gradient routing), сохраняя только член KL, повторно используя существующие вычисления перекрестной энтропии, не требуя матрицы логитов для всего словаря и совместим с ядрами слияния перекрестной энтропии.
Лайки 9 · Звезды на GitHub 6
Причины
- Минимальные изменения: для стандартного MTP это одна строка кода, без добавления параметров; авторы заявляют об отсутствии измеримых дополнительных накладных расходов; поддерживаются как MTP в стиле DeepSeek, так и в стиле Meta.
- Авторы утверждают, что на каждом масштабе от 167M до 1.2B достигается более низкая перплексия следующего слова, более быстрое спекулятивное декодирование и более высокая точность при нулевом обучении (zero-shot); после дообучения модель размером 1.2B демонстрирует более сильные способности к математическим и кодовым рассуждениям, а по большинству метрик не уступает прямому минимизированию KL.
- Воспроизводимость: четыре масштаба × три цели (NTP, MTP, Soft-MTP), всего 12 контрольных точек размещены на Hugging Face; конфигурации обучения и скрипты оценки спекулятивного декодирования доступны; в каталоге tests/ можно проверить совпадение градиентов с явным KL в точности float64; лицензия MIT.
Использование
- Код: git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5 (предварительно токенизированный набор данных FineWeb-Edu, обучение на 31.5B токенов для каждой конфигурации, минимум 4 шардов, при сборке требуется около 120 ГБ дискового пространства); оптимизатор Muon требует PyTorch версии 2.9 или выше.
- Обучение: torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp (замените на hard_mtp / direct_kl / ntp для сравнения).
- Готовые модели: AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True), токенизатор — оригинальный gpt2; для оценки спекулятивного декодирования используйте python -m src.eval_spec --ckpts ….
- Чтобы интегрировать в свой код обучения, посмотрите функцию combined_mtp_loss() в файле src/models/loss.py.
Риски
- Выпущены только небольшие исследовательские модели размером менее 1.2B, обученные исключительно на FineWeb-Edu; это не готовые диалоговые модели; эффективность на больших масштабах не проверена.
- Основная ценность проявляется на этапе предобучения; чтобы получить выгоду, нужно обучать модель с нуля; по умолчанию обучение и оценка интегрированы с Weights & Biases.
- В аннотации не указано конкретное ускорение спекулятивного декодирования; результаты приведены по данным авторов; в README отмечается, что ускорение по времени выполнения зависит от GPU, поэтому сравнения следует проводить на одном и том же оборудовании.