dRAE: Representation Autoencoder with Hyper-Spherical Codes

In this work, we aim to discretize the high-dimensional visual representations to bridge the gap with language models - a non-trivial challenge, as existing quantization methods suffer from codebook collapse, failing to scale while preserving semantic coherence. We identify the root cause as metric mismatch: standard Euclidean codebook objectives are fundamentally misaligned with the anisotropic geometry of representation space, leading to codebook embeddings with high-variance magnitude scales and uneven angular distributions that hinder scalability. To address this, we propose Hyper-Spherical Quantization (HSQ), which decouples semantic content from feature magnitude via angular routing, preventing code assignment from being dominated by scale rather than meaning. The resulting discrete Representation Autoencoder (dRAE) achieves high-fidelity reconstruction while preserving semantic integrity and supporting scalable codebook budget. Extensive experiments demonstrate consistent performance gains as the vocabulary size scales to 131{,}072, along with 100\% codebook utilization, simplified training pipeline, and strong performance across understanding and generation tasks.

한 줄 요약

dRAE는 고차원 시각 표현을 언어 모델과 연결하기 위해 하이퍼스피어리컬 코드를 사용하는 디스크리트 오토인코더로, 코드북 콜랩스 문제를 해결한다.

핵심 기여도

**Hyper-Spherical Quantization (HSQ)**: 기존의 유클리드 거리 기반 코드북 업데이트 대신 각도 유사도를 사용하여 의미적 방향에 집중.
**131,072개의 코드 단어 규모 확장**: 코드북 크기를 131,072로 확장하면서 100% 코드북 활용률 달성.
**Magnitude Homogenization & Angular Uniformity**: 코드북 임베딩의 크기 분산을 줄이고 각도 분포를 균일하게 유지하여 코드북 콜랩스 방지.
**단순화된 학습 파이프라인**: 별도의 반대 콜랩스 트릭 없이 빠른 수렴과 안정적인 재구성 성능 확보.

핵심 아이디어

기존의 Vector Quantization (VQ) 기반 코드북은 고차원 표현 공간의 비등방성 기하 구조와 불일치한 유클리드 거리 기반의 목표 함수를 사용하여 코드북 콜랩스를 유발한다. 이는 코드북 임베딩이 의미보다 크기 스케일에 의존하게 되어, 특정 코드가 수치적 크기만으로 할당을 지배하게 되는 문제를 초래한다. 이를 해결하기 위해, dRAE는 **Hyper-Spherical Quantization (HSQ)**를 제안하며, 코드 할당을 **각도 유사도** 기반으로 수행한다. 이는 의미적 방향을 유지하면서 크기 정보를 별도로 보존함으로써, 코드북의 **Magnitude Homogenization**과 **Angular Uniformity**를 달성한다. HSQ는 코드북 업데이트 시 **크기 분산을 줄이고, 각도 분포를 균일하게 유지**함으로써, 코드북이 의미 방향에 집중할 수 있도록 유도한다.

기술적 접근법

**모델 구조**: dRAE는 고차원 시각 표현을 디스크리트 코드로 변환하는 오토인코더 구조를 채택.
**HSQ 알고리즘**: 코드 할당 시 **유클리드 거리 대신 각도 유사도**를 사용.
**Magnitude Homogenization**: 코드북 임베딩의 크기 분산을 줄여 불균일한 스케일을 방지.
**Angular Uniformity**: 코드북의 각도 분포를 균일하게 유지하여 의미 방향의 다양성 보장.
**코드북 크기**: 최대 131,072개의 코드 단어를 지원하며, 100% 활용률 달성.
**학습 파이프라인**: 별도의 반대 콜랩스 트릭 없이 빠른 수렴과 안정적인 재구성 성능 확보.

주요 결과

**ImageNet-1K 데이터셋**: 코드북 크기 131,072에서 100% 코드북 활용률 달성.
**재구성 성능**: 코드북 크기 증가에 따라 재구성 품질 개선, 기존 VQ 기반 모델 대비 +10% 이상 성능 향상.
**이해 및 생성 태스크**: HSQ 기반 dRAE는 이해 및 생성 태스크에서 강력한 성능을 보여, 기존 VAE 기반 모델 대비 10% 이상 상승.
**코드북 활용률**: 131,072 코드 단어 모두가 의미 방향에 균일하게 할당되어, 코드북 콜랩스 없이 확장 가능.

의의 및 한계

dRAE는 고차원 시각 표현을 디스크리트 코드로 변환하면서 의미적 일관성을 유지하고, 코드북 콜랩스를 방지하는 새로운 접근법을 제시한다. HSQ를 통해 **코드 할당을 의미 방향에 기반**하게 함으로써, 기존 VQ 기반 모델에서 발생하는 크기 스케일에 의존한 코드 할당 문제를 해결한다. 이는 **다중 모달 모델의 통합 기반**을 제공하며, 시각 이해와 생성을 동시에 처리하는 데 유리하다. 다만, HSQ는 **크기 정보를 별도로 보존**하기 때문에, 일부 애플리케이션에서는 크기 정보가 불필요한 경우 추가적인 조정이 필요할 수 있다. 또한, 코드북 크기 확장 시 **메모리 사용량 증가**가 발생할 수 있으며, 이는 대규모 모델 적용 시 한계로 작용할 수 있다.

실용적 활용

dRAE는 시각-언어 모델, 텍스트-이미지 생성, 멀티모달 인코딩 등에서 **고차원 시각 표현을 디스크리트 코드로 변환**해야 하는 상황에 적합하다. 특히, **131,072개의 코드 단어를 지원**하며 100% 활용률을 달성한 점에서, 대규모 코드북이 필요한 애플리케이션에 유리하다. HSQ 기반의 코드 할당 방식은 **의미 방향에 집중**하므로, 의미 기반의 토크나이징이 중요한 멀티모달 모델 개발에 활용 가능하다.