Квантизация ASR-моделей

Мотивация запуска проекта заказчиком: всё большую популярность приобретают “умные” устройства: часы, колонки, камеры, холодильники. Как правило, ресурсы на таких устройствах ограничены, что приводит к проблеме адаптации нейросетевых моделей под конкретные машины, когда необходимо упростить или модифицировать архитектуру сети с целью уменьшения размера модели и ускорения её работы на стадии применения. Однако, ограничения на нейросеть могут быть и более существенными. Так, вычислитель, на котором предполагается запускать модель, может поддерживать вычисления только в низкобитной точности. В таком случае необходима квантизация модели. В настоящее время, с одной стороны, квантизация — это активно развивающаяся область. С другой стороны, мало работ посвящено квантизации нейросетевых моделей, основанных на трансформерах. В проекте требуется предложить честный метод низкобитной квантизации SOTA архитектуры ASR трансформера без существенной потери качества итоговой модели на общедоступном наборе данных для обучения и валидации ASR моделей LibriSpeech.

Описание исходной ситуации:

Есть архитектуры моделей ASR трансформеров и их реализации;
Есть методы квантизации нейросетевых моделей (не обязательно трансформеров);
Квантизация трансформеров — относительно слабо исследованная область;
Нередки случаи публикации статей, где используется нечестная квантизация. Контроль честности квантизации — отдельный вопрос, требующий внимания;
Применение квантизации “в лоб” сильно ухудшает качество нейросети;
Некоторые модули в модели более чувствительны к квантизации, чем другие (например, Embedding слой, SoftMax).

Цели проекта:

подобрать стратегию квантизации ASR модели, основанной на трансформерах, при необходимости внеся также изменения в архитектуру сети, с тем чтобы качество квантизованной модели (метрика WER) было не сильно хуже, чем качество модели с полной точностью (допускается ухудшение качества в несколько процентов).

Решение MIL Team: взять доступную реализацию SOTA архитектуры модели ASR на основе трансформеров и встроить в неё квантизацию. Для успешной и честной квантизации необходимо, во-первых, реализовать квантованные версии для всех используемых внутри модели Torch модулей. Во-вторых, нужен удобный инструмент по замене оригинальных Torch модулей на квантованные. В-третьих, надо провести ряд экспериментов по выбору лучшей стратегии квантизации.

Для построения модели были использованы:

реализованная в открытом репозитории Fairseq архитектура ASR трансформера, представленная в статье Transformers with convolutional context for ASR.
набор данных LibriSpeech для обучения модели по распознаванию речи, состоящий из пар аудио и текстовых файлов с английской речью.

Результаты моделирования: under NDA.

Заказчик: under NDA

Технологический стек: Python (PyTorch, torchaudio, Fairseq, SentencePiece, DeepLabV3, SRCNN)