Меньше шума, больше смысла: опыт SESAME для Speech Enhancement / Хабр

Привет, Хабр! Мы — Александр Забурдаев, Ксения Золина, Даниил Марценюк, Илья Трофименко, Илья Кулешов — студенты мастерской по сервисам и платформам ИИ в Инженерно-математической школе НИУ ВШЭ и VK. Хотим поделиться опытом, который мы накопили, пытаясь перестать кормить весь речевой сигнал одной и той же нейросети целиком, и дать разным экспертам разбираться с разными типами шума.

Речь пойдёт о задаче улучшения качества речи (Speech Enhancement, SE) и о том, как получилось интегрировать в неё разреженную смесь экспертов (sparse Mixture of Experts, MoE). Это архитектура, которая в языковых моделях уже стала стандартом, а в обработке аудио до сих пор не прижилась.

Модель мы назвали SESAME — Speech Enhancement with Sparse Adaptive Mixture of Experts. Она построена на базе MP-SENet. В ней стандартные полносвязные блоки внутри трансформера заменены на разреженный MoE-модуль. Код проекта доступен на GitHub.

Зачем вообще понадобилась очередная модель шумоподавления, когда их и так десятки — расскажем ниже. А заодно разберём, какие трюки из мира LLM в аудио работают, а какие нет (например, Expert Choice) — и почему.

Постановка задачи

Задача шумоподавления формулируется просто: есть зашумлённый сигнал $y\left[n\right]=s\left[n\right]+d\left[n\right]$ , где $s\left[n\right]$ — чистая речь, а $d\left[n\right]$ — шум произвольной природы: стационарный фон, импульсные помехи, реверберация, конкурирующие голоса, артефакты сжатия. Нужно построить оценку $\hat{s}\left[n\right]$ , которая максимально близка к чистой речи по восприятию и разборчивости.

Большинство современных компактных моделей работает в частотно-временной области: к сигналу применяется STFT, получается комплексный спектр $Y(t,f)=\vert \ Y(t,f) \ \vert \cdot e^{i argY(t,f)}$ , сеть предсказывает улучшенное представление, а сигнал восстанавливается обратным преобразованием (iSTFT).

Проблема монолитной архитектуры

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

На бумаге это нормально. В плотной архитектуре все входы обрабатываются одним и тем же набором параметров. Хотя нелинейная сеть может адаптировать преобразование к входу, её параметрическая ёмкость остаётся общей для всех акустических режимов. MoE вводит более явную conditional specialization: разные токены могут использовать разные подмножества параметров.

В языковых моделях похожую проблему давно решили разреженной смесью экспертов. Плотный FFN-блок внутри трансформера заменяется набором параллельных экспертных подсетей, и лёгкий роутер для каждого токена активирует лишь $k\ll E$ экспертов. Ёмкость растёт, а стоимость инференса — почти нет.

В Speech Enhancement sparse MoE исследован заметно слабее, чем в языковых моделях и некоторых других speech-задачах.

Гипотеза: разреженная смесь экспертов — более удачный индуктивный сдвиг для SE, чем плотный универсальный FFN. Это позволяет экспертам специализироваться на узких акустических подзадачах — стационарный vs нестационарный шум, низкий vs высокий SNR, звонкие vs глухие участки, — не оплачивая это пропорциональным ростом вычислений на инференсе.

Требования и метрики

Прежде чем выбирать архитектуру, зафиксируем ограничения и метрики.

Модель должна:

работать на частоте дискретизации 16 кГц (стандарт для wideband-оценки);
превосходить по качеству базовую MP-SENet при сопоставимом общем числе параметров;
использовать при инференсе меньше активных параметров, чем плотная модель того же размера — в этом, собственно, и весь смысл MoE.

Метрики — классические для SE:

PESQ (wideband) — основная перцептивная метрика, от −0,5 до 4,5.
CSIG/CBAK/COVL — композитные метрики (сохранность речи, качество подавления фона, общее качество).
SSNR — сегментное отношение сигнала и шума.

Разреженность и поведение роутера дополнительно диагностируем внутренними метриками: энтропией маршрутизатора и нагрузкой экспертов.

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

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

Кратко, как устроена MP-SENet:

Зашумлённый сигнал проходит STFT → амплитуда $\vert \ Y \ \vert$ (со степенным сжатием $\vert \ Y \ \vert^{\beta}$ , $\beta\approx 0,3$ ) и фаза.
Двухканальное представление подаётся в свёрточный Dense Encoder (→ 64 признака).
Дальше — цепочка из четырёх TS-блоков (Time-frequency Self-attention). Каждый блок — пара трансформерных подблоков: один с вниманием вдоль времени, другой — вдоль частоты. Внутри подблока — FFN в виде двунаправленной GRU + линейная проекция.
На выходе два декодера: Mask Decoder предсказывает амплитудную маску (мультипликативно применяется к зашумлённой амплитуде), Phase Decoder — развёрнутую фазу через арктангенс.
Обучение — с метрическим дискриминатором, который предсказывает прокси-оценку PESQ (GAN-схема).

Монолитная природа тяжеловесных TS-блоков — это и есть то ограничение, которое мы решили преодолеть.

Разработка метода SESAME

Граф вычислений SESAME почти полностью повторяет MP-SENet: те же STFT-параметры ( $N_{FFT}=400$ , hop=100, окно Ханна, 16 кГц), тот же энкодер, те же декодеры. Модификации затрагивают только TS-блоки, в части которых стандартный FFN заменён на разреженный MoE-FFN.

Архитектура SESAME (MP-SENet с добавленными MOE-модулями)

Разберём ключевые решения по порядку. Каждое из них проверялось отдельно — без этого любое улучшение можно списать на «просто добавили параметров».

Шаг 1: сам модуль MoE-FFN и двухступенчатость

Самое наивное решение — скопировать FFN (в MP-SENet это BiGRU + Linear) целиком для каждого эксперта. Проблема: BiGRU — самая тяжёлая часть блока, и при экспертах размер модели взлетает несоразмерно.

Поэтому модуль сделан двухступенчатым:

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

Выход GRU конкатенируется с глобальным профилем шума (о нём ниже), и уже этот вектор попадает в роутер.

Шаг 2: роутинг, обусловленный шумом

Это, пожалуй, главное архитектурное отличие SESAME от классических языковых MoE.

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

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

Без этого вектора роутер пытается угадать режим шума по локальному спектральному паттерну, что неоднозначно — и в наших экспериментах маршрутизация заметно ухудшалась.

Примечание: на самом деле в нём содержится сжатое представление как шума, так и полезного сигнала. (Так как на входе у нас есть только зашумлённая спектрограмма.)

Шаг 3: Token Choice, а не Expert Choice — и это принципиально

В мире LLM популярна стратегия Expert Choice: каждый эксперт сам выбирает себе топ- токенов. Это даёт почти идеальную балансировку нагрузки «из коробки». Первой мыслью было взять именно её.

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

В аудио каждый TF-токен соответствует физическому фрагменту сигнала. Пропуск даже небольшой доли токенов приводит к разрывам во временной размерности — и на слух это превращается в специфические артефакты: высокочастотные щелчки и резкие скачки спектральной плотности в местах разрывов. В наших замерах Expert Choice дал наибольшее падение PESQ среди всех протестированных стратегий.

Поэтому в SESAME используется Token Choice: каждый токен гарантированно отправляется ровно к экспертам из , веса выбранных экспертов перенормируются, и сброса токенов (token dropping) нет вовсе.

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

Шаг 4: где размещать MoE — и почему не везде

Естественный соблазн — поставить MoE во все четыре TS-блока, чтобы усилить всю сеть. Мы попробовали. Получилось хуже.

Две причины:

Техническая. При 8 экземплярах MoE-FFN суммарный вспомогательный балансировочный лосс становится сопоставим с основным, зашумляя градиент. На масштабе ~3,5M параметров несколько экспертов в глубоких слоях обучаются почти идентичному поведению — разрешающая способность смеси просто не используется.
Содержательная. Ранние слои ближе к входу, и именно там различия между типами шума (стационарный, импульсный, речеподобный) выражены ярче всего. Глубокие слои оперируют уже абстрактными признаками, где эта разница частично нивелирована, — там обычного плотного FFN хватает.

Итог: MoE стоит только в первых двух TS-блоках (TSB₀, TSB₁), по одному экземпляру во временной и частотной ветвях каждого — итого 4 MoE-FFN. Два последних блока оставлены плотными. Это согласуется и с общей практикой MoE-трансформеров: MoE в подмножестве слоёв обычно выигрывает у MoE везде.

Шаг 5: многоуровневая балансировка нагрузки

Главная головная боль при обучении Token Choice — вырождение экспертов (expert collapse). На ранних шагах роутер сваливает почти все токены на удачно инициализированных экспертов. Те улучшаются быстрее и притягивают ещё больше токенов — положительная обратная связь. Вскоре половина экспертов мертва.

Решить одним механизмом не вышло. Поэтому в SESAME работают сразу три механизма для балансировки нагрузки на экспертов на разных уровнях:

Switch Transformer loss (градиентная). Штрафует отклонение нагрузки от равномерной: $\mathcal{L}_{bal}=\alpha\cdot E\cdot\sum \limits_{i}{f_iP_i}, \alpha=0,01$ , где — доля токенов у эксперта (без градиента), — средняя вероятность выбора (с градиентом).
Адаптивное смещение логитов (безградиентное). Буфер обновляется на каждом шаге: $b_i+=\eta\cdot(\frac{K}{E}-f_i), \eta=0,001$ . Перегруженный эксперт получает штраф к логитам и реже выбирается в дальнейшем
Router Z-loss (численная стабильность). Штраф за рост логитов: $\mathcal{L}z=\beta\cdot\mathbb{E}\left[\log^2\sum \limits _{i}^{}e^{l_i}\right], \beta=0,001$ . Без неё softmax на больших логитах численно нестабилен

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

Шаг 6: функция потерь

Базовый набор потерь унаследован из MP-SENet без изменений:

ошибка амплитуды;
фазовая ошибка — с anti-wrapping функцией, потому что фаза периодична и сравнивать её в лоб нельзя;
ошибка комплексной составляющей;
STFT-консистентность;
L1 норма разности сигналов во временной области;
метрический GAN-сигнал.

К ним добавляется суммарный вспомогательный лосс ${\mathcal{L}_{aux}}^{total}$ от всех четырёх MoE-FFN.

Коэффициенты менять не стали — чтобы не смешивать эффект MoE с эффектом перетюнинга лосса.

Итоговая конфигурация

Параметр	Значение
Базовая архитектура	MP-SENet
Эксперты	4
Активные на токен	2 (Token Choice)
MoE-FFN экземпляры	4 (в TSB₀, TSB₁)
Скрытый размер эксперта	256
Профиль шума	64-мерный, только в роутер
Всего параметров	3,53 M
Активные параметры на инференсе	2,87 M

Ключевая цифра — последняя: при общем объёме 3,53 M на каждый токен реально работает лишь 2,87 M, то есть почти на 20% меньше, чем у плотной модели сопоставимого размера (MP-SENet large, 3,59 M).

Тренировочные данные и сетап

Основной корпус — классический VoiceBank+DEMAND: 11 572 обучающих пары и 824 тестовых, ресемплирование до 16 кГц. Это стандарт для сравнения SE-моделей, на нём же обучалась базовая MP-SENet, так что сравнение честное.

Обучение: AdamW, lr $5\cdot10^{-4}$ , $\beta_1=0,8$ , $\beta_1=0,99$ , косинусное затухание с разогревом на 58k шагов. Покомпонентный анализ (абляции) гонялся на укороченном бюджете в 100 эпох, итоговая модель — на полном, 400 эпох.

Поиск конфигурации: что не сработало

Прежде чем показать финальные цифры, стоит рассказать, какие очевидные улучшения улучшениями не оказались. Все абляции — 100 эпох, фиксированный сид, одинаковое всё, кроме одной варьируемой детали.

Модификации функции потерь — все провалились

Мы пробовали сделать обучающий сигнал умнее тремя способами:

SSL-лосс на признаках HuBERT ( $W_{SSL}=0,2$ ). Идея: замороженная речевая модель кодирует перцептивно релевантные признаки, которых нет в спектральных потерях. На практике PESQ упал на 0,07.
Multi-scale дискриминатор во временной области (HiFi-GAN-стиль, даунсемплинг ×1/×2/×4) вместо спектрального Metric Discriminator. PESQ не изменился, но SSNR просел с 10,69 до 10,50.
VAD-маскирование потерь (Silero VAD, порог 0.5): штрафовать модель только на речевых сегментах, не наказывая за паузы. PESQ упал на 0,04.

Ни одно в итоговую конфигурацию не вошло. Мораль: на этом датасете и при этом бюджете эпох навороты лосса только мешают.

Стратегии роутинга и число экспертов

№	Конфигурация	PESQ↑	ΔPESQ
00	Базовая SESAME (E=4, k=2, TSB₀–₁, Token Choice)	3,52	—
05	Expert Choice,	3,38	−0,14
06	MoE во всех 4 TS-блоках, E=4, k=2	3,44	−0,08
07	MoE в TSB₀–₁, E=8, k=2	3,47	−0,05
08	MoE в TSB₀–₁, E=16, k=4	3,42	−0,10

Результаты:

Expert Choice — худший вариант (−0,14). Это и есть эмпирическое подтверждение тезиса выше: пропуск токенов в аудио = слышимые артефакты.
MoE везде хуже, чем MoE в двух блоках (−0,08) — балансировочный лосс зашумляет градиент, а глубокие эксперты дублируют друг друга.
Больше экспертов — не лучше. При целевая нагрузка падает до $\frac{K}{E}=0,25$ , балансировка становится жёстче при том же бюджете параметров на эксперта — обучение дестабилизируется. с $\frac{K}{E}=0,5$ оказался самым устойчивым: каждый токен активирует ровно половину пула, и малые дисбалансы не разгоняются так быстро, как при 25%.

Итоговая конфигурация (, , Token Choice, MoE в первых двух блоках) — лучшая по основной метрике среди всех рассмотренных.

Специализация: эксперты реально делят работу, а не дублируют

Главный скептический вопрос к любой MoE: а вдруг прирост — просто от лишних параметров, а роутер на деле гонит все токены через одних и тех же экспертов? Чтобы это проверить, мы смотрели на две диагностические величины.

Энтропия роутера $H=-\mathbb{E}\sum\nolimits_{i}{p_i\log p_i}$ . Максимум для 4 экспертов — $\ln4\approx 1,39$ (равномерное распределение), близко к нулю — вырождение

Нагрузка эксперта — доля токенов, ушедших к нему. Целевое значение $\frac{K}{E}=0,5$ .

Что получилось в обучении:

На ранних шагах нагрузка гуляет от 0,2 до 0,9, но после разогрева (~5–8 тыс. шагов) все четыре кривые сходятся к 0,5 и держатся с отклонением не больше ±0,01.
Энтропия стартует у максимума $\ln4$ и стабилизируется на плато ~1,29 — это ~93% от теоретического максимума. Во всех четырёх MoE-FFN установившиеся значения лежат в узком коридоре 1,28–1,31.

О чём это говорит? Роутер не вырождается в использование пары экспертов — все четыре в деле и нагружены почти поровну. Но при этом энтропия чуть ниже максимума: значит, роутер не кидает токены случайно, а формирует осмысленные предпочтения для конкретных токенов, сохраняя сбалансированную среднюю нагрузку. Иными словами — эксперты действительно специализируются, а не дублируют друг друга. Именно это, а не «просто больше весов», и даёт прирост.

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

После определения лучшей конфигурации модель дообучалась до полного бюджета в 400 эпох.

VoiceBank+DEMAND

^Метод	^Парам.	^{Активн.}	^PESQ↑	^CSIG↑	^CBAK↑	^COVL↑	^SSNR↑
^{Зашумлённое аудио}	^—	^—	^1,97	^3,35	^2,44	^2,63	^1,68
^PrimeK-Net	^1,41M	^1,41M	^3,60	^4,80	^4,01	^4,33	^10,77
^{MP-SENet (base)}	^2,26M	^2,26M	^3,59	^4,79	^4,00	^4,32	^10,71
^{MP-SENet (large)}	^3,59M	^3,59M	^3,61	^4,81	^4,02	^4,34	^10,82
^SESAME	^3,53M	^2,87M	^3,64	^4,84	^4,04	^4,37	^10,93

SESAME берёт лучшие значения по всем пяти метрикам. Прирост по PESQ кажется скромным в абсолютных числах, но для моделей этого класса на этом корпусе пробить планку 3,60 — уже заметно. При этом показательно улучшение композитных метрик и SSNR: остаточный шум давится лучше при сохранении качества самой речи.

Главное же — в столбце «Активн.». SESAME сравнима по общему числу параметров с MP-SENet large (3,53 M против 3,59 M), но благодаря разреженной маршрутизации на каждый токен реально работает лишь 2,87 M — почти на 20% меньше вычислений в соответствующих слоях, при лучшем качестве. Улучшение достигнуто не за счёт роста модели, а за счёт специализации.

Подводим итоги

Разреженная смесь экспертов работает для улучшения речи. На VoiceBank+DEMAND SESAME превосходит базовую MP-SENet, её утяжелённую версию и PrimeK-Net, достигая PESQ 3,64 — при этом на 20% меньше активных параметров, чем у плотной модели сопоставимого размера
Роутинг для текста ≠ роутинг для аудио. Expert Choice, дающий идеальный баланс в LLM, здесь проигрывает из-за пропуска TF-токенов — а каждый токен в аудио физичен, и его пропуск слышен. Token Choice с гарантией обработки каждого токена экспертами — единственный рабочий вариант
Конфигурация не «чем больше, тем лучше». Больше экспертов (, ) и MoE во всех блоках ухудшают результат. , , MoE в первых двух TS-блоках — устойчивый оптимум
Эксперты реально специализируются. Диагностика по энтропии роутера (~1,29 из 1,39) и нагрузке (стабильно у 0,5) подтверждает: все четыре эксперта в работе, равномерно нагружены, но при этом формируют осмысленные предпочтения. Прирост — от разделения труда, а не от лишних весов
Умные потери на этом датасете не помогли. SSL-лосс на HuBERT, multi-scale дискриминатор и VAD-маскирование в условиях эксперимента только снижали качество

Что дальше: хочется глубже разобраться со специализацией — обучить на более крупных и разнородных корпусах (DNS во весь рост, мультиязычная речь) и попытаться интерпретируемо привязать отдельных экспертов к конкретным акустическим классам шума. Анализ текущей модели показал, что выбор агентов в ней не интерпретируем для человека.

Меньше шума, больше смысла: опыт SESAME для Speech Enhancement