기술적 요약: 마스킹 확산 모델이 반전 저주를 완화하는 이유에 대한 이론적 분석

문제 제기

자기회귀 언어 모델 (ARM) 은 **반전 저주 (reversal curse)**에 시달립니다. 즉, 한 순서로 관계적 사실을 학습한 후 (예: "A 는 B 이다"), 해당 순서가 반전된 질의 (예: "B 는 A 이다") 에 대해 자주 답변하지 못합니다. 최근 마스킹 확산 모델 (MDM) 은 반전 저주가 현저히 약하게 나타나는 모습을 보이지만, 그 근본적인 메커니즘은 여전히 불명확했습니다.

일반적인 설명은 MDM 의 **임의 순서 마스킹 학습 목표 (any-order masked training objective)**가 이를 완화한다고 보는데, 이는 모델이 순방향과 역방향 모두에서 마스킹된 예시를 볼 수 있게 허용하기 때문입니다. 그러나 저자들은 이 설명이 불완전하다고 주장합니다. 논문의 Figure 1 에서 보듯, "[M] 은 B 이다"와 같은 순방향 마스킹 예시 (목표는 A) 를 관찰하는 것은 특정 위치 구성 ($i < j$) 에서 조건부 $p_\theta(A|B)$에 대한 감독 신호를 제공합니다. 이는 왜 학습된 증거가 $j' < i'$인 다른 위치 구성에 해당하는 역방향 질의 "B 는 [M] 이다"로 전이되어야 하는지 본질적으로 설명하지 못합니다.

방법론

이 논문은 역방향 전이를 가능하게 하는 메커니즘을 분리해 내기 위해 단순화된 단층 (one-layer) 마스킹 확산 모델에 대한 이론적 분석을 제공합니다. 방법론의 핵심은 **저장 - 라우팅 분해 (Storage-Routing Decomposition)**입니다:

1. 저장 - 라우팅 분해: 저자들은 입력 토큰 $B$가 목표 토큰 $A$의 로짓 (logit) 에 기여하는 바를 두 가지 구성 요소로 분해합니다:

   • 공유 저장 (Shared Storage, $Y_{A,B}$): 모델 파라미터에 저장된 증거를 나타내는 위치 불변 (position-invariant) 항입니다. 상대적 위치 인코딩 (RoPE 또는 ALiBi 등) 하에서 토큰 $B$의 값 벡터는 그 위치에 독립적이므로, $Y_{A,B}$는 $B$가 마스킹 앞이나 뒤에 나타나든 상관없이 공유됩니다.
   • 어텐션 라우팅 (Attention Routing, $\alpha_{i,j}$): 마스킹된 위치가 토큰 $B$를 얼마나 강하게 주시하는지를 결정하는 위치 의존적 어텐션 가중치입니다.

2. 이론적 분석:

   • 저장 역학 (Storage Dynamics): 저자들은 순방향 마스킹 학습 (B 로부터 A 복원) 이 공유 저장 항 $Y_{A,B}$를 증가시켜, 역방향 질의에서 재사용 가능한 증거가 되게 함을 증명합니다.
   • 라우팅 상관관계 (Routing Correlation): 역방향 질의가 이 저장된 증거에 여전히 접근할 수 있는지 분석합니다. RoPE(회전 위치 임베딩) 와 ALiBi를 사용하여, 순방향 및 역방향 구성에 대한 어텐션 점수가 **양적으로 상관관계 (positively correlated)**를 유지함을 증명합니다. 구체적으로 RoPE 의 경우, 기대 상관관계는 회전 행렬의 트레이스에 의존하는 양의 값 하한으로 제한되며, ALiBi 의 경우 콘텐츠 항이 동일하므로 상관관계는 완벽합니다 (1.0).
   • 그래디언트 정렬 (Gradient Alignment): 순방향 및 역방향 손실의 공유 저장 열 $Y_{:,B}$에 대한 그래디언트가 양적으로 정렬됨을 보여줍니다. 결과적으로 순방향 손실을 최소화하기 위해 취해진 그래디언트 단계는 1 차적으로 역방향 손실을 감소시킵니다.

3. 실험적 검증:

   • 제어된 단층 실험: 소문자 - 대문자 쌍으로 구성된 합성 데이터셋을 사용하여 단층 MDM 과 ARM 을 학습시켰습니다. 어텐션 점수 상관관계, 그래디언트 코사인 유사도, 그리고 역방향 예측 확률을 측정했습니다.
   • 위치 인코딩 개입: 학습된 절대 임베딩, RoPE, ALiBi 를 비교하여 위치 불변 저장 및 특정 라우팅 기하학이 역방향 전이에 직접적인 영향을 미침을 검증했습니다.
   • 대규모 검증: 관계 반전 벤치마크 (Parent-Child, Person-Description, T-REx) 에서 다층 확산 LLM(LLaDA, Dream) 과 자기회귀 베이스라인 (LLaMA, Qwen) 을 파인튜닝하여, 이론적 서명이 심층 모델에서도 유지되는지 확인했습니다.

주요 기여

1. 이론적 메커니즘: 이 논문은 MDM 의 반전 저주 완화가 단순히 임의 순서 목표 때문이 아니라, 다음을 통해 순방향 및 역방향 조건부가 **결합 (coupling)**됨에 기인함을 규명합니다:
   • 입력 - 목표 증거의 위치 불변 저장.
   • 상대적 위치 인코딩 하의 양적으로 상관된 어텐션 라우팅.
   • 공유 저장 파라미터에 대한 양적으로 정렬된 그래디언트.
2. 공식적 증명: 저자들은 다음을 보여주는 증명을 제공합니다:
   • 순방향 학습이 공유 증거 항 $Y_{A,B}$를 강화함.
   • 상대적 위치 인코딩 (RoPE, ALiBi) 이 순방향 및 역방향 어텐션 점수 간의 양적 상관관계를 보존함.
   • 그래디언트 정렬로 인해 순방향 업데이트가 1 차적으로 역방향 손실을 감소시킴.
3. 실증적 검증: 이 연구는 제어된 단층 설정과 대규모 다층 확산 모델 모두에서 이러한 서명을 검증하여, 이론적 메커니즘이 향상된 역방향 예측 확률로 이어짐을 보였습니다.

결과

• 단층 검증:
  • 어텐션 상관관계: 순방향 및 역방향 어텐션 점수가 증거 토큰에 할당될 때, 서로 다른 상대적 거리에 걸쳐 강력하게 양적으로 상관관계를 유지합니다 (Fig. 3, 왼쪽).
  • 그래디언트 정렬: 순방향 및 역방향 그래디언트 간의 코사인 유사도는 양적으로 유지됨 (약 >0.6) 을 보이며, 이는 정렬된 최적화 방향을 나타냅니다 (Fig. 3, 중앙).
  • 예측 전이: 순방향 전용 학습 하에서, 역방향 질의에서 올바른 목표 토큰의 확률이 현저히 증가합니다 (Fig. 3, 오른쪽).
  • 베이스라인 실패: 반면, 단층 ARM 은 순방향은 학습하지만 역방향으로 전이하지 못합니다 (Table 1).
• 위치 인코딩 영향:
  • 절대 임베딩: 값 벡터를 위치 의존적으로 만들어 공유 저장 항을 깨뜨리므로 역방향 전이가 poor 합니다.
  • RoPE: 부분적인 어텐션 상관관계로 인해 비자명하지만 불완전한 역방향 전이를 가능하게 합니다.
  • ALiBi: 콘텐츠 점수를 완벽하게 보존하여 가장 강력한 라우팅 상관관계를 유지하므로, 가장 강력한 전이 (거의 완벽함) 를 산출합니다 (Table 2).
• 대규모 검증:
  • 다층 확산 모델 (LLaDA, Dream) 은 동일한 정성적 서명을 보입니다: 양적 어텐션 상관관계, 양적 그래디언트 정렬, 그리고 순방향 전용 파인튜닝 중 향상된 역방향 확률 (Fig. 4).
  • ARM(추론 지향 모델인 Qwen-3.5 포함) 은 높은 순방향 정확도를 보이지만 역방향 정확도에서는 붕괴됩니다 (Table 7).

중요성 및 주장

이 논문은 "임의 순서 학습이 도움이 된다"는 관찰을 넘어, MDM 이 반전 저주를 완화하는 이유에 대한 **기작적 설명 (mechanistic explanation)**을 제공한다고 주장합니다. 그 중요성은 역방향 전이에 필요한 구체적인 아키텍처 및 학습 역학을 규명하는 데 있습니다:

• 저장: 증거는 위치 불변 방식으로 저장되어야 합니다 (값 벡터가 아닌 어텐션 점수에만 영향을 미치는 상대적 위치 인코딩을 통해 달성됨).
• 라우팅: 모델은 역방향 구성에서 이 증거를 검색하기 위해 상관된 경로를 유지해야 합니다.
• 역학: 학습 과정은 공유 증거에 대해 순방향 및 역방향 그래디언트를 자연스럽게 정렬해야 합니다.

저자들은 자신의 주장에 대해 겸손하게 다음과 같이 언급합니다:

• 이 이론은 역방향 전이가 왜 발생할 수 있는지 설명하고 전이 효과의 부호를 예측하지만, 복잡하고 다중 토큰 환경에서 완벽한 역방향 정확도 (확률이 1.0 에 근접) 를 보장하지는 않습니다.
• "반전 저주" 완화는 저장 - 라우팅 메커니즘에 특화된 것이며, 반드시 MDM 이 우수한 일반 논리적 추론 능력이나 추상적 역전 능력을 갖췄음을 의미하지는 않습니다.
• 단층 이론은 최소 기작적 설명으로 기능하며, 이러한 서명이 심층 모델에서 유지되는 것은 이론적 보장이 아닌 실증적 발견입니다.

요약하자면, 이 논문은 MDM 의 반전 저주 완화가 **파라미터 수준의 결합 (parameter-level coupling)**에서 비롯됨을 확립합니다. 여기서 공유 저장과 상관된 라우팅은 순방향으로 학습된 증거가 순서가 반전된 질의 하에서도 접근 가능하게 하며, 이는 표준 자기회귀 학습에서는 근본적으로 붕괴되는 속성입니다.