기술 요약: 마스킹된 이산 시퀀스 모델에서의 쌍별 상호 정보 신경 추정

1. 문제 정의

마스킹 확산 모델 (MDM) 은 텍스트, 단백질, 스도쿠와 같은 이산 시퀀스를 생성하는 강력한 생성 모델로, 자기회귀 (AR) 모델의 고정된 회귀 순서를 피합니다. 그러나 표준 MDM 은 주로 주변 조건부 분포($p(x_i | x_{\text{context}})$) 만 노출하며 변수 간 의존성을 명시적으로 표현하지는 않습니다.

이러한 명시적 의존성 모델링의 부재는 두 가지 주요 과제를 야기합니다:

1. 해석 가능성: 변수들이 서로 어떻게 관련되는지에 대한 모델의 내부 신념 구조를 이해하기 어렵습니다.
2. 병렬 디코딩의 효율성: 현재 병렬 디코딩 전략 (예: Mask-Predict, EB-Sampler) 은 일반적으로 어떤 토큰들을 동시에 마스킹 해제할지 결정하기 위해 주변 신뢰도(엔트로피) 에 의존합니다. 이 접근법은 쌍별 의존성을 고려하지 못합니다. 서로 조건부로 의존하는 토큰들 (높은 상호 정보) 을 서로에 대한 조건 없이 동시에 마스킹 해제하면 스도쿠 규칙이나 단백질 구조적 제약과 같은 전역적 불일치가 발생합니다. 이는 종종 순차적 디코딩으로의 후퇴를 강요하거나 저품질 생성으로 이어집니다.

상호 정보 (MI) 의 전통적인 계산은 밀도 추정이 필요하기 때문에 고차원 환경에서 계산적으로 처리 불가능합니다.

2. 방법론

저자들은 사전 훈련된 MDM 의 은닉 상태로부터 직접 쌍별 조건부 상호 정보($I(X_i; X_j | C)$) 를 추정하기 위한 신경 프레임워크를 제안합니다. 이 접근법은 세 가지 주요 구성 요소로 이루어집니다:

A. 정답 MI 계산 (지도 신호)

가벼운 추정기를 훈련하기 위해 저자들은 사전 훈련된 MDM 의 자체 조건부 분포에 기반하여 "정답" MI 를 계산하는 정확하지만 비용이 많이 드는 방법을 먼저 정의합니다.

• 정의: 컨텍스트 $C$(마스킹 해제된 토큰) 에 대해, 두 개의 마스킹된 위치 $i$와 $j$ 사이의 MI 는 결합 분포 $P(X_i, X_j | C)$와 주변 분포들의 곱 사이의 KL 발산으로 정의됩니다.
• 계산 전략: MDM 이 주변 분포를 출력하므로, 저자들은 교란 기반의 브루트 포스 프로빙 전략을 사용합니다:
  1. 기본 패스: 마스킹된 시퀀스에 모델을 실행하여 주변 분포 $P(X_i | C)$를 얻고 개별 엔트로피 $H(X_i | C)$를 계산합니다.
  2. 조건부 패스: 각 위치 $i$와 가능한 모든 토큰 $v$에 대해, $X_i = v$로 고정하고 순전파를 실행하여 조건부 분포 $P(X_j | X_i=v, C)$를 얻습니다.
  3. 계산: 조건부 엔트로피 $H(X_j | X_i, C)$를 계산하고 엔트로피 감소량으로 MI 를 유도합니다: $I(X_i; X_j | C) = H(X_j | C) - H(X_j | X_i, C)$.
• 비용: 이는 $1 + N \cdot |V|$번의 순전파를 필요로 하므로 추론에는 비실용적이지만 훈련 데이터 생성에는 적합합니다.

B. 신경 MI 추정기

가벼운 신경망 ($f_\phi$) 이 사전 훈련된 MDM 의 고정된 은닉 상태 ($h$) 로부터 MI 행렬을 직접 근사하도록 훈련됩니다.

• 아키텍처: 추정기는 은닉 상태 $h \in \mathbb{R}^{N \times D}$를 입력받아 모든 위치에 대한 추정된 쌍별 MI 를 나타내는 대칭 행렬 $\hat{I} \in \mathbb{R}^{N \times N}$을 출력합니다.
• 훈련 목적: 모델은 마스킹된 인덱스에 대해 예측된 행렬 $\hat{I}$와 정답 행렬 $M_{GT}$ 사이의 평균 제곱 오차 (MSE) 를 최소화하도록 훈련됩니다.

C. MI 기반 병렬 샘플링

저자들은 마스킹 해제된 토큰들 간의 조건부 독립성을 보장하기 위해 예측된 MI 행렬을 활용하는 병렬 디코딩을 위한 탐욕적 선택 알고리즘을 도입합니다.

• 전략: 단순히 가장 낮은 엔트로피 (가장 높은 신뢰도) 를 가진 토큰을 선택하는 대신, 알고리즘은 컨텍스트가 주어졌을 때 상호 독립적인 토큰 뱅크 $S$를 선택합니다.
• 알고리즘:
  1. 엔트로피가 증가하는 순서로 마스킹된 인덱스를 정렬합니다 (가장 높은 신뢰도부터).
  2. 후보들을 순회하며 의존성 비용을 계산합니다: $d(i | U) = \sum_{j \in U} \hat{I}_{i,j}$, 여기서 $U$는 이미 선택된 토큰들의 집합입니다.
  3. 토큰 $i$의 총 비용 (엔트로피 + $\lambda \times$ 의존성 비용) 이 남은 예산 $\gamma$ 내에 있을 때만 토큰 $i$를 선택합니다.
  4. 비용이 너무 높으면 (이미 선택된 토큰들과 높은 MI 를 의미함), 해당 토큰은 순차적 단계로 연기됩니다.
• 결과: 이를 통해 높은 상관관계를 가진 변수들은 순차적으로 처리되고, 조건부 독립적인 부분집합들은 병렬로 처리됩니다.

3. 주요 기여

1. 신경 MI 추정 프레임워크: 추론 중 비용이 많이 드는 밀도 추정이나 정답 계산을 필요로 하지 않고 MDM 은닉 상태로부터 직접 쌍별 조건부 MI 를 추정하는 방법.
2. MI 기반 병렬 디코딩: 추정된 MI 를 사용하여 조건부 독립적인 변수 부분집합을 식별하는 새로운 샘플링 전략으로, 전역 일관성을 유지하면서 안전한 병렬화를 가능하게 함.
3. 해석 가능성 도구: MI 맵은 명시적 프로그래밍 없이 스도쿠 규칙이나 단백질 접힘 의존성과 같은 학습된 제약 조건을 드러내는 모델의 내부 신념 구조를 시각화하는 도구로 작용함.

4. 실험 결과

이 접근법은 스도쿠(구조화된 논리) 와 단백질 시퀀스 생성(ESM-C 사용) 두 가지 도메인에서 평가되었습니다.

스도쿠

• 설정: 100,000 개의 퍼즐로 훈련; 1,000 개의 보지 못한 어려운 퍼즐에서 평가.
• 성능:
  • 순차적 베이스라인: 평균 53.9 회 순전파, 61.6% 정확도.
  • 단순 병렬 (k=7): 9.0 회 순전파, 그러나 정확도는 36.8% 로 하락.
  • MI 기반 (γ=0.3): 15.2 회 순전파, 63.6% 정확도(순차적 베이스라인을 상회).
  • MI 기반 (γ=0.6): 9.7 회 순전파, 56.2% 정확도.
• 관찰: MI 기반 샘플러는 순차적 디코딩에 비해 순전파 횟수를 3~5 배 줄였으며, 단순 병렬 방법에 비해 정확도를 유지하거나 향상시켰습니다.

단백질 시퀀스 (ESM-C)

• 설정: 500 개의 무작위 단백질 (길이 50~100) 을 생성하고, 500 개의 UniRef50 참조 샘플과 Jensen-Shannon 발산 (JSD) 을 비교.
• 성능:
  • 순차적: 74.8 회 순전파, JSD 0.093.
  • 단순 병렬 (k=12): 6.2 회 순전파, JSD 0.218 (품질 저하 심각).
  • MI 기반 (γ=4): 10.0 회 순전파, JSD 0.174.
• 관찰: MI 기반 샘플링은 단순 병렬 베이스라인보다 더 나은 속도 - 정확도 트레이드오프를 달성했으며, 순차적 방법에 비해 순전파 횟수를 거의 한 자릿수 수준으로 크게 줄이면서도 엔트로피 기반 방법보다 생성 품질을 더 잘 유지했습니다.

5. 의의 및 주장

이 논문은 변수 의존성을 명시적으로 모델링하는 것이 이산 확산 모델의 잠재력을 최대한 끌어내는 데 필수적이라고 주장합니다.

• 격차 해소: 이 연구는 순차적 샘플링의 높은 품질과 병렬 디코딩의 효율성 사이의 격차를 해소합니다.
• 내부 표현: MI 맵은 MDM 이 명시적 프로그래밍 없이도 스도쿠 규칙이나 단백질 의존성과 같은 견고한 구조적 제약을 자연스럽게 습득하며, 이러한 것들을 추정기를 통해 추출할 수 있음을 보여줍니다.
• 효율성: 이 방법은 조건부 독립적인 부분집합을 식별하는 MI 기반 병렬 디코딩을 가능하게 하여, 순차적 디코딩에 비해 추론 시간 순전파 횟수를 3~5 배 수준으로 감소시킵니다.

인정된 한계:
저자들은 예측기가 완벽하지 않으며 상당한 설정과 훈련 (훈련 데이터에 대한 정답 MI 의 실시간 계산 포함) 이 필요하다고 지적합니다. 향후 연구로는 훈련 단계 중 정답 계산의 계산 비용을 피하기 위해 최적의 예측기 아키텍처와 개선된 커리큘럼 훈련 전략을 조사할 것을 제안합니다.