Preprint. Under review
https://arxiv.org/abs/2506.21170
Viacheslav Meshchaninov, Egor Chimbulatov, Alexander Shabalin, Aleksandr Abramov, Dmitry Vetrov
Introduction
Autoregressive(AR) 언어 모델은 텍스트 생성 분야에서 지배적으로 자리매김하고 있다. 이는 다음 토큰 예측 과제로서 teacher forcing 방식을 사용하여 최적화하는데, 이러한 순차적 분해 방식은 여러 구조적 bottlenecks을 야기한다.
- 디코딩은 본질적으로 순차적인 방식이기 때문에, 생성할 시퀀스의 길이가 길어질수록 생성 시간도 선형적으로 증가한다.
- exposure bias 로 인해 초기의 작은 오류가 전체 문맥에 영향을 미치며, 이는 중간에 수정이 불가능하다.
- Local log-likelihood를 최대화하려는 학습 목표는 모델이 사실성보다는 유창성을 우선시하게 만들며, 이로 인해 hallucination 현상이 발생한다.
- 토큰 단위로 결정을 내리기 때문에, 모델은 전역적인 계획을 갖기 어려우며, 이로 인해 장기적인 논리나 서사적 일관성을 유지하는 데 한계를 보인다.
대조적으로, CV 분야에서는 diffusion model의 도입으로 생성 패러다임이 크게 변화하였다. 특히 latent diffusion의 등장은 이미지나 비디오를 먼저 latent representation으로 압축한 후, 이 latent space에서 디퓨전 과정을 수행하는 방식으로 전환점을 마련했다. low-dimensional space에서 연산을 수행함으로써 계산 비용을 줄이면서도 이미지나 비디오 합성 분야에서 큰 성공을 거두었다. 그러나 최근 연구들은 디퓨전이 latent space의 geometry에 매우 민감하다는 사실을 밝혀냈다. 즉, latent representation 설계에 따라 학습 안정성이 떨어지고 생성 품질을 저하시킬 수 있으며, 이는 diffusable representation을 구성하기 위한 원칙 기반의 메커니즘 설계 필요성을 시사한다.
저자들은 latent diffusion의 성공에 착안하여, 기존의 토큰 수준 표현이 제공하는 생성 품질을 유지하거나 능가하면서, 텍스트 정보를 얼마나 압축된 잠재 공간으로 줄일 수 있을지에 대한 질문을 제기한다. 저자들은 텍스트를 저차원 공간으로 매핑하는 Autoencoder를 설계하고, 그 공간에서 디퓨전 모델을 학습한다. 단순한 토큰 복원 방식은 불안정한 잠재 구조를 초래하지만, 견고성과 매끄러움을 고려한 학습 목표는 안정된 manifold를 형성해 고품질 생성을 가능케 한다.
이에 따라 COSMOS(COmpressed and SMOoth latent Space)라는 새로운 training recipe를 통해 압축되고 매끄러운 잠재 공간에서 디퓨전 모델을 학습할 수 있음을 보였으며, 이는 기존 토큰 수준 디퓨전과 유사하거나 그 이상의 품질을 달성한다.
Methodology
Frozen text encoder. 저자들은 pretrained BERT $E_{\text{text}}$를 텍스트 인코더를 모든 학습 단계에서 frozen 된 상태로 유지된다. 길이가 $L$인 토큰 시퀀스 $\mathbf{w} = (w_1, \ldots, w_L)$에 대해 인코더는 $\mathbf{h} = E_{\text{text}}(\mathbf{w}) \in \mathbb{R}^{L \times d}$를 출력하여 이때 $d=768$이다.
Compressor. 가변길이의 hidden state $\mathbf{h}$를 고정된 크기의 latent로 압축하기 위해, 본 연구에서는 Percevier Resampler architecture를 사용한다. 이는 12개의 레이어로 구성된 transformer이며, 각 블록의 self-attention을 cross-attention으로 대체하였다. 구체적으로 학습 가능한 벡터 $\mathbf{u} \in \mathbb{R}^{N \times d}$를 random initialized 하며, 이때 $N \ll L$ 이다. 각 블록에서 이 벡터 $\mathbf{u}$는 query로 사용되며 전체 시퀀스의 정보를 집약하면서도 의미적으로 구조화된 잠재 표현을 생성한다. 추가적으로 저자들은 sequence axis에서만 압축을 수행한다.
Latent normalization. 가우시안 디퓨전 과정 시작 전 held-out corpus에서 global 평균과 표준편차 $(\mu, \sigma) \in \mathbb{R}^{N \times d}$ 를 추정하고 각 latent feature를 normalization을 수행한다. 구체적으로 $\mathbf{z} \leftarrow \frac{\mathbf{z} - \mu}{\sigma}$로 정규화하며 분산 보전 디퓨전 과정을 더 잘 수행해 줄 수 있다.
Latent diffusion model. 디퓨전 모델은 다음 식을 통해 최적화되며, $N$이 작기 때문에 디퓨전 모델은 더 빠르게 실행된다.
$$z_t = \sqrt{\alpha_t} z_0 + \sqrt{1 - \alpha_t} \, \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I)$$
$$\mathcal{L}_{\text{DM}} = \mathbb{E}_{z_0 \sim p_{\text{data}}, \, t \sim \mathcal{U}[0,1], \, \varepsilon \sim \mathcal{N}(0, I)} \left[ \left\| z_0 - z_\theta(z_t, t) \right\|_2^2 \right]$$
Decompressor. 구조적으로 Compressor와 대칭구조이며 고정 길이의 latent representation의 길이 $L$을 다시 확장하여 $\hat{\mathbf{h}} \in \mathbb{R}^{L \times d}$를 생성하며 이때 $L_{max}=512$로 제한된다.
Token predictor. 생성된 $\hat{\mathbf{h}}$는 linear projection과 softmax를 거쳐 각 벡터를 어휘 집합 상의 확률 분포로 변환하고 이를 통해 최종 텍스트를 생성한다.
Learning a compact text latent space
저자들은 먼저 latent sequence length $N$에 따른 토큰 수준 복원 정확도를 측정하였다. 이때 objective는 cross-entropy로 설정하였으며, 위의 그림과 같이 $N=16$개부터 100%에 달하는 복원 정확도를 달성하였다. 하지만 저자들은 높은 복원 정확도 자체만으로는 latent manifold가 생성 모델링에 적합하지 않다고 주장한다.
왜냐하면 robust한 디퓨전 기반 생성을 위해서는 compressed space가 추가적인 조건들을 만족해야 하기 때문이다. 저자들은 latent manifold가 smoothness와 robustness를 결여할 경우, 가우시안 디퓨전 모델은 고품질 텍스트를 생성할 수 있는 latent vector를 샘플링하지 못한다는 것을 관찰하였다. 이를 해결하기 위해, 저자들은 Fig. 1. 아키텍처와 같이 학습 과정에 세 가지 상호보완적인 전략을 적용하였다.
MSE regularisation on encoder activations. 저자들은 cross-entropy loss에 더해, $h$와 $\hat{h}$ 사이에 MSE loss를 추가로 도입하였다. 이 보조 loss를 통해 문맥 의미를 담고 있는 $h$를 더 잘 보존하도록 하였다.
Activation-space perturbations. compressor가 단순히 semantic 정보를 보존하는 것뿐만 아니라, 추가적인 feature를 추출하도록 학습시키기 위해, 본 연구에서는 perturb-and-recover 학습 절차를 적용하였다. 구체적으로 변형된 view인 $\mathbf{h}'$를 샘플링하여 compressor–decompressor 파이프라인을 통과시킨 후 $\hat{\mathbf{h}}'$를 얻고 $\text{MSE}(\mathbf{h}, \hat{\mathbf{h}}')$를 최소화하도록 학습시킨다. 이때 다음 두 가지 perturbation 방식이 각 미니배치 내에서 동일한 확률로 적용된다:
(a) Random masking: $h$의 vector 중 30%를 0으로 설정함.
(b) Gaussian noise: 사전에 계산된 통계로 정규화한 $h$에 노이즈를 주입함.
$$\mathbf{h}' = \delta \, \mathbf{h} + \sqrt{1 - \delta^2} \, \varepsilon$$
이러한 data augmentation는 오토인코더가 부분적 information loss에 대해 더 robust하게 작동하도록 만들고, 인접한 latent 사이에 smooth interpolation를 유도한다.
Latent-space augmentation. 저자들은 추가적으로 latent vector $z$ 자체에 직접적인 augmentation을 적용한다. 잠재 벡터의 수 $N$은 토큰 수 $L보다 훨씬 작기 때문에, 전체 잠재 벡터를 마스킹하면 과도한 정보 손실이 발생하게 된다. 이에 따라 훈련 중에는 각 latent vector 내 개별 feature에 대해 고정된 비율 $p$만큼 무작위로 0으로 설정함으로써 fine-grained sparsification을 수행한다. 이러한 방식은 각 잠재 벡터 내부의 인접한 feature들이 redundant cue를 담도록 유도하여 일부 feature가 제거되더라도 전체 표현을 해석할 수 있게 만든다. 결과적으로 latent space 상 작은 perturbation에도 모델을 더욱 robust 하게 만든다.
Latent-space properties that facilitate diffusion training
저자들은 text-latent manifold를 탐색하며, 내재적 속성 중 어떤 요소가 디퓨전 모델의 성능을 좌우하는 지를 분석한다. 모든 실험은 128개 토큰 길이 텍스트를 16개로 압축하는 오토인코더를 사용하였다. 저자들은 실험을 통해 manifold의 smoothness과 perturbation에 대한 robustness라는 두 가지 특성이 디퓨전 학습에 있어 핵심적인 역할을 한다는 점을 강조한다.
Smoothness of the latent manifold.
오토인코더 학습 후 각 텍스트 $w$는 latent vector $z$에 대응되며 이로부터 원래의 텍스트를 복원할 수 있다. 하지만 저자들은 이 latent space 내 어떠한 텍스트에도 매핑되지 않는 벡터들이 존재할 수 있으며, 이러한 영역의 데이터 분포가 불명확할 경우 디퓨전 모델이 잘못된 분포를 학습하거나 일반화에 실패할 수 있음을 지적한다. 왜냐하면 디퓨전 모델은 텍스트의 분포를 잠재 벡터의 분포를 통해 간접적으로 학습하게 되는데, 이 분포는 오직 학습 중 관찰된 제한된 영역만을 기반으로 수행하기 때문이다. 따라서 latent space에서의 데이터 density가 얼마나 smooth하게 변화하는지에 따라 보지 못한 영역까지 일반화가 가능한지를 결정짓는다. 이를 분석하기 위해 저자들은 일련의 실험을 진행하였다.
실험 절차
① trainig corpus에서 두개의 문장을 랜덤 하게 선택하고 각각 latent vector $\mathbf{z}^{(1)}, \, \mathbf{z}^{(2)} \in \mathbb{R}^{N \times d}$ 를 인코딩함
② linear interpolation을 통해 중간점을 구성함으로써 보지 못한 영역으로 벡터를 이동시킴
$$\mathbf{z}^{\mu} = \mu \, \mathbf{z}^{(1)} + (1 - \mu) \, \mathbf{z}^{(2)}, \quad \mu \in [0, 1]$$
③ $ \mathbf{z}^{\mu}$에 디퓨전 노이즈를 적용하여 $\mathbf{z}^{\mu}_t$를 생성하고, time step $t$에 따라 노이즈 수준이 다른 경우의 영향을 관찰함
④ $\mathbf{z}^{\mu}_t$로 부터 $\hat{\mathbf{w}}$를 예측하고 원래의 텍스트로 디코딩함
⑤ GPT-2 PPL을 측정하여 생성된 텍스트의 자연스러움을 평가함
실험 결과 COSMOS는 전체 latent space에 걸쳐 smooth하고 generalizable 표현을 형성하는 반면, CE baseline은 학습 중 관찰된 지점에서는 잘 작동하지만, interpolation region에서는 급격히 PPL이 증가하며 품질이 저하되는 현상이 확인되었다. 이는 COSMOS가 더 smooth latent manifold를 형성하여 낯선 영역에서도 디퓨전 모델이 안정적으로 작동할 수 있게 함을 보여주는 근거이다.
Reducing the train–inference mismatch.
저자들은 text latent space에서 디퓨전 모델을 학습할 때 train-inference mismatch 발생하는 경향이 있음을 실험적으로 관찰하였다. 즉 샘플링 시점에서 생성된 latent vector는 그것을 디코딩한 후 이를 다시 인코딩했을 때와 상당히 다를 수 있다.
$$\mathbb{E}(D(\hat{\mathbf{z}})) \neq \hat{\mathbf{z}}$$
이러한 불일치는 두 가지 문제를 야기하는데, ①디코더를 더 이상 신뢰할 수 없는 문제와 ②디퓨전 모델 반복적으로 학습 중 한 번도 본 적 없는 벡터들을 입력으로 사용하게 되어, 시간이 지남에 따라 오류가 누적될 수 있다.
이에 따라 저자들은 Decoder와 Diffusion model이 모두 latent space 내 perturbation에 대해 robut해야 한다고 주장한다. 이를 검증하기 위해, 본 저자들은 앞서 제안한 학습 전략이 train-inference mismatch를 실질적으로 줄일 수 있는지를 평가하는 실험을 진행하였다.
Decoder robustness. 디코더가 perturbation에 얼마나 robust 한 지 평가하기 위해 실제 텍스트의 latent vector $z$에 가우시안 노이즈를 주입한다: $\mathbf{z}_{\text{noised}} = \mathbf{z} + \sigma \, \varepsilon, \quad \text{where } \varepsilon \sim \mathcal{N}(0, I)$ 이 후 이를 디코딩한 뒤 원래 텍스트와 BLEU를 측정하였다.
실험 결과, 디코더의 robust에는 두 가지 요인이 가장 크게 기여함을 확인하였다.
① Decompressor의 MSE loss: 이는 마지막 레이어에서 자주 발생하는 norm exploding 문제를 방지하여, 안정적인 복원을 가능하게 한다.
② latent masking: compressor가 정보를 latent feature 전반에 고르게 분산시키도록 유도하고, 일부 feature가 손실되더라도 복원 가능성을 유지할 수 있게 한다.
이러한 학습 전략은 디코더가 latent space 내 purterbation에 대해 매우 robust 하다는 것을 시사한다.
Diffusion robustness during generation. 저자들은 생성 과정 중 디퓨전 모델이 작은 perturbation에 얼마나 민감한지를 평가하기 위해, reverse process의 후반부 단계에서 노이즈를 주입한 뒤 샘플링을 계속 진행하는 실험을 수행하였다. 노이즈가 주입되지 않은 경우와의 차이는 최종 출력 간의 MSE로 측정된다: $\left\| \hat{\mathbf{z}} - \hat{\mathbf{z}}_{\text{shifted}} \right\|_2^2$
실험 결과 제안한 latent space에서 학습된 디퓨전 모델은 기존 CE baseline보다 훨씬 더 안정적인 반응을 보였으며, 노이즈가 주입된 경우에도 생성 경로의 일관성을 잘 유지하는 것으로 나타났다.
Direct mismatch measurement. 저자들은 샘플링된 latent vector \hat{\mathbf{z}}와 이를 다시 디코딩 후 다시 인코딩한 결과 $\mathbb{E}(D(\hat{\mathbf{z}}))$ 사이의 MSE를 통해 train-inference mismatch를 직접 측정하였다.
실험 결과 제안된 학습 기법을 하나씩 적용할수록 이 오차가 점진적으로 감소하며, 학습 절차가 오토인코더를 디퓨전 노이즈에 더 강건하게 만든다는 것을 보여준다. 결과적으로, 저자들의 학습 전략은 디코더와 디퓨전 모델 모두의 안정성을 향상하며, train–inference 간의 격차를 효과적으로 줄였다.
Results
각 기법의 성능 기여를 분석한 결과, 제안된 모든 구성 요소들이 디퓨전 모델의 성능 향상에 기여함을 확인하였다. 이에 따라, 이후 모든 실험에서는 해당 결과를 기반으로 선정된 hyper-parameter 조합을 최적 설정으로 고정하여 사용한다.
latent vector의 개수 N에 따른 성능 변화를 분석한 결과, 그림과 표로 미뤄봤을 때 $N=32$(4배 압축) 설정에서 성능 저하 없이 높은 품질을 유지할 수 있음을 밝혔다.
다양한 생성 모델과의 비교 실험 결과, COSMOS는 기존 latent diffusion은 물론 autoregressive 모델 대비 대부분의 평가 지표에서 우수한 성능을 보였고, 특히 긴 문맥 과제에서 빠른 생성 속도를 유지하는 것이 특징이다.
