기술 요약: DiHAL – 기하학적 유도 숨김 상태 대체

문제 정의

연속 확산 언어 모델은 역사적으로 자동회귀 트랜스포머보다 뒤처져 왔습니다. 확산 모델은 반복적 노이즈 제거를 통해 이미지 생성에서 탁월한 성능을 발휘하지만, 토큰의 이산적 특성으로 인해 이를 텍스트에 적용하는 것은 어렵습니다. 이산 토큰 손상, 마스크 확산, 또는 토큰 임베딩과 학습된 잠재 공간에서의 연속 확산과 같이 언어에 확산을 적응시키려는 기존 접근법들은 종종 "연속에서 이산으로"의 복구 병목 현상에 직면합니다. 연속 표현 공간에서의 작은 오류는 잘못된 토큰 복구를 초래하여 생성 품질을 저하시킵니다.

저자들은 성능 격차가 이산성 때문만이 아니라 부적합한 기하학을 가진 연속 공간에 확산을 적용하기 때문이라고 가설을 세웁니다. 트랜스포머 언어 모델이 연속 숨김 상태 위에서 작동하지만, 이러한 상태는 계층마다 기하학적 특성이 다양합니다. 제기된 핵심 질문은 다음과 같습니다: 사전 훈련된 트랜스포머 내부에서 확산 기반 노이즈 제거 및 복구에 필요한 기하학을 갖춘 표현 공간은 어디에 존재하는가?

방법론: DiHAL

본 논문은 독립적인 확산 모델을 훈련하는 대신 사전 훈련된 트랜스포머에 확산 브릿지를 통합하는 "위치 파악 및 대체" 프레임워크인 DiHAL(언어 생성용 확산 - 트랜스포머 하이브리드 아키텍처) 을 제안합니다.

1. 이론적 동기

저자들은 접근 방식을 랑베빈 동역학과 집중 이론에 근거합니다. 그들은 "확산에 친화적인" 표현 공간이 세 가지 기하학적 속성을 만족해야 한다고 주장합니다:

• 빠른 수축: 공간은 노이즈가 있는 샘플이 목표 분포로 빠르게 수축할 수 있어야 합니다 (강한 로그 오목성/곡률에 의해 지배됨).
• 안정성: 노이즈 제거 과정은 점수 추정 오류에 대해 견고해야 합니다 (역시 곡률에 의해 지배됨).
• 낮은 유효 복잡도: 표현의 변동은 저차원 매니폴드에 집중되어야 하므로 확산 모델이 학습하기 쉬워야 합니다.

2. 위치 파악: 기하학적 유도 계층 선택

숨김 상태의 실제 밀도와 헤시안 행렬을 사용할 수 없으므로, DiHAL 은 각 트랜스포머 계층을 평가하기 위해 경험적 기하학적 대리 변수를 사용합니다:

• 국소 곡률 ($\hat{m}_{curv}$): 국소 이웃 공분산의 최대 고유값의 역수를 통해 추정됩니다. 높은 값은 컴팩트한 이웃을 나타냅니다.
• 전역 단조성 ($\hat{m}_{mono}$): 정규화된 정밀도 행렬의 레일리 몫을 통해 추정됩니다. 높은 값은 전역 방향적 강성을 나타냅니다.
• 유효 랭크 ($\hat{k}$): 공분산의 트레이스를 그 스펙트럼 노름으로 나눈 값 ($\text{tr}(\Sigma)/\|\Sigma\|$) 으로 추정됩니다. 낮은 값은 저차원 매니폴드에서의 집중을 나타냅니다.

선택 점수는 이러한 대리 변수들을 결합하여 (곡률은 보상하고 유효 랭크는 패널티를 부여) 최적의 삽입 계층 $\ell^*$을 식별하도록 계산됩니다.

3. 대체: 숨김 상태 확산 브릿지

최적의 계층 $\ell^*$이 식별되면:

• 하부 트랜스포머 접두사 (계층 $1$부터$\ell^*$까지) 는 조건부 확산 브릿지로 대체됩니다.
• 브릿지는 소스 모델의 입력에서 파생된 임베딩 기반 조건을给定하여 숨김 상태 $h_{\ell^*}$를 재구성하도록 훈련됩니다.
• 상부 트랜스포머 계층 ($\ell^*+1$부터 $N$까지) 과 원래 언어 모델 (LM) 헤드는 유지되고 고정됩니다 (또는 최소한으로 미세 조정됩니다).
• 추론 시, 브릿지는 조건을 노이즈 제거하여 $\hat{h}_{\ell^*}$를 생성한 후, 이를 유지된 상부 계층에 전달하여 토큰 예측을 수행합니다. 이는 직접적인 토큰 복구를 피하고 문제를 숨김 상태 재구성으로 재구성합니다.

주요 기여

1. 확산 삽입 공식화: 본 논문은 사전 훈련된 트랜스포머에 확산을 삽입하는 것을 기하학적 유도 인터페이스 선택 문제로 공식화하며, 확산에 친화적인 공간을 식별하기 위한 실용적인 계층별 대리 변수 (국소 컴팩트성, 전역 강성, 유효 랭크) 를 제안합니다.
2. 고정 기하학 점수: 저자들은 8B 규모 백본 전반에 걸쳐 브릿지화 가능성 (재구성 품질) 과 강력하게 상관관계를 보이는 고정 기하학 점수를 도입합니다. 이를 통해 모든 후보 계층에 대한 브릿지 훈련을 철저히 수행하지 않고도 효과적인 삽입 계층을 식별할 수 있습니다.
3. DiHAL 아키텍처: 하부 트랜스포머 계층을 조건부 확산 브릿지로 대체하고 상부 계층 및 LM 헤드를 재사용하는 하이브리드 아키텍처입니다. 진단적 매칭 예산 설정 하에서 DiHAL 은 숨김 상태 복구가 임베딩 기반 및 잠재 기반 연속 확산 베이스라인보다 우수함을 입증합니다.

실험 결과

저자들은 Dolma v1.7 말뭉치를 사용하여 Llama-3.1-8B-Instruct 및 Qwen3-8B 백본에서 DiHAL 을 평가했습니다.

• 계층 선택: 기하학 점수가 효과적인 삽입 계층을 성공적으로 예측했습니다. Llama-3.1-8B 의 경우 점수가 3 계층을 선택했고, Qwen3-8B 의 경우 2 계층을 선택했습니다. 이러한 계층들은 임베딩 인터페이스 근처에 위치하여 초기 숨김 상태가 확산에 적합한 임베딩과 유사한 기하학적 구조를 유지함을 시사합니다.
• 브릿지화 가능성과의 상관관계: 고정 기하학 점수는 검증 브릿지 손실과 강한 스피어만 상관관계 ($\rho \approx 0.91$–$0.93$) 를 보여, 기하학적 대리 변수가 확산 브릿지 훈련이 성공적으로 이루어질 수 있는 계층을 저렴하게 식별할 수 있음을 확인했습니다.
• 진단 비교: 베이스라인 (Diffusion-LM, SED, LD4LG, CoDAR) 에 대한 매칭 예산 비교에서 DiHAL 은 가장 낮은 생성 퍼플렉시티 (Gen.PPL) 와 가장 높은 다양성을 달성했습니다. 구체적으로 Qwen3-8B 에서 DiHAL 은 Gen.PPL 을 144.83 (CoDAR) 에서 136.02 로, 다양성을 0.4777 에서 0.5913 으로 개선했습니다.
• 전체 훈련: 완전히 훈련되었을 때, 기하학으로 선택된 계층들은 검증 손실 오라클 (모든 브릿지를 명시적으로 훈련하여 선택된 계층) 과 경쟁력을 유지했으며, "최악의 계층" 제어군 및 CoDAR 베이스라인보다 현저히 우수했습니다.

중요성 및 주장

본 논문은 숨김 상태 기하학이 언어 모델 내에서 확산이 효과적으로 작동할 수 있는 위치를 결정하는 중요한 요소라고 주장합니다. 결과는 언어에서 연속 확산의 실패가 단순히 이산성의 결과뿐만 아니라 노이즈 제거 요구 사항과 표현 공간 기하학 사이의 불일치 때문일 수도 있음을 시사합니다.

주요 교훈은 다음과 같습니다:

• 문제 재구성: 언어를 위한 연속 확산은 토큰을 직접 복구하여 텍스트를 생성할 필요가 없습니다. 대신 사전 훈련된 트랜스포머가 이미 디코딩 방법을 알고 있는 내부 표현을 재구성할 수 있습니다.
• 하이브리드 모델의 실현 가능성: 삽입 지점이 기하학적 기준에 따라 선택된다면, 사전 훈련된 트랜스포머의 일부를 확산 브릿지로 대체하는 것이 실현 가능합니다.
• 한계: 저자들은 겸손하게 DiHAL 이 아직 독립적인 확산 언어 모델은 아니며, 토큰 예측을 위해 유지된 트랜스포머 접미사에 의존한다고 지적합니다. 또한, 노이즈 제거 단계로 인해 추론 지연 시간은 현재 원래 백본보다 높습니다. 이 작업은 이러한 규모에서 자동회귀 생성을 완전히 대체하기보다는, 확산이 기존 모델에 어디서 진입할 수 있는지에 대한 이해를 위한 한 걸음으로 제시됩니다.