Кодирование двустадийное:
1️⃣ Сначала эмбеды прогоняют через замороженную GPT-2
2️⃣ Затем переводят в целевое латентное пространство через обучаемый энкодер (в котором приятнее диффундировать)
Декодирование тоже двустадийное:
1️⃣ Прогоняем через обучаемый латентный декодер
2️⃣ Выход декодера подаем уже в обучаемый токен декодер, чтобы получить токены
При обучении вход декодера замумляется, но система энкодер + декодер пытаеся реконструировать исходный эмбед (GPT2). Суть данной аугментации в том, чтобы сделать декодер более устойчивым к шуму, а также предотвратить коллапс латентного пространства.
Диффузионная модель просто оптимизирует L2 лосс для
x0 оценки. Она может обусловливаться на предыдущую оценку x0. Все учится end-to-end с диффузионным лоссом, лоссами реконструкции и кросс-энтропией на предсказание следующего токена (лосс влияет только на token decoder, после него stopgrad).Для того чтобы все хорошо училось важно
🔥 Разогреть энкодер, сначала обучая его только на реконструкцию, а лишь затем включить диффузию.
🕔 Адаптивное расписание сэмплирования шагов времени. Оно подбирается так, чтобы в каждом бине был примерно одинаковый лосс в сумме.
❗️ Различие с ELF
ELF тоже оперирует в латентном пространстве, но фиксированном (Т5 эмбеды).
🧪 Эксперименты
В качестве базового энкодера используется GPT-2 и DiT как диффузионная тушка. В ablation гоняют на отрезках из OpenWebText длиной 128 токенов.
Показывают, что все компоненты метода - разогрев энкодера, обучаемый энкодер поверх GPT-2, расписание шагов диффузии и правильная величина шума важны для качества.
В финальных экспериментах увеличивают длину контекста до 1024 и сравниваются с разными современными подходами по текстовой диффузии
👉 MDLM
👉 Duo
👉 CANDI
👉 FLM
И оказываются лучше по соотношению между перплексией и энтропией (разнообразием). Учится оно несколько дольше большинства бейзлайнов, но зато сэмплирует быстрее на инференсе (за счет того, что декодировать надо только в конце?).
💡 Выводы
Интересный результат. Но как всегда в этой отрасли, встает вопрос о масштабируемости на уровень более-менее SOTA моделек.