Сегодня разбираем работу о важной проблеме в Audio Language Models. Если дать модели один и тот же запрос в текстовом и аудиоформате, качество ответа на аудио обычно заметно хуже. Это видно на бенчмарках: аудиодомен почти всегда проседает относительно текстового.
Большие аудиоязыковые модели обычно строятся поверх текстовых LLM: к ним добавляют аудиоэнкодер и модуль выравнивания между модальностями. Предполагается, что этого достаточно, чтобы аудио и текст оказались в одном семантическом пространстве. Но на практике это работает неидеально: при одинаковом смысле входа модель рассуждает по аудио хуже, чем по тексту.
В этой работе авторы хотят передать в аудиомодальность текстовые reasoning-способности и сделать так, чтобы при аудиовходе модель рассуждала как можно ближе к тому, как она рассуждает при текстовом.
Для этого берут один и тот же вопрос, получают его текстовую и аудиоверсию, прогоняют через модель и сравнивают распределения. Это нужно, чтобы генерация при аудиовходе была максимально похожа на генерацию при текстовом.
Важно, что всё происходит on-policy. Для аудиовхода модель сначала сэмплирует ответ, а дальше сравнение идёт именно по этой траектории. Эти же префиксы подаются в модель с текстовым входом, чтобы посмотреть, как она продолжила бы ту же последовательность.
То есть здесь не идут по «правильной» текстовой траектории учителя, а обучают модель на тех состояниях, в которые она реально попадает при генерации по аудио. За счёт этого можно исправлять ошибки, которые возникают именно в аудиорежиме.
Для сравнения распределений используют reverse KL, а именно KL(p_audio || p_text).
Такой лосс особенно наказывает случаи, когда аудиоветвь даёт высокую вероятность токенам, которые текстовая ветвь считает маловероятными. Модель в первую очередь отучают генерировать продолжения, которые сильно расходятся с текстовым поведением.
Авторы также показывают, что расхождения между модальностями возникают неравномерно. Основные ошибки появляются в начале генерации, а затем тянутся по всей последовательности. Кроме того, сильный вклад дают не все токены, а лишь небольшое число наиболее важных шагов.
Из этого рождается основной трюк работы — взвешивание token-level loss. Во-первых, усиливают вклад токенов с наибольшим расхождением между аудио- и текстовым распределениями: берут top-K токенов, в работе K = 20. Во-вторых, больший вес дают ранним токенам, потому что ошибка в начале рассуждения чаще всего ломает всю дальнейшую траекторию. В итоге reverse KL взвешивается и по важности токена, и по его позиции.
Но этого авторам кажется недостаточно, потому что токенного уровня может не хватить, так как модель всё равно способна прийти к неправильному ответу в целом. Поэтому добавляется sequence-level.
Здесь генерируются ответы для аудио, и judge-модель проверяет, совпадают ли они семантически с текстовым ответом. Этот сигнал превращается в reward, и дальше обучение идёт через GRPO. В итоге комбинируются два сигнала: по токенам и по всей последовательности, а финальный objective — это их сумма.
По экспериментам видно, что подход стабильно уменьшает разрыв между аудио и текстом и делает аудиоответы ближе к текстовым. Обучаются только на математическом датасете, но улучшения переносятся и на general knowledge задачи.
По сравнению с обычной дистилляцией метод меньше ломает дополнительные аудиоспособности модели, то есть не ухудшает понимание звуков, музыки и прочего.
Ярослав Ведерников