Speech Info

Chunk-wise Attention Transducers for Fast and Accurate Streaming Speech-to-Text

В NVIDIA есть несколько сотрудников, которые стабильно пишут интересные статьи об ASR в целом и RNN-T в частности. Примеры таких работ — FastConformer, TDT, WIND. Сегодня расскажем о CHAT, суть которого также в улучшении RNN-T. Но сначала вспомним, что это такое.

Recurrent Neural Network Transducer — архитектура для распознавания и перевода речи (а в одной статье внезапно предлагают использовать её и для синтеза), состоящая из энкодера, prediction network и joint network. Работает следующим образом:

1. Энкодер принимает на вход звук, чтобы выдать последовательность эмбеддингов.

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

3. Joint Network использует эмбеддинг от Prediction Network и один из эмбеддингов от энкодера, чтобы предсказать следующий токен.

4. Полученный токен подаём назад в Prediction Network, чтобы получить новый эмбеддинг. Если же был предсказан специальный токен <BLANK>, то оставляем эмбеддинг от Prediction Network в покое и берём уже следующий по порядку эмбеддинг от энкодера.

5. Повторяем шаги 3 и 4, пока не кончатся эмбеддинги от энкодера или пока Joint Network не предскажет <EOS>.

Сегодняшняя статья строится на двух логичных и справедливых утверждениях:

1. В реальных системах распознавания речи звук поступает чанками, а не отдельными токенами.

2. Для предсказания следующего токена в Joint Network можно и полезно использовать более одного эмбеддинга от энкодера за раз.

Руководствуясь первым, авторы предлагают использовать не стандартную для LLM треугольную каузальную маску, а блочно-треугольную. С неё, помимо возможности смотреть назад, токены в рамках блока (чанка) могут смотреть друг на друга. Сама идея не тянет на новаторскую, но она ощутимо подкрепляет следующую.

Нововведение статьи основано на втором утверждении. Обычно Joint Network незамысловатый: сумма, конкатенация или линейный слой с нелинейностью для агрегации эмбеддингов и голова для предсказания следующего токена. Авторы для агрегации решили использовать cross-attention, где эмбеддинг от Prediction Network становится Q, а чанк (!) эмбеддингов от энкодера становится K и V (к этому чанку также конкатенируется токен из нулей чтобы модель могла использовать его для генерации токена <BLANK>). Таким образом мы получаем чанк, токены которого смотрели друг на друга на протяжении всего энкодера и который используется в Joint Network целиком.

Это даёт победу сразу по нескольким направлениям:

🔴Обучение RNN-T требует построения решетки из всех пар эмбеддингов от энкодера и от Prediction Network, чтобы считать лосс по всем возможным траекториям. Метод уменьшил «энкодерную» сторону этой решётки в число раз, равное размеру чанка (в статье — 12).

🔴Пиковое использование GPU-памяти уменьшилось почти в два раза и обучение ускорилось на 36%.

🔴Инференс также ускорился на 69%, потому что мы используем чанки целиком и достаточно проставить 1 <BLANK> для всего чанка вместо каждого эмбеддинга от энкодера.

🔴ASR WER уменьшился на 6,3% и AST BLEU вырос на 18% относительно аналогичных классических RNN-T. Авторы объясняют это тем, что возможность использовать более одного эмбеддинга от энкодера за раз даёт необходимый (особенно для перевода) контекст, позволяющий решать задачу более качественно. Я (автор обзора) считаю, что также благодаря использованию чанка целиком, модели не нужно паковать всю полезную контекстную информацию в каждый токен и она может извлечь и упаковать больше информации в чанк.

Бонус: на недавно прошедшей ICASSP нашему человеку повезло столкнуться с авторами этой статьи. На вопросы «Действительно ли необходимо добавлять токен чисто из нулей для предсказания <BLANK>? Как вы это поняли?» один из авторов ответил: «Inspiration, I had a kind of feeling I should add zeros», — и дальше не углублялся.

Николай Коновальчук ❣ Специально для Speech Info