Разреженная смесь экспертов в обработке звука: разбор архитектуры SESAME для Speech Enhancement

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

Я подробно изучил технические особенности этой разработки, проверил эмпирические результаты экспериментов и разобрал, почему популярные в языковых моделях (LLM) приёмы роутинга не всегда подходят для звукового сигнала. В этой статье представлен подробный разбор устройства SESAME, механизмов маршрутизации и практических уроков, полученных при её обучении.

Постановка задачи шумоподавления

С математической точки зрения задача улучшения речи формулируется достаточно просто. Имеется входной зашумлённый сигнал $y$, состоящий из чистой речи $x$ и аддитивной помехи $n$ случайного характера:

$$y = x + n$$

Шум $n$ может принимать самую разную форму: от монотонного гула вентилятора и импульсных щелчков до реверберации помещения, артефактов сжатия и фоновых голосов сторонних людей. Цель алгоритма — построить максимально точную оценку $\hat{x}$, которая приближается к оригинальному голосовому сигналу как по объективным метрикам, так и по естественности звучания.

Современные компактные нейросети обычно работают в частотно-временной области. С помощью коротковременного преобразования Фурье (STFT) входной сигнал переводится в комплексную спектрограмму $Y$. Модель предсказывает восстановленный спектр, после чего обратное преобразование (iSTFT) возвращает обработанный аудиосигнал во временную область.

Проблема монолитных архитектур и переход к смеси экспертов

Наиболее эффективные компактные модели последнего времени — такие как MP-SENet, PrimeK-Net или SEMamba — являются монолитными. Это означает, что один и тот же фиксированный набор весов обрабатывает абсолютно каждый частотно-временной токен (TF-токен) спектрограммы. Для сети нет разницы, попал ли в данный фрагмент чистый гласный звук, резкий импульсный шум или участок паузы — вычисления выполняются одним и тем же преобразованием.

Хотя глубокие нелинейные сети способны в некоторой степени адаптироваться к разным входам, их параметрическая ёмкость остается ограниченной. В мире крупных языковых моделей эту проблему давно решили с помощью разреженной смеси экспертов (Sparse MoE). В таких архитектурах стандартные полносвязные слои (FFN) заменяются набором параллельных подсетей («экспертов»), а специальный роутер для каждого токена выбирает лишь небольшую их часть. Это позволяет нарастить общую ёмкость модели, не увеличивая число активных операций на шаг инференса.

В задаче Speech Enhancement концепция Sparse MoE до сих пор оставалась малоизученной. Основная гипотеза авторов SESAME заключается в том, что разреженная смесь экспертов задаёт более эффективный индуктивный сдвиг для аудио: эксперты обучаются специализироваться на конкретных акустических сценариях (стационарный шум против импульсного, низкий SNR против высокого, звонкие согласные против глухих), не перегружая процессор или видеокарту лишними вычислениями.

Требования и ключевые метрики качества

При проектировании SESAME были зафиксированы строгие рамки и критерии оценки:

• Работа с частотой дискретизации 16 кГц (стандарт для широкополосной речи);
• Превосходство над базовой моделью MP-SENet при равном или близком общем объёме параметров;
• Уменьшение числа активных параметров на инференсе по сравнению с плотными аналогами аналогичного размера.

Качество оценивалось по общепринятому набору метрик: PESQ (wideband, главная перцептивная оценка от -0.5 до 4.5), композитным метрикам CSIG (сохранность речи), CBAK (подавление фона), COVL (общее качество), а также сегментному отношению сигнала к шуму (SSNR). Дополнительно отслеживались внутренние параметры роутера — энтропия распределения и нагрузка на отдельных экспертов.

Базовая модель: почему выбор пал на MP-SENet

В качестве фундамента архитектуры была выбрана модель MP-SENet. Это решение обусловлено чистой модульной структурой исходной сети и её высокими результатами в классе фазозависимых систем. Кроме того, FFN-блоки внутри MP-SENet можно заменить на MoE-модули без изменения внешнего тракта STFT, энкодера и декодеров.

Стандартный пайплайн MP-SENet включает следующие этапы:

1. Преобразование STFT дает сжатую по амплитуде спектрограмму и фазовую составляющую.
2. Свёрточный Dense Encoder формирует 64-мерное представление признаков.
3. Последовательность из четырёх TS-блоков (Time-Frequency Self-Attention) обрабатывает контекст поочерёдно вдоль временной и частотной осей. Внутри каждого подблока применяется FFN на базе двунаправленной GRU и линейной проекции.
4. Два декодера восстанавливают сигнал: Mask Decoder отвечает за амплитудную маску, а Phase Decoder предсказывает развёрнутую фазу.
5. Обучение происходит с участием метрического дискриминатора в GAN-режиме для оценки перцептивного качества.

Главным узким местом MP-SENet оставалась монолитность TS-блоков, которую и призван устранить метод SESAME.

Архитектура SESAME: интеграция MoE в процесс обработки

Вычислительный граф SESAME сохраняет ключевые параметры MP-SENet (размер окна Ханна, шаг hop=100, частоту 16 кГц, энкодер и декодеры). Главные изменения затронули структуру FFN внутри первых TS-блоков.

Шаг 1. Двухступенчатый модуль MoE-FFN

Прямолинейное копирование стандартного FFN (состоящего из BiGRU и линейного слоя) для каждого эксперта привело бы к раздуванию модели, так как рекуррентный блок BiGRU требует наибольшего числа параметров. Чтобы избежать этого, авторы разделили MoE-FFN на два уровня:

• Первая ступень (общая): Единая двунаправленная GRU (64 → 128 в каждом направлении, итого 256) обрабатывает входной поток. Её веса являются общими для всех экспертов, поэтому рекуррентный временной контекст просчитывается ровно один раз.
• Вторая ступень (специализация): Каждый эксперт представляет собой лёгкую двухслойную перцептронную сеть: Linear(256 → 256) → GELU → Linear(256 → 64) с собственными весами. Таким образом, специализируются только дешёвые полносвязные «головы».

Шаг 2. Маршрутизация на основе глобального профиля шума

Одно из ключевых нововведений SESAME — учет общего акустического фона при роутинге. Локальный TF-токен сам по себе не содержит исчерпывающей информации о природе шума во всей аудиозаписи. Один и тот же фрагмент спектра может оказаться как элементом речи, так и тональной помехой или щелчком.

Чтобы дать роутеру контекст, из сжатой амплитудной спектрограммы формируется глобальный профиль звука: спектрограмма усредняется по времени, проходит через слой Linear + GELU и превращается в 64-мерный вектор $g$. Этот вектор подаётся исключительно в роутер, не затрагивая основные вычисления экспертов. Это позволяет маршрутизатору принимать решения с учётом общего характера звуковой сцены.

Шаг 3. Отказ от Expert Choice в пользу Token Choice

В сфере обработке текста популярна стратегия Expert Choice, при которой каждый эксперт самостоятельно выбирает топ-$C$ наиболее подходящих токенов, что гарантирует идеальный баланс нагрузки. Однако для аудио эта схема оказалась неприменимой.

При Expert Choice часть токенов неизбежно остаётся без обработки экспертами (пропускается). В тексте пропущенный токен может пройти через остаточные (residual) связи без потери смысла. В аудиосигнале каждый TF-токен физически связан с фрагментом звучания. Пропуск даже доли токенов приводит к разрывам спектра, которые на слух воспринимаются как высокочастотные щелчки и фазовые артефакты.

В SESAME используется исключительно стратегия Token Choice: каждый токен гарантированно направляется к $k$ экспертам из пула $E$ (в итоговой модели $k=2, E=4$), веса экспертов нормируются, а сброс токенов полностью отсутствует.

Шаг 4. Выборочное размещение MoE-блоков

Эксперименты показали, что добавление MoE во все четыре TS-блока снижает итоговое качество. На это есть две причины:

1. Вычислительная: При 8 модулях MoE-FFN вспомогательный лосс балансировки начинает перекрывать основной градиент обучения, что дестабилизирует сеть.
2. Физическая: Начиная с ранних слоев, модель разделяет сигналы по типам шума. На глубоких уровнях признаки становятся абстрактными, и разница между типами помех сглаживается — там достаточно стандартного плотного FFN.

В результате MoE-модули были установлены только в первые два TS-блока ($TSB_0$ и $TSB_1$) — по одному в частотную и временную ветви (всего 4 MoE-FFN). Третий и четвёртый блоки остались монолитными.

Шаг 5. Трёхуровневая балансировка нагрузки

Для предотвращения ситуации, когда роутер отправляет все токены нескольким «любимым» экспертам (expert collapse), в SESAME применили комбинацию трёх методов:

1. Switch Transformer loss: Градиентный штраф за отклонение распределения токенов от равномерного.
2. Адаптивное смещение логитов: Безградиентная корректировка, при которой перегруженный эксперт получает штраф к логитам на каждом шаге и выбирается реже.
3. Router Z-loss: Штраф за чрезмерный рост абсолютных значений логитов, обеспечивающий численную стабильность функции softmax.

Шаг 6. Итоговая конфигурация модели

В таблице ниже приведены ключевые параметры финальной сборки SESAME:

За счет разреженной структуры при общем объёме весов в 3.53 млн на каждый токен фактически работает только 2.87 млн параметров. Это почти на 20% ниже, чем у плотной модификации MP-SENet large (3.59 млн).

Результаты экспериментов и аблиационные исследования

Модель обучалась на датасете VoiceBank+DEMAND (11 572 обучающих и 824 тестовых аудиозаписи, 16 кГц) с использованием оптимизатора AdamW и косинусного затухания скорости обучения.

Что не сработало в ходе экспериментов

Исследователи проверили несколько популярных гипотез, которые на практике показали ухудшение результатов на 100-эпоховых тестах:

• Дополнительный SSL-лосс на базе HuBERT: Ожидалось, что предобученная речевая модель улучшит перцептивные признаки, но PESQ просел на 0.07.
• Multi-scale дискриминатор во временной области: Замена спектрального дискриминатора на временной (в стиле HiFi-GAN) снизила метрику SSNR с 10.69 до 10.50.
• VAD-маскирование потерь: Использование Silero VAD для подсчёта ошибки только на речевых сегментах привело к снижению PESQ на 0.04.
• Стратегия Expert Choice: Показала худший результат с падением PESQ на 0.14 из-за выпадающих токенов.
• Увеличение числа экспертов ($E=8, E=16$): Разрастание пула экспертов привело к дестабилизации обучения и снижению качества. Конфигурация $E=4, k=2$ оказалась наиболее устойчивой.

Специализация экспертов

Для оценки реального разделения труда исследователи замерили энтропию роутера и фактическую загрузку экспертов. Начиная с 8-й тысячи шагов нагрузки всех 4 экспертов стабилизировались на уровне 0.5 (±0.01), а энтропия выровнялась на значении ~1.29 (при теоретическом максимуме 1.39 для 4 экспертов). Это подтверждает, что роутер не вырождается и задействует всех экспертов, формируя устойчивые предпочтения для разных типов токенов.

Сравнение с аналогами на VoiceBank+DEMAND (400 эпох)

SESAME обошла базовые и расширенные версии MP-SENet по всем ключевым метрикам. Модель достигла показателя PESQ 3.64, затрачивая при этом на 20% меньше вычислительных ресурсов на шаг инференса по сравнению с аналогичной по размеру MP-SENet large.

Выводы

Разработка SESAME показала, что разреженная смесь экспертов успешно работает в задачах улучшения речи. Главные практические выводы исследования:

• Прямой перенос текстовых методик MoE на аудио не работает — алгоритм Expert Choice непригоден для звука из-за пропуска токенов, ведущего к слышимым артефактам. Token Choice является единственным надёжным выбором.
• Размещение MoE-блоков должно быть точечным: наибольший эффект достигается на первых слоях сети, обрабатывающих первичные акустические признаки.
• Специализация экспертов позволяет улучшить перцептивные метрики речи без увеличения вычислительной нагрузки при инференсе.

В дальнейшем разработчики планируют протестировать архитектуру на более крупных и разнообразных датасетах (включая DNS Challenge), а также изучить возможность интерпретируемой привязки конкретных экспертов к понятным акустическим категориям шума.

Источник: habr.com