SESAME: как разреженная смесь экспертов улучшает шумоподавление речи

Студенты мастерской по сервисам и платформам ИИ в Инженерно-математической школе НИУ ВШЭ и VK представили модель SESAME (Speech Enhancement with Sparse Adaptive Mixture of Experts). Она решает задачу улучшения качества речи (Speech Enhancement, SE) — подавления шумов, мешающих разборчивости аудиосигнала. В качестве базовой архитектуры взята MP-SENet, одна из лучших компактных моделей в этом классе. Ключевое отличие — замена стандартных полносвязных блоков внутри трансформера на разреженный MoE-модуль. Код проекта опубликован на GitHub.

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

Проблема монолитных архитектур, таких как MP-SENet, PrimeK-Net и SEMamba, в том, что один и тот же фиксированный набор весов применяется ко всем частотно-временным токенам сигнала. Неважно, что представляет собой токен — фрагмент чистой речи, шипение или резкий щелчок, — обработка идёт одинаково. В языковых моделях эту проблему давно решили с помощью разреженной смеси экспертов: плотный FFN-блок заменяется набором параллельных экспертных подсетей, а лёгкий роутер активирует лишь некоторых из них для каждого токена. Ёмкость модели растёт, а стоимость инференса почти не увеличивается.

В Speech Enhancement sparse MoE исследован заметно слабее, чем в NLP. Авторы SESAME выдвинули гипотезу, что разреженная смесь экспертов — более удачный индуктивный сдвиг для SE, чем плотный универсальный FFN. Эксперты могут специализироваться на узких акустических подзадачах: стационарный шум против нестационарного, низкий SNR против высокого, звонкие участки против глухих. При этом вычисления на инференсе не растут пропорционально числу параметров.

Перед выбором архитектуры авторы зафиксировали требования. Модель должна работать на частоте дискретизации 16 кГц, превосходить по качеству базовую MP-SENet при сопоставимом общем числе параметров и использовать при инференсе меньше активных параметров, чем плотная модель того же размера. Метрики — классические для SE: PESQ (перцептивная оценка), CSIG/CBAK/COVL (сохранность речи, качество подавления фона, общее качество), SSNR (сегментное отношение сигнала и шума). Дополнительно диагностировались разреженность и поведение роутера через энтропию маршрутизации и нагрузку экспертов.

Выбор MP-SENet в качестве базы объясняется прагматично: у неё чистая модульная структура, одни из лучших результатов в классе компактных фазозависимых архитектур, а FFN-модули можно заменить на MoE, не трогая энкодер, декодеры и STFT-пайплайн. В MP-SENet зашумлённый сигнал проходит STFT, затем свёрточный Dense Encoder, после чего идёт цепочка из четырёх TS-блоков (Time-frequency Self-attention). Каждый блок — пара трансформерных подблоков с вниманием вдоль времени и частоты, внутри которых FFN реализован как двунаправленная GRU с линейной проекцией. На выходе два декодера: Mask Decoder предсказывает амплитудную маску, Phase Decoder — развёрнутую фазу. Обучение идёт с метрическим дискриминатором, предсказывающим прокси-оценку PESQ в GAN-схеме.

Авторы заменили FFN в TS-блоках на разреженный MoE-модуль. В статье они делятся опытом, какие трюки из мира LLM работают в аудио, а какие нет. Например, Expert Choice — метод, где эксперты сами выбирают токены, — оказался неэффективным для SE. Причины авторы разбирают в деталях, подчёркивая, что перенос архитектурных решений между модальностями требует проверки.

Результаты работы показывают, что SESAME превосходит базовую MP-SENet по качеству при меньшем числе активных параметров на инференсе. Это подтверждает гипотезу о преимуществе разреженной смеси экспертов для шумоподавления. Для отрасли это шаг к более эффективным моделям обработки аудио, которые могут работать на устройствах с ограниченными вычислительными ресурсами, сохраняя высокое качество.

Проект интересен не только техническими результатами, но и открытостью: код доступен на GitHub, что позволяет другим исследователям воспроизвести эксперименты и развивать подход. Авторы — студенты, что говорит о растущем уровне подготовки молодых специалистов в области ИИ в России.