Restructuring Vector Quantization with the Rotation Trick

Christopher Ré продолжает делать интересное, правда на уже не в рамках HazyResearch, и на сей раз он добрался до VQ-VAE

авторы поднимают вопрос того, что по факту Straight-Through Estimator (STE) во время бекворда третирует производную кодбук вектора по отношению к энкодер аутпуту как identity matrix

но в этом может заключаться и загвоздка, когда бекворд доходит до энкодера, ибо по сути STE не привносит свой импакт по изменению параметров энкодера → положение энкодер аутпутов по отношению к кодбукам не влияет на градиент энкодера

авторы предлагают оч красиво и просто такое исправить - через матрицу поворота, а в итоге метод назвали rotation trick. он не меняет аутпуты при форвард пассе, но когда дело доходит до бекворда, то на градиент уже влияет как раз эта матрица поворота, которая высчитывается в соответствии с энкодер аутпутом и кодбук вектором

и вычисляют к тому же они не наивно, а сразу через хаусхолдера, при том, по заверениям авторов, разницы по времени с классическим вк вае нет

сравнивали на видео и картинках + еще добавили адверсариал сетап вместе с диффузией. много где обыгрывает классику по всем трем фронтам - реконструкции, “утилизации” кодбука и ошибке квантизации

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

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

👀LINK