기술 요약: SMOOTHIE (토큰 임베딩에 대한 확산 평활화)

문제 제기

확산 모델은 이미지 및 오디오와 같은 연속 영역에서 최첨단 결과를 달성했으나, 언어의 이산적 특성으로 인해 텍스트 생성에 적용될 때 상당한 어려움에 직면합니다. 기존 접근법은 일반적으로 두 가지 범주로 나뉘며, 각각 고유한 한계를 가지고 있습니다:

1. 잠재 공간의 가우시안 확산: 이러한 방법은 연속 토큰 임베딩에 가우시안 노이즈를 적용합니다. 의미 구조를 보존한다는 점 (유클리드 의미 공간 가설에 기반) 에서 유리하지만, 이산적 디코딩 단계에서 어려움을 겪으며, 종종 연속 벡터를 토큰으로 명시적으로 매핑해야 하는데 이는 불안정하거나 정보 손실이 발생할 수 있습니다.
2. 카테고리/심플렉스 확산: 이러한 방법은 어휘 확률 심플렉스나 카테고리 분포에서 직접 작동합니다. 텍스트의 이산적 특성을 존중하고 명시적 디코딩의 필요성을 제거하지만, 노이즈 추가 과정에서 토큰 간의 의미 관계를 일반적으로 무시합니다. 이로 인해 정보의 열화가 더 불안정하고 덜 의미 있게 발생합니다.

핵심 문제는 텍스트의 이산적 특성을 존중하면서도, 순방향 과정에서 의미적 일관성을 유지하는 방식으로 점진적으로 의미 정보를 제거하는 확산 프레임워크가 부재하다는 점입니다.

방법론: SMOOTHIE

저자들은 유클리드 거리를 기반으로 토큰 임베딩을 평활화함으로써 연속 및 이산 확산 간의 간극을 메우는 새로운 프레임워크인 SMOOTHIE(토큰 임베딩에 대한 확산 평활화) 를 제안합니다.

1. 순방향 확산 과정 (노이즈 추가)

SMOOTHIE 는 임베딩이나 카테고리 확률에 직접 노이즈를 추가하는 대신, 대상 시퀀스의 각 토큰 $w_i$ 를 해당 토큰의 임베딩과 어휘 내 모든 토큰 임베딩 간의 음의 제곱 유클리드 거리 벡터로 표현합니다:
$$D_0 = \left( -\|E_{w_i} - E_j\|^2 \right)_{i,j=1}^{m,V}$$
여기서 $E$ 는 고정된 임베딩 행렬 (예: BERT 에서 가져옴) 이고, $m$ 은 시퀀스 길이, $V$ 는 어휘 크기입니다.

순방향 과정은 가우시안 노이즈를 추가하여 이러한 거리 벡터를 점진적으로 교란시킵니다:
$$q(D_t | D_0) = \mathcal{N}\left(D_t \mid \frac{1}{\sigma_t^2}D_0, \delta^2 I\right)$$
여기서 $\sigma_t$ 는 대역폭을 제어하는 노이즈 스케줄러이며, $\delta$ 는 확률적 성질을 제어합니다. 노이즈가 추가된 벡터 $D_t$ 는 소프트맥스 함수를 사용하여 어휘에 대한 확률 분포로 변환됩니다: $p_t = \text{softmax}(D_t)$.

주요 통찰: 노이즈가 증가함에 따라 (즉, $\sigma_t$ 가 커짐), 초기에 올바른 토큰에 집중되어 있던 확률 질량이 다른 토큰으로 점차 분배됩니다. 중요한 점은 입력이 유클리드 거리에 기반하기 때문에, 확률 질량이 먼 토큰에 도달하기 전에 먼저 의미적으로 유사한 토큰(유클리드 거리가 더 작은 토큰) 으로 퍼진다는 것입니다. 이는 "열 방출" 원리를 모방하여, 의미 정보가 무질서하게 제거되는 것이 아니라 점진적으로 제거되도록 보장합니다.

2. 역방향 확산 과정 (디노이징)

역방향 과정은 신경망 $f_\theta$ 가 노이즈가 추가된 입력 $p_t$ 에서 원래 거리 벡터 $D_0$(또는 스케일링된 버전) 를 예측하는 마르코프 체인으로 모델링됩니다. 모델은 예측된 거리 벡터와 실제 거리 벡터 간의 평균 제곱 오차 (MSE) 를 최소화하도록 훈련됩니다:
$$L(\theta) = \mathbb{E}[\|D_0 - g_\theta(p_t, t)\|^2]$$
저자들은 거리 벡터를 예측하는 것이 기대 임베딩을 예측하는 것과 동등함을 증명했으며, 이를 통해 모델은 클린 임베딩 추정을 출력한 후 최근접 이웃 탐색 또는 argmax 를 통해 다시 토큰으로 변환할 수 있습니다.

3. 샘플링 및 솔버

저자들은 표준 DDPM 솔버와 달리 이전 잠재 상태를 의존하는 대신 모델의 예측에 순방향 과정을 직접 적용하여 잠재 상태를 업데이트하는 특정 샘플링 방식 (식 9) 을 도입했습니다.
$$D_{t-1} = \frac{1}{\sigma_{t-1}^2}D_0(f_\theta(p_t, t)) + \tilde{\delta}\epsilon$$
이 접근법은 업데이트 단계에서 이전 숨겨진 상태를 무시합니다. 저자들은 텍스트 임베딩의 희소 데이터 분포에 있어 이것이 유익하다고 주장하며, 잠재 표현이 데이터 매니폴드에서 벗어나는 위험을 줄인다고 설명합니다.

4. 아키텍처 및 조건부

• 아키텍처: UNet 스타일의 스킵 연결이 있는 트랜스포머 디코더 레이어를 기반으로 합니다.
• 조건부: 시퀀스 간 (sequence-to-sequence) 작업의 경우, 추가적인 트랜스포머 인코더가 소스 시퀀스를 처리하며 크로스 어텐션을 통해 통합됩니다.
• 자기 조건부: 모델은 안정성과 품질을 향상시키기 위해 자기 조건부 (self-conditioning) 를 활용하여 이전 예측을 입력으로 다음 단계에 제공합니다.

주요 기여

1. 새로운 확산 프레임워크: SMOOTHIE 는 유클리드 거리를 기반으로 토큰 임베딩을 평활화함으로써 텍스트의 이산적 특성을 존중하면서 동시에 의미 정보를 점진적으로 제거하는 최초의 텍스트 확산 방법입니다.
2. 심플렉스 확산의 일반화: 저자들은 이론적으로 거리 기반 잠재 공간이 심플렉스 기반 잠재 공간 (SSD-LM 및 TESS 에서 사용됨) 을 일반화함을 증명합니다. 심플렉스 확산은 거리 척도가 자명한 경우 (동일 토큰은 0, 그 외는 1) 의 SMOOTHIE 특수 사례임이 입증되었습니다.
3. 개선된 샘플링 솔버: 이 논문은 업데이트 시 이전 숨겨진 상태를 무시하는 솔버를 제안하며, 이는 시퀀스 간 작업에서 표준 DDPM 및 DDIM 솔버보다 경험적으로 우수한 성능을 보입니다.
4. 경험적 우위: 광범위한 실험을 통해 SMOOTHIE 가 기존 확산 기반 모델 (임베딩 기반 및 심플렉스 기반 변형 포함) 을 능가하며, 여러 작업에서 BART 및 GPT-2 와 같은 자기회귀 모델과 비교 가능한 품질을 달성함을 보여줍니다.

실험 결과

저자들은 SMOOTHIE 를 무조건적 생성 작업 (ROCStories) 하나와 시퀀스 간 작업 (Quora Question Pairs, Quasar-T, XSum, ParaDetox) 네 가지에서 평가했습니다.

• 성능: SMOOTHIE 는 다른 확산 기반 베이스라인 (DiffuSeq, TESS, SSD-LM, TEncDM 등) 과 동일한 프레임워크 내에서 훈련된 대체 잠재 공간 구현 (임베딩 기반 및 심플렉스 기반 확산) 보다 일관되게 우수한 성능을 보였습니다.
• 지표:
  • XSum: BERTScore 68.8 및 ROUGE-L 26.0 을 달성하여 다른 확산 방법을 능가하고 자기회귀 베이스라인에 근접했습니다.
  • ROCStories: 확산 모델 중 가장 높은 MAUVE 점수 (80.7) 를 달성하여 GPT-2(78.9) 와 거의 비슷하면서도 높은 다양성을 유지했습니다.
  • ParaDetox: 텍스트 독소 제거를 위한 가장 높은 J-Score(51.7) 를 달성했습니다.
• 어블레이션 연구:
  • $\tilde{\delta}$ (확률적 성질): 시퀀스 간 작업에서는 $\tilde{\delta}$ 의 낮은 값 (예: 0.1) 이 퍼플렉시티와 일관성을 개선한 반면, 무조건적 생성에서는 다양성을 유지하기 위해 약간 더 높은 값이 최적이었습니다.
  • 솔버: 제안된 솔버 (식 9) 는 모든 작업에서 DDPM 및 DDIM 보다 일관되게 우수한 성능을 보였습니다.
  • 임베딩: 모델은 사전 훈련된 BERT 임베딩에 의존합니다. 어블레이션 연구에 따르면 GPT-2 임베딩은 유사한 성능을 보였으나, GloVe 임베딩은 더 낮은 품질을 초래하여 임베딩 품질의 중요성을 강조했습니다.

중요성 및 주장

이 논문은 SMOOTHIE 가 의미 보존과 이산 구조 간의 트레이드오프를 해결함으로써 텍스트 확산 분야에서 중요한 진전을 이루었다고 주장합니다. 유클리드 거리를 통해 노이즈 추가 과정에 의미 관계를 직접 통합함으로써, 정보가 점진적이고 의미 있게 제거되는 "자연스러운 디코딩 과정"을 가능하게 합니다.

저자들은 겸손하게도 이 방법이 그래프나 단백질 서열과 같은 고유한 유사성을 가진 다른 카테고리 도메인에 적용 가능하지만, 현재 MSE 손실의 안정성 문제로 인해 엔드투엔드 훈련 대신 사전 훈련된 임베딩에 의존하고 있다고 지적합니다. 또한 매 단계에서 쌍별 거리를 계산하는 계산 병목 현상을 인정하지만, 성능 향상이 비용을 정당화하며, 높은 다양성과 제어 가능성이 필요한 작업에 대해 자기회귀 모델에 대한 견고한 대안을 제공한다고 주장합니다.

이 연구는 잠재 공간의 선택과 특정 확산 솔버가 모델 아키텍처 자체만큼이나 중요하며, SMOOTHIE 가 연속 및 이산 확산 접근법의 강점을 모두 활용하는 통합 프레임워크를 제공함을 시사합니다.