기술 요약: 삽입 기반 생성을 위한 변분 학습 (Variational Learning for Insertion-based Generation)

1. 문제 정의

자기회귀(Autoregressive) 시퀀스 모델은 일반적으로 고정된 좌측-우측(left-to-right) 순서로 토큰을 생성한다. 이는 많은 작업에서 효과적이지만, 이러한 인수 분해(factorization)는 모든 데이터 구조에 내재된 것은 아니며, 출력의 일부가 미래의 컨텍스트에 의존할 수 있는 생물학적 서열 설계, 프로그램 합성, 구조적 계획과 같은 실제 세계의 생성 문제와는 일치하지 않을 수 있다.

마스크 확산 모델(masked diffusion models, 비단조적 생성)과 같은 기존의 대안들은 고정된 길이의 그리드 상에서 마스킹된 부분집합을 업데이트함으로써 유연성을 제공한다. 그러나 이러한 접근 방식은 두 가지 주요 한계점을 가진다:

1. 정적 캔버스(Static Canvas): 이들은 규정된 고정 길이 그리드에서 작동하므로, 외부에서 길이를 지정해야 하며 가변 길이 생성을 간접적으로만 수행할 수 있다.
2. 순서 불가지론(Order Agnosticism): 이들의 훈련 목적 함수는 가능한 업데이트 스케줄을 무시하거나 암묵적으로 가정한다. 이로 인해 모델은 데이터의 기저 구조와 일치하는 작은 부분집합이 존재하더라도, 팩토리얼 단위의 방대한 생성 순서에 대한 조건부 확률을 근사해야 하는 상황에 놓인다.

최근 연구들은 이러한 문제들을 각각 독립적으로 다루어 왔다: 삽입 기반 모델은 임의의 삽입을 통해 가변 길이 시퀀스를 허용하며, 순서 학습(learning-order) 방법은 순열에 대한 분포를 도입한다. 그러나 이를 결합하는 것은 쉽지 않은데, 삽입 모델에서는 삽입 위치가 고정된 좌표가 아니라 부분 시퀀스와 함께 진화하며, 생성 순서와 중간 상태가 밀접하게 결합되어 있기 때문이다.

2. 방법론: 삽입 과정 (Insertion Process, IP)

저자들은 다음을 공동으로 학습하는 확률적 프레임워크인 **삽입 과정(IP)**을 제안한다:

1. 어디에 삽입할 것인가: 삽입 결정에 대한 데이터 기반 선호도 (슬롯 선택).
2. 무엇을 삽입할 것인가: 토큰 분포.
3. 언제 종료할 것인가: 정지 규칙 (정책 기반 또는 분류기 기반).

2.1. 순열을 통한 재매개변수화 (Reparametrization via Permutations)

핵심적인 이론적 기여는 삽입 궤적(insertion trajectories)과 타겟 시퀀스 인덱스의 순열(permutations) 사이의 **전단사 대응(bijective correspondence)**을 공식화한 것이다.

• 과제: 상대적 슬롯 인덱스($z_i$)가 궤적에 따라 달라지기 때문에, 삽입 궤적에 대한 직접적인 변분 추론은 어렵다. 이는 상대적 행동 공간에서 안정적인 순서 후험 분포를 정의하기 어렵게 만든다.
• 해결책: 저자들은 타겟 시퀀스 $y_L$을 생성하는 모든 유효한 삽입 궤적을 고유한 순열 $\sigma \in S_L$에 매핑한다. 여기서 $\sigma_i$는 $i$ 단계에서 채워지는 요소의 인덱스이다.
• 결정론적 매핑: 단계 $i$에서의 상대적 삽입 슬롯은 순열 접두사($\sigma_{\le i}$)의 결정론적 함수가 된다. 구체적으로, $z_i = 1 + |\{j \in \sigma_{<i} : j < \sigma_i\}|$이다.
• 우도 인수 분해 (Likelihood Factorization): 이러한 변수 변환은 궤적 의존적 행동에 대한 주변화(marginalization)를 순열에 대한 정확한 합산으로 변환하여, 데이터 우도의 정확한 재매개변수화를 가능하게 한다:
  $$p(y_L) = \sum_{\sigma \in S_L} p(y_L, \sigma)$$

2.2. 변분 추론 및 훈련

모델을 훈련하기 위해, 저자들은 생성 순서(순열 $\sigma$)를 잠재 변수로 취급하여 데이터 로그 우도에 대한 증거 하한(ELBO)을 최대화한다.

• 변분 후험 분포 (Variational Posterior): 저자들은 플락-루스(Plackett-Luce, PL) 모델로 매개변수화된 암보화된 변분 후험 분포 $q_\theta(\sigma | y_L)$를 도입한다.
• 최적화: 잠재 순서가 이산적(discrete)이므로 목적 함수는 미분 불가능하다. 저자들은 다음을 사용하여 ELBO를 최적화한다:
  • Gumbel-Top-k 트릭: PL 분포로부터 순열을 효율적으로 샘플링하기 위함.
  • Leave-One-Out을 적용한 REINFORCE (RLOO): 그래디언트를 추정하기 위함.
  • 분산 감소 (Variance Reduction): 순열-궤적 구조를 활용함으로써, 고정된 접두사에 조건화된 다음 인덱스 $\sigma_i$에 대한 내부 기댓값을 폐쇄형(closed form)으로 정확하게 계산하여 몬테카를로 추정량의 분산을 줄인다.

2.3. 네트워크 아키텍처

• 생성 디코더 (Generative Decoder): 트랜스포머 기반 디코더가 부분 시퀀스(경계 토큰 포함)를 처리하여 문맥적 임베딩을 생성한다.
  • 위치 헤드 (Location Head): 임베딩을 바탕으로 후보 삽입 슬롯(종료 동작 포함)의 점수를 매긴다.
  • 내용 헤드 (Content Head): 선택된 슬롯에 조건화된 토큰 분포를 예측한다.
• 추론 인코더 (Inference Encoder): 양방향 트랜스포머 인코더가 전체 관찰된 시퀀스를 처리하여 각 글로벌 인덱스에 대한 점수를 출력하며, 이 점수들은 순열 샘플링을 위한 플락-루스 분포를 매개변수화한다.

3. 주요 기여

1. 가변 길이 삽입을 위한 확률적 프레임워크: 순열과의 전단사 매핑을 사용하여 가변 길이 시퀀스에 대한 삽입 순서를 학습하는 최초의 프레임워크이며, 이를 통해 정확한 우도 재매개변수화를 가능하게 한다.
2. 순열 주변화 우도 (Permutation-Marginalized Likelihood): 데이터 우도가 순열의 합으로 표현될 수 있음을 보여주는 이론적 유도를 통해, 잠재 순서 변수를 가진 변분 추론을 용이하게 한다.
3. 효율적인 훈련 알고리즘: 이산적이고 미분 불가능한 잠재 순서의 특성을 다루기 위해 플락-루스 모델, Gumbel-Top-k 샘플링, 그리고 RLOO 그래디언트를 사용하는 분산 감소 훈련 방법을 제공한다.
4. 네이티브 가변 길이 지원: 고정 캔버스 확산 모델과 달리, IP는 가변 길이 생성을 네이티브하게 지원하며 데이터 기반의 삽입 순서 선호도를 학습한다.

4. 실험 결과

저자들은 두 가지 도메인, 즉 목표 조건부 계획(goal-conditioned planning)과 분자 문자열(molecular string) 생성에서 IP를 평가하였다.

4.1. 목표 조건부 계획 (Goal-Conditioned Planning)

• 벤치마크: 미로 계획 (Braided, Imperfect, Perfect) 및 Star Graph 계획.
• 베이스라인: 고정 순서 ARM (FO-ARM), 임의 순서 ARM (AO-ARM), 마스크 확산 모델 (MDM), 순서 학습 ARM (LO-ARM).
• 결과: IP는 모든 미로 패밀리와 그래프 벤치마크에서 최첨단(SOTA) 성능을 달ensis했다.
  • Perfect Maze (고도로 제약된 경우): FO-ARM은 성능이 크게 저하된 반면, IP는 높은 정확도를 유지했다 (예: Hard Perfect Maze에서 FO-ARM의 27.2% 대비 IP는 99.9% 달성).
  • IP는 모든 베이스라인에 비해 시퀀스 및 토큰 정확도를 실질적으로 개선하고 편집 거리(edit distance)를 줄였으며, 이는 장거리 제약이 있는 작업에서 인스턴스별 삽입 순서를 학습하는 것이 매우 중요하다는 것을 입증한다.

4.2. 분자 문자열 (SMILES) 생성

• 벤치마크: GuacaMol 무조건부 및 조건부 생성.
• 결과:
  • 유효성(Validity) 및 신규성(Novelty): IP는 높은 유효성(97.2–97.4%)과 가장 높은 신규성(~95%)을 달성하였으며, 이는 무작위 순서 삽입(AO-IP, 유효성 88.2%)을 크게 상회하는 수치이다. 이는 IP가 화학적 구문 규칙을 성공적으로 내재화했음을 나타낸다.
  • 분포적 지표 (Distributional Metrics): 분포 유사성(KL 발산, FCD) 측면에서는 FO-ARM이 약간 더 우세했으나, IP는 화학적 공간을 더 효과적으로 탐색함으로써 우수한 트레이드오프를 제공했다.
  • 조건부 생성: 고정된 파편 사이에 삽입이 필요한 작업(예: 링커 설계)에서 IP는 약 100%의 유효성을 유지한 반면, 고정 순서 모델은 이러한 제약 조건에서 어려움을 겪었다.
• 해석 가능성 (Interpretability): 학습된 삽입 스케줄 분석 결과, 일관되고 해석 가능한 생성 순서가 나타났다: (1) 분자 골격(고리/괄호) 구축, (2) 원자 삽입, (3) 종료. 이는 명시적인 문법적 사전 지식 없이도 발생하였다.

5. 의의 및 주장

본 논문은 데이터 의존적 삽입 순서를 학습하는 것이 표준적인 좌측-우측 구조가 없는 도메인에서 모델링 품질과 일반화 능력을 크게 향상시킨다고 주장한다. 삽입 궤적과 순열 사이의 연결을 공식화함으로써, 저자들은 삽입 기반 모델을 훈련하기 위한 실행 가능하고 편향되지 않은 변분 방법을 제공한다.

저자들은 IP를 비선형적이거나 자연스러운 생성 순서가 모호한 영역에서의 생성 모델링을 향한 단계로 규정한다. 또한, 향후 연구에서 이 방법을 단백질 설계와 같은 다른 도메인에 적용할 수 있음을 명시하며, IP가 복잡한 구조적 작업에서 일반화 및 유효성을 개선하지만, 단순한 분포 매칭 지표에서는 반드시 고정 순서 모델보다 뛰어난 것은 아니라는 점을 언급하며 유연성과 엄격한 분포 충실도 사이의 트레이드오프를 강조한다.