Автор: Александр Гончаренко, СТО enot.ai

Введение

В 2022 году вышла модель от OpenAI, которая, по их заявлениям, решила задачу Speech2Text. Хоть у этого заявления и были противники, модель хорошо себя зарекомендовала. Я ей пользовался, когда нужно было сделать s2t «нестандартных» текстов (обзор новостей из мира игр).

И спустя год Hugging Face ****выкатили более лёгкую версию, которую назвали Distil-Whisper. Авторы говорят, что им удалось ускорить модель в 6 раз, подрезав параметры на 50% и потеряв всего 1% WER на out of distribution (OOD) тестовых данных. В этой статье мы разберёмся, как у них получился такой результат.

Distil-Whisper: основные составляющие

В качестве основных компонентов Distil-Whisper можно выделить следующие:

  1. Грамотная инициализация модели-ученика;
  2. Дистилляция по словам (авторы это называют word-level KD);
  3. Дистилляция не только по словам, но и по наиболее вероятной последовательности, сгенерированной учителем (авторы называют это sequence-level KD и псевдо-лейбеллинг);
  4. Фильтрация предсказаний, полученных при помощи псевдо-лейбеллинга;
  5. Ускорение самой большой модели и только декодера.

Кратко про архитектуру

Whisper — это Seq2seq модель, енкодер и декодер которой являются трансформерами. На вход в енкодер ей подаются мел спектрограммы, а токены предсказываются уже в авторегрессионном режиме в декодере.

Вся вычислительная сложность сосредоточена именно в декодере, который гоняется в “цикле for”. Его и оптимизируют авторы. Отсюда и получается ускорение x6 при выбрасывании 50% параметров.

Рисунок 1. Архитектура сети Whisper

Рисунок 1. Архитектура сети Whisper

Грамотная инициализация

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

В качестве примера рассмотрим такую ситуацию: есть модель из 6-и слоев, мы хотим натренировать модель из 3-х слоев. Для этого берем 0, 3, 5 (первый и последний включены, а по середине опционален). Если из 2-х слоев, то 0 и 5 (максимально разнесенные по архитектуре). Если надо взять всего 1 слой, то берем 0, не задумываясь. Это исследовалось в одной из предыдущих работ авторов. При этом саму конфигурацию слоёв мы не меняем.

KD по словам

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