SLAY-ASR, или как я перестал волноваться и полюбил тренировать модели

Как добавить аудио-модальность в LLMку максимально экономно? Рассказываю про серию попыток добиться совместимости эмбеддингов разной природы.

...Зачем?

Изначально этот проект был лишь одной задачкой в таск-трекере другого, более масштабного проекта – приложения «Killah» (да-да, как песня Леди Гаги), амбициями которого было создание редактора текста с локальными языковыми моделями. Мы планировали сделать фичу: надиктовываешь текст, и, в противовес обычным приложениям для диктовки, типа «Whisper Flow», этот текст ещё и редактируется в соответствии с твоими голосовыми командами. То есть в моменте можно было бы сказать «Давай остановись, ой, там поставь запятую, ой, убери эту всю часть, давай заново» – и текст бы генерировался уже отредактированным. К тому же, в голове звенела тарелками красивая идея о подаче аудио-информации на вход модели без потерь, в виде латентов: это бы снарядило LLM-редактора пониманием, например, что за интонация у говорящего, чтобы делать более точные правки. Но вот эта вот отдельная часть затянулась – на очень долгие полгода, и в приложение мы по итогу добавили обычное распознавание речи через Whisper. Про чудесное приложение я напишу в следующем хаброблоге.

Ресёрч

Как добавлять модальности к языковым моделям? Посмотрим на популярное: как обстоят дела с VLM. В визуальных мультимодальных моделях берётся изображение и кодируется отдельным энкодером. Далее векторы из этого энкодера подаются в Cross-Attention в виде матриц K и V, а Q приходит от текстовых токенов. Получается модель, которая во время тренировки сопоставляет визуальные и текстовые эмбеддинги, понимая картинки. То же самое делают и с аудио всякие «Квены». Но кросс-аттеншн работает хорошо тогда, когда у нас модель и энкодер обучаются с нуля вместе – оч дорого и долго. А вот чтобы взять уже готовый, предобученный аудио-энкодер, готовую языковую модель, и совместить их, делают по-другому: добавляют проектор – нейросеть между энкодером и LLM, которая, по сути, является переводчиком: эмбеддинги одной модели трансформирует в эмбеддинги другой.

Нашёл я, значит, по этой теме «потрясающие» работы с манящими названиями «An Embarrassingly Simple Approach for LLM with Strong ASR Capacity», «Prompting Large Language Models with Speech Recognition Abilities», «Large Language Models Are Strong Audio-Visual Speech Recognition Learners», «Instruction-Following Speech Recognition» – все они рассказывают о том, как же легко натренировать такую мультимодальную модель: «Забудьте свои кросс-энкодеры, забудьте сложные сетапы и кластеры видеокарт! Возьмите обычный перцептрон и натренируйте его как адаптер! Никаких сложных лоссов, никаких проблем. У нас всё просто получилось.»

Проблема в том, что эти методы не работают. По крайней мере, не воспроизводятся с аналогичными языковыми моделями и энкодерами. На то, чтобы это понять и уверенно заявить, у меня ушло немало времени. Поэтому я искренне не знаю, как они проводили обучение, но просто так, через ванильную кросс-энтропию и обычный MLP-проектор учить такие вещи нереалистично.

После краткосрочного периода отчаяния я решил во что бы то ни стало заставить этот проектор работать и сделать ✨SLAY-ASR ✨, по аналогии со «SLAM-ASR» – моделью в одной из вышеупомянутых работ.

Строим архитектуру

Языковая модель ожидает на вход последовательность эмбеддингов из своей таблицы – каждый токен (слово, кусок слова, пунктуация) кодируется как вектор фиксированной размерности, который модель знает с предобучения. Речь устроена совсем иначе: аудиоэнкодер выдаёт непрерывный поток представлений, по одному на каждые несколько миллисекунд звука, и их несравнимо больше, чем токенов в соответствующем тексте. Чтобы их совместить, нужен «мост» – проектор, который не только переводит один тип векторов в другой, но и агрегирует длинную аудиопоследовательность в разумное число токенов.

Начну с выбора энкодера – это отдельное решение, потому что две популярные устоявшиеся модели и, соответственно, два популярных подхода к кодированию аудио, и устроены они принципиально по-разному.

Wav2Vec2 обучался self-supervised на сырых аудиозаписях: он учится предсказывать маскированные участки сигнала. Его выученные репрезентации различают акустику: тембр, темп, интонацию, фонетику. Хорошо захватывает общие свойства речевого сигнала, но не оптимизировался под конкретную задачу транскрипции.

Whisper обучался с учителем – на задаче ASR с парами аудио, текст. Его внутренние представления более ориентированы на лингвистическое содержание: что именно говорится, а не как. В итоге для задачи транскрипции эмбеддинги Whisper оказываются стабильнее и осмысленнее – поэтому на нём я и остановился.

Использовал Whisper Medium: она выдаёт векторы размерностью 1024, примерно по одному на каждые 20 миллисекунд аудио. Чтобы сократить длину последовательности, конкатенировал каждые пять подряд идущих векторов: вместо одного 20ms-фрейма получался один 100ms-фрейм размерностью 5120 – в пять раз длиннее.

Эти длинные векторы подаются в проектор – двухслойный MLP с архитектурой 5120 → 2000 → 2000 → 2048, где 2048 – это размерность эмбеддингов Gemma 3 4B, которую я использую как LLM для редактирования. Эксперименты с промежуточными размерами кардинально на качество не влияли.

Gemma 3 4B: мультиязычная, относительно компактная, без ризонинга и прочей ненужной истории. Более того, я использовал только претрейн версию, потому что в основном проекте мне нужно было дообучать её на своих сценариях.

Коктейль из датасетов

Начинал с LibriSpeech Clean – начитанного студийного английского. Всё окей, но модель не была робастной к малейшиму шуму и не умела распознавать русский. Нужно было разнообразие.

Тогда я написал WeightedInterleaveDataset – систему динамического потокового смешивания источников: шарды загружаются с HuggingFace Hub на лету и миксуются в реальном времени с весами, пропорциональными отношению размеров датасетов:

1. LibriSpeech (train.360) – база, чистый английский, ~54%

2. LibriSpeech (train.100) – дополнительный чистый английский, ~15%

3. Russian LibriSpeech – русский язык, ~28%

4. DisfluencySpeech – речь с оговорками, паузами и заиканиями, ~2% – так модель с самого начала училась на разнообразных сигналах.

Первая конфигурация тренировки

Технически то, что делается в наивном подходе – это teacher forcing: на вход LLM подаётся инструкция + аудио-векторы, за ними сразу идёт правильная транскрипция – всё в одной последовательности. Cross-Entropy лосс считается только по токенам транскрипции: инструкция и аудио замаскированы и штрафа не несут. Модель не угадывает ответ – мы показываем правильные токены один за другим и учим предсказывать следующий, оставляя перед глазами модели все правильные предыдущие токены. Это стандартный режим обучения авторегрессионных LLM.

Настроил инфраструктуру: mixed precision, gradient accumulation, bitsandbytes-квантизация, очистка memory cache. Много вылетал по памяти из-за накопленных градиентов.

Да и на метриках работало плохо: модель всё время выдавала несвязный текст, не имеющий отношения к записи. WER (доля правильно предсказанных слов) опускалась до, кхм, 300% и вставала – модель нащупывала optimization plateau, при котором достигнут хоть какой-то минимум, но дальше не двигалась.

Вторая конфигурация – LoRA

Первым делом я решил разморозить языковую модель, ведь именно это сильно помогало в статьях (правда, в этих статьях модель и до этого училась хорошо, но надежда умирает последней). Так в конструкцию добавилась LoRA (Low-Rank Adaptation) добавляет в каждый трансформерный слой небольшие пары матриц с сильно урезанным числом параметров, которые обновляются вместо исходных весов. Эффект от этого такой, что модель адаптируется к новым данным без переписывания всех своих весов – и без взрыва памяти. С этих пор LoRA остаётся включённой во всех конфигурациях.

Тоже не помогло.

Проблема глубже, чем недостаточная гибкость модели: даже при размороженной LLM аудио-сигнал через CE слишком слаб, чтобы преодолеть «prior» языковой модели – её внутренние представления о том, как должна выглядеть информация на входе. Даже если в аудио-эмбеддингах есть лингвистическая информация, она не обязана лежать там, где LLM умеет её читать. LLM предобучена на определённой геометрии токенов – и проектор, обучаясь только через CE по транскрипции, не получает достаточно прямого сигнала о том, куда именно нужно перенести аудио-векторы в этом пространстве.

Конфигурация 3 – Дискретизация и многостадийное обучение

Так, значит гипотеза в том, что модель не понимает входящие токены вообще. Так значит надо ей помочь! Можно ли вместо того, чтобы учить проектор выдавать произвольные непрерывные векторы, заставить его выдавать эмбеддинги, которые Gemma уже знает – из её собственной таблицы? Это идея из VQ-VAE (Vector Quantized Variational Autoencoder) – архитектуры для работы с дискретными латентными пространствами. В ней каждый вектор на выходе энкодера заменяется ближайшим вектором из конечного набора, которому там и придумали название «кодбук».

В моём случае кодбук – это таблица эмбеддингов самой Gemma, около 256-ти тысяч векторов.

Ещё немножечко теории, перед тем, чтобы это внедрить: операция квантования – взятие ближайшего вектора из таблицы эмбеддингов – недифференцируема: через неё не проходит градиент, который нужен для обучения. Обходится это через Straight-Through Estimator: при обратном распространении ошибки операция квантования просто игнорируется – градиент просто идёт к проектору напрямую, как будто квантования не было. Грубо, но работает.

«Но ведь это убивает всю идею красиво перетекающих латентов из модели в модель без потерь» – да. И с этим кое-что решили сделать ребята в работе «Bridging the Modality Gap: Softly Discretizing Audio Representation for LLM-based Automatic Speech Recognition» – у них же получилось натренировать эту систему! Ресёрчеры предложили выстроить несколько стадий квантизации, от известных векторов из кодбука до полностью свободно формирующихся эмбеддингов. Некий такой curriculum learning получается – модель, как в школе, привыкает к простым штукам в начале, чтобы потом спокойнее выучиться на сложных.

Stage 1 – Hard Discretization

Жёсткое квантование: один аудио-вектор → один ближайший эмбеддинг из таблицы Gemma. LLM получает ровно те векторы, с которыми работала при предобучении – ничего незнакомого.

Stage 2 – Soft (Top-k) Discretization

Разрешаю смотреть на топ-k (в реализации поставил десятку) ближайших кодов и брать их взвешенную сумму по softmax-вероятностям. Так модель получает плавную смесь эмбеддингов сразу и начинает понимать, что токены могут представляться и другими векторами, но эти векторы пока что напоминают известные.

Stage 3 – Continuous

Убираю кодовую книгу совсем: проектор подаёт непрерывный вектор напрямую в LLM. Это финальный режим. Расчёт на то, что если предыдущие стадии достаточно адаптировали LoRA-веса, модель сможет читать и непрерывные векторы из знакомой геометрической окрестности.

Stage 0 – Языковой прогрев

К сожалению, ничего из описанного серьёзно не помогает модели, и на всех четырёх стадиях WER не опускается ниже ста процентов. И это меня беспокоило не только потому что это много, а потому что это блин, БОЛЬШЕ СТА ПРОЦЕНТОВ ОШИБОК. То есть модель как минимум всегда выдает текст длиннее того что подаётся ей на вход. Скорее всего это из-за того, что аудио-токенов просто больше, чем аналогичных текстовых. Что я решил с этим сделать:

Моя версия Gemma 3 4B не дообучена под инструкции. Когда я подавал ей промпт «транскрибируй» и ждал ответ в нужном формате, она не понимала, что делать – потому что не встречала такого формата при предобучении. Добавил Stage 0: аудио в нём не подаётся вовсе, просто учу модель переписывать поданный текст точь-в-точь по инструкции:

Transcribe the following dictated text: [текст]

и вставлять в конец токен <EOS>, чтобы училась кончать вовремя.

LoRA-слои адаптировались к формату задачи. WER на нулевой стадии был около двадцати процентов. На самом деле, это всё ещё очень много для задачи «тупо перепиши этот текст», но что имеем, то имеем.

Это в целом улучшило метрики на следующих стадиях, но всё равно не помогло стабильно пробить плато в 100% WER.

Регуляризации

Подход с квантованием векторов предполагал добавление регуляризатора, заставляющего сближать эмбеддинги проектора к эмбеддингам из кодбука. Это буквально метод из VQ-VAE – никакой ментальной нагрузки чтобы принять решение его добавить не потребовалось.

Commitment loss

Это штраф, добавляющийся в лосс-функцию за то, что выход проектора уходит далеко от ближайшего эмбеддинга в таблице. Штраф выражается в виде MSE между выходом проектора для одного аудио-фрейма и ближайшим к нему вектору из таблицы эмбеддингов. Это регуляризация, потому что она не учит проектор транскрибировать – только удерживает его выходы в знакомой LLM геометрической окрестности.

Визуализации

Чтобы понять, почему после всех тренировочных прогонов лосс не падал куда надо, я начал строить визуализации и смотреть, как вообще расположены аудио-эмбеддинги относительно текстовых после выхода из проектора.

Впервые выпала возможность воспользоваться «в бою» классикой: t-SNE – алгоритмом снижения размерности, который при проекции в 2D старается сохранить локальную структуру данных. Векторы, которые близки в исходном высокоразмерном пространстве, окажутся близкими и на картинке. А вот расстояния между кластерами на t-SNE ничего не значат, потому что он пытается эти расстояния увеличить для лучшей различимости кластеров. Визуализировал я два набора векторов для каждого примера: спроецированные аудио и квантизованные текстовые эмбеддинги.

И вот что я увидел:

1. Коллапс: проектор схлопывается – все точки, иллюстрирующие спроецированные эмбеддинги, с каждым разом сближаются друг с другом всё больше и больше. Это может значить, что модель решает пойти лёгким путём и все данные представить одним набором очень похожих векторов, которые достаточно хорошо могут универсально описать все данные. Получается, модель одинаково забивает на содержание и генерирует текст произвольно.

2. Геометрический разрыв: спроецированные векторы занимают область пространства, которая вообще не пересекается с аудио-эмбеддингами.

Попытки юзать контрастивные методы я откладывал: мне казалось, что они в прошлом. Я почему-то опасался, что контрастив принципиально не сможет восстановить нужное распределение – он сколлапсирует информацию, которая не нужна для разделения позитивных и негативных пар, но нужна для понимания оттенков речи. К тому же, в популярных VLM и ALM типа BLIP2, LLaVA и Falcon3-Audio контрастив используется только для предобучения энкодера, а не для финального выравнивания с LLM.

Поэтому я начал пытаться управлять этой геометрией явно – через регуляризации, в дополнение к commitment loss.

SWD (Sliced Wasserstein Distance)

Проблему геометрического разрыва я начал решать через Sinkhorn Optimal Transport (OT) loss – очень точный, но вычислительно сложный алгоритм нахождения транспортного плана между конкретными точками двух множеств. Добавляя его в лосс можно было помочь точкам притягиваться друг к другу и наблюдать это на визуализациях. Правда, OT притягивал больше отдельные точки, чем весь кластер вместе, да и вычислял он свои расстояния очень медленно. Поэтому очень быстро он был заменён хорошей альтернативой: SWD.

Как работает это «расстояние Вассерштайна»: берётся тысяча случайных единичных направлений в пространстве эмбеддингов. Для каждого направления проецируются все аудио-векторы – получается набор чисел – и отдельно все текстовые векторы. Расстояние между двумя такими одномерными распределениями вычислить легко: отсортируй оба и сравнивай поэлементно (это расстояние Вассерштайна в одном измерении). SWD – среднее таких расстояний по всем направлениям, и на каждом шаге обучения направления новые. Это чуткий способ нахождения расстояния, не требующих попарных сравнений.

Со временем на t-SNE его эффект становился виден, когда два облака векторов начали приближаться друг к другу и пересекаться. Однако одновременно с этим проецирующиеся эмбеддинги начали группироваться в кластеры и сжиматься, расходясь с нужными текстовыми векторами.

Посмотрев на другие метрики, коих я логировал десятками в отчаянных попытках сделать хоть какие-то выводы, я заметил, что из 256 000 доступных эмбеддингов модель активно использовала только 50–100. Большинство аудиозаписей с очень разным содержанием отображались в одни и те же несколько токенов, то есть выходы проектора действительно были неразнообразны.

Code Entropy Loss

Чтобы заставить модель использовать больше разных эмбеддингов, я добавил Code Entropy Loss: максимизировал энтропию распределения по используемым токенам из таблицы. Если все двести тысяч кодов используются равномерно – энтропия максимальна.

Но и здесь модель нашла лазейку: начала равномерно использовать все коды, но случайно. На t-SNE облако стало равномерным, но подкластеры рассыпались слишком далеко и не совпадали с точками текстовых эмбеддингов.

Посмотрев на распределение значений по отдельным осям во всех спроецированных эмбеддингах (ещё одна из десятков метрик), увидел, что многие оси имеют почти нулевую дисперсию – по этим направлениям все аудио-векторы одинаковы. Если можно так сказать, это representation collapse в рамках отдельных осей: модель нашла несколько направлений, которые варьируются, и забила на остальные. Значит, добавляем ещё регуляризаций!

VICReg

Добавил VICReg (Variance-Invariance-Covariance Regularization) – штраф за низкую дисперсию вдоль каждой координатной оси эмбеддинга:

Здесь – значения -й координаты по всем векторам в батче, – размерность эмбеддинга (2048), – стандартное отклонение. Штраф ненулевой только тогда, когда стандартное отклонение по каждой оси меньше одного – то есть лосс толкает каждую ось быть хоть сколько-то разнообразной.

Модель обошла и это: начала генерировать высокую дисперсию по всем осям, но векторы разбегались хаотично, и информации о звуке в них как будто бы не было.

Балансировка весов

Всё это время я крутил веса лоссов, смотрел на t-SNE и примеры генераций. Типичный сценарий был такой: поднимаю entropy loss – использование кодов растёт, WER тоже. Опускаю entropy loss, убираю VICReg – collapse возвращается. Добавляю SWD – форма выравнивается, содержание не улучшается. Ни одна комбинация не давала стабильной сходимости на протяжении нескольких эпох подряд: каждый регуляризатор исправлял одну проблему и создавал другую. Модель постоянно находила новую лазейку.

Финал: контрастив спасает мир

Помните, как я в начале отложил контрастивные методы в долгий ящик? Так вот, после ещё тщетных попыток добавить CTC-выравнивание, сменить MLP на Q-Former и подобной дичи, я решил всё же попробовать полноценное контрастивное обучение – не как регуляризатор с маленьким весом поверх CE, а как основной сигнал выравнивания.

Разница принципиальная. Все предыдущие регуляризаторы работали со всеми эмбеддингами сразу и независимо, меняя общую статистику распределений. Это косвенный, слабый сигнал, потому что эмбеддинги не знают и не думают о существовании своих соседей, и поэтому двигаются несогласованно. Контрастив же, оказывается блин, использует информацию о парах отдельно в каждом батче, и задаёт конкретные попарные отношения: этот аудио-вектор должен быть близко к эмбеддингу своей транскрипции и далеко от всех остальных в батче. Это прямой геометрический сигнал – и он кратно сильнее с большим батчем, потому что больше негативных примеров означает более жёсткое требование к положению каждого вектора.

Лёгким движением набитой руки добавил InfoNCE с симметричным лоссом – оптимизировал оба направления: аудио к тексту и текст к аудио по вот такой страшной формуле:

Здесь мы берём среднее из транскрипций в батче, сопоставляя положительную пару из аудио-вектора и его транскрипции с суммой всех косинусных сходств аудио-вектора с остальными транскрипциями в батче.

С большим размером батча лосс начал падать – ровно, монотонно, без скачков. Это было похоже на чудо после месяцев рваных графиков. WER опустился до 35%.

Это далеко от современных коммерческих ASR-систем. Но это первый прогон, где модель слышит – и видно, что дальше нужно просто докинуть вычислений и данных, а не перепридумывать архитектуру.

Вы только посмотрите на пример из логов последнего прогона:

• Референс: «воистину еврейки молодой мне дорого душевное спасенье»

• Предсказание: «воистину венрейки молодой мне поро душевное спасенье»

  Или:

• Предсказание: «воистину еврейки молодой мне дороги душевное спассье молодой мнерый душевную спасенье»

Фонетика на месте, слова же – почти. Там, где модель ошибается, она мило ошибается как человек с плохим слухом. Но главное, что она уже слышит.