CoSTA сжимает речевые токены Step-Audio 2 mini через BPE: 10K–320K добавленных единиц

Речевые диалоговые модели генерируют ответ в виде длинных последовательностей дискретных речевых токенов, и каждый токен стоит одного шага авторегрессии. Исследователи из Института автоматизации Китайской академии наук и Школы искусственного интеллекта Университета CAS предложили CoSTA — способ сократить число таких шагов, не трогая ни исходный речевой токенизатор, ни декодер, который превращает токены в звук. Для этого в словарь Step-Audio 2 mini добавляют от 10K до 320K BPE-единиц: каждая заменяет часто встречающуюся пару или цепочку соседних речевых токенов. Модель учат выдавать эти сжатые единицы, а перед синтезом их детерминированно разворачивают обратно в исходные токены.

Почему речевые токены хорошо сжимаются

В большинстве речевых диалоговых моделей одно и то же пространство токенов служит и для генерации, и для синтеза. Из-за этого последовательности получаются длинными: ответы идут медленнее, усложняется выравнивание текста и речи, а лимит генерации расходуется быстрее. Это ограничивает длительность разговора и произносимого ответа.

Альтернатива — токенизаторы с меньшей частотой токенов. Но слишком агрессивное снижение частоты ухудшает качество, а замена токенизатора в готовой модели может потребовать переобучить токенизатор, парный декодер и саму модель. К тому же генерация и синтез по-прежнему остаются привязаны к одному представлению.

Авторы заметили, что речевые токены содержат много повторяющихся локальных шаблонов, поэтому к ним применим BPE (byte pair encoding): частые соседние пары объединяются в новые единицы. Сжатие обратимо: развёртка возвращает исходную последовательность без потерь. Эффект проверили на восьми конфигурациях токенизаторов на китайской диалоговой речи: словарь BPE учили на обучающей части и фиксировали для отложенной выборки. Чем больше добавлено единиц, тем короче последовательность, причём зависимость примерно лог-линейна. Она сохраняется на разных источниках данных и языках.

Как адаптируют готовую модель

CoSTA состоит из трёх шагов постобучения. Сначала словарь расширяют, а эмбеддинги и строки LM-головы новых единиц инициализируют средним по строкам исходных токенов, из которых единица состоит. Это даёт более устойчивую стартовую точку, чем случайная инициализация, а параметры исходного словаря остаются неизменными.

Затем идёт адаптация произношения на сжатых TTS-данных. Модель Step-Audio 2 mini чередует текстовые и речевые токены, поэтому BPE применяют только к речевым, а текстовые не трогают. Обучение поэтапное: на первом этапе на 50% данных настраивают только эмбеддинги и LM-голову, остальные параметры заморожены, и это калибрует новые строки словаря. На втором этапе адаптация расширяется на остальную модель. Последний шаг — адаптация к диалогу на сжатых данных формата «речь — речь».

На инференсе перед неизменным речевым декодером вставляется лишь лёгкая детерминированная развёртка BPE. Исходные токенизатор и декодер остаются прежними, поэтому совместимость с существующим конвейером сохраняется.

Что показали эксперименты

Метод проверили на двух архитектурах речевых диалоговых моделей. На Step-Audio 2 mini CoSTA повышает скорость генерации расширенных речевых токенов на китайском и английском, сохраняя качество ответов по автоматическим и человеческим оценкам. Это приводит к меньшей задержке полного ответа, более высокой пропускной способности и более длительному взаимодействию при ограниченном бюджете генерации. Точные коэффициенты ускорения в доступном фрагменте текста не приведены.

Размер словаря задаёт практический компромисс между степенью сжатия, эффективностью декодирования и качеством, так что его нужно подбирать. Эксперименты на другой архитектуре показали перенос подхода и его совместимость с низкочастотной токенизацией: два способа сокращения последовательностей можно сочетать.

Код опубликован на GitHub в репозитории yingzhuodeng/CoSTA.