Дискретные токены по-прежнему тяжело моделировать и реконструировать: подробнее об этом мы говорили в посте о фреймворке DiTAR. А сегодня разберём статью об очень похожем решении — DiSTAR, фреймворке для zero-shot text-to-speech.
Авторы вдохновлялись LLaDA, где используют диффузионный подход. Но вместо обычных одноуровневых текстовых токенов используют RVQ (residual vector quantization):
- Их битрейта достаточно для качественной реконструкции.
- Из-за дискретности их обучение так же стабильно и интерпретируемо, как и классических LLM.
Авторы утверждают, что известные подходы к моделированию RVQ жертвуют либо скоростью инференса, либо долгосрочной консистентностью. Чтобы решить эту проблему, авторегрессионную языковую модель с маскированной диффузией учат предсказывать токены в дискретном RVQ-пространстве:
1. Формируют из RVQ-токенов патчи, как в DiTAR, и подают их на вход агрегатору — трансформерному энкодеру, который сжимает каждый патч до одного вектора.
2. Авторегрессионный каузальный трансформер учится моделировать непрерывные представления — на каждом шаге обрабатывает сжатый вектор из агрегатора с учётом предыдущего контекста. Новый вектор подаёт в диффузию в качестве контекста.
3. Маскированная диффузионная модель предсказывает следующий патч RVQ-токенов по скользящему окну предшествующих токенов и выходу трансформера.
RVQ-токены порождают гораздо меньше ошибок реконструкции, чем одноуровневый VQ. Моделировать их можно разными способами. В поисках трейд-оффа между качеством работы модели и скоростью её инференса авторы нашли работоспособный вариант с адекватным компьютом: связка авторегрессионного трансформера с маскированной диффузией позволяет совместно моделировать временные и послойные зависимости RVQ.
Александр Плахин