FastWave: компактная диффузионная модель для ускорения обработки аудио

Команда студентов мастерской по сервисам и платформам ИИ Инженерно-математической школы НИУ ВШЭ совместно с исследователями VK представила FastWave — диффузионную модель для задачи Bandwidth Extension (BWE). Целью разработки стало создание инструмента, способного повышать частоту дискретизации любого аудио до студийных 48 кГц при минимальных вычислительных затратах. В отличие от существующих решений, которые часто требуют мощных GPU и длительного времени обучения, FastWave весит всего 1,3 млн параметров и демонстрирует высокую скорость инференса.

Разработчики выбрали в качестве базовой архитектуры NU-Wave 2, так как она поддерживала режим any-to-48 kHz и обладала относительно компактным размером. Однако исходная модель требовала около двух недель обучения на двух GPU A100. Для оптимизации процесса команда применила методологию EDM (Elucidating the Design Space of Diffusion Models), разработанную исследователями NVIDIA. Адаптация EDM к задаче аудио-суперразрешения включала использование σ-параметризации, предобуславливания входов и выходов, а также лог-нормального распределения уровней шума при обучении. Эти изменения позволили достичь пиковых показателей метрик, затратив втрое меньше итераций по сравнению с исходной моделью.

Дополнительные оптимизации были внесены на уровне архитектуры. Обычные свёртки были заменены на depthwise separable convolutions, что снизило параметрическую сложность. Для компенсации потери кросс-канального взаимодействия в блоках модели были внедрены слои Global Response Normalization (GRN). В результате число параметров сократилось с 1,8 млн до 1,3 млн, а вычислительная сложность за один вызов уменьшилась с 19,0 до 12,9 GFLOPs. При этом качество реконструкции звука сохранилось на высоком уровне.

Ключевым преимуществом FastWave является сокращение числа шагов инференса. Если оригинальный NU-Wave 2 требует 8 вызовов модели для генерации аудиофрагмента, то FastWave стабильно выдаёт хорошие результаты уже за 4 шага. Это напрямую удваивает скорость обработки. Эксперименты проводились на датасете VCTK, содержащем около 44 часов аудио от 110 спикеров. Обучение модели заняло до 30 часов на одной GPU V100, что значительно быстрее по сравнению с недельным процессом для базовой модели.

Сравнительный анализ показывает, что FastWave демонстрирует результаты, сопоставимые с моделями SOTA. По метрике SNR модель уступает лишь незначительно крупным аналогам, таким как FlowHigh и NU-Wave 2. При этом по количеству параметров FastWave в 38 раз меньше FlowHigh. Что касается скорости, то при обработке звука в четыре шага модель работает в шесть раз быстрее реального времени. Это делает её пригодной для потоковых приложений на устройствах с GPU.

Проект был реализован под руководством доцента ФКН НИУ ВШЭ Максима Каледина и исследователя AI VK Александра Тараканова. Результаты работы оформлены в научную статью, которая принята к публикации на конференции Interspeech 2026. Исходный код и предобученные веса модели доступны для скачивания на GitHub, что позволяет исследователям и разработчикам воспроизвести эксперименты и интегрировать модель в собственные проекты.