Insertion Based Sequence Generation with Learnable Order Dynamics

Motivation

• Known: 기존 masked diffusion model(MDM)은 고정 길이 시퀀스에서 절대 위치 기반 unmasking으로 생성을 수행하며, FlexMDM과 같은 insertion 기반 모델은 가변 길이 시퀀스를 다루기 위해 mask token의 insertion과 unmasking을 교차 수행하지만 목표 스케줄이 데이터와 무관하게 고정되어 있다.
• Gap: 그래프나 분자와 같이 구조화된 시퀀스는 생성 순서에 따라 국소적 예측 용이성이 크게 달라지는데, 고정된 스케줄을 사용하면 모델이 모든 가능한 생성 순서에 확률을 분산시켜야 하므로 비효율적이며, 데이터에 특화된 최적 순서를 학습할 방법이 없었다.
• Why: 생성 순서를 데이터에 맞게 학습할 수 있다면 action space에 대한 불확실성을 줄여 sample quality를 높일 수 있고, 이는 그래프·분자 등 구조화된 이산 데이터 생성 전반에 적용 가능한 tractable하고 확장성 있는 프레임워크로 이어질 수 있다.
• Approach: discrete flow matching(DFM) 프레임워크를 가변 길이 시퀀스 및 projected partial sequence에 대해 일반화하고, Kumaraswamy CDF 기반의 tractable한 스케줄 파라미터화와 REINFORCE leave-one-out estimator를 이용한 joint training objective를 제안하여 생성기(generator)와 목표 순서 동역학(target order dynamics)을 동시에 학습한다.

Achievement

Figure 4a, LoFlexMDM indeed learns to generate in the

1. 가변 길이 시퀀스에 대한 DFM 일반화: Discrete Flow Matching을 projected partial sequence 공간으로 확장하여 insertion과 unmasking을 아우르는 조건부 rate matching과 ELBO를 유도했다.
2. Tractable한 학습 가능 스케줄: Kumaraswamy CDF를 이용해 insertion time과 unmasking time의 밀도를 파라미터화함으로써 event time을 병렬 샘플링하면서도 closed-form likelihood를 확보했다.
3. 효율적인 joint training: REINFORCE leave-one-out estimator를 활용해 전체 trajectory를 시뮬레이션하지 않고도 generator와 target rate model을 동시에 학습할 수 있는 방법을 제시했다.
4. 성능 향상 입증: De Novo 분자 생성에서 최대 17.5%p, fragment-constrained 분자 생성에서 최대 6.7%p의 sample quality 개선을 FlexMDM 대비 달성했다.
5. 학습된 순서의 구조적 타당성 확인: 그래프 순회 및 분자 생성 과제에서 LoFlexMDM이 태스크 구조에 부합하는 생성 순서를 학습함을 정성적으로 보였다.

How

Figure 1. Left: The left subfigure illustrates how the endpoint-dependent target process grows the final sequence x1 fro

• 보조 공간(auxiliary space) Z를 partial sequence와 위치 튜플의 곱 공간으로 정의하고, [D](dropped) 토큰을 이용한 fcontract 매핑으로 ¯Z와 Z 사이의 동형사상(isomorphism)을 구성한다.
• Projected DFM 프레임워크를 통해 auxiliary process 위에서 정의된 target rate를 원래 공간 X로 projection하여 학습 가능한 generator rate Rθ_t와 매칭시킨다 (Bregman divergence 최소화).
• 목표 스케줄 f^ϕ_Tini, f^ϕ_Tumi|Tini를 Kumaraswamy CDF로 파라미터화하여 각 위치별 insertion time과 unmasking time을 병렬적으로 샘플링 가능하게 하고 closed-form 밀도를 얻는다.
• Auxiliary network ϕ는 clean sequence x1을 입력받아 위치별 insertion rate λ^ϕ_in, unmasking rate λ^ϕ_um을 출력하고, generator network θ는 transformer backbone(DDiT)과 LM head, 두 개의 rate head(insertion, unmasking·K)를 통해 xt로부터 대응하는 rate를 예측한다.
• 두 네트워크(θ, ϕ)를 REINFORCE leave-one-out estimator 기반의 ELBO objective로 공동 학습하여, 전체 궤적 시뮬레이션 없이도 tractable하게 학습 가능하도록 한다.

Originality

• 기존 insertion 기반 diffusion model(FlexMDM 등)이 고정된 데이터 독립적 스케줄만을 사용했던 것과 달리, 데이터 의존적이고 학습 가능한 insertion/unmasking rate를 도입해 생성 순서 자체를 학습 대상으로 삼은 최초의 접근이다.
• Discrete Flow Matching을 projected partial sequence 및 가변 길이 시퀀스로 확장한 이론적 일반화(Projected DFM)를 제시했다.
• Kumaraswamy CDF를 이용한 tractable schedule parameterization과 REINFORCE leave-one-out을 결합해, 전체 궤적을 시뮬레이션하지 않고도 generator와 target order dynamics를 함께 학습하는 새로운 training objective를 고안했다.

Limitation & Further Study

• REINFORCE 기반 estimator는 일반적으로 분산이 크고 학습 안정성 문제가 있을 수 있어, 대규모·복잡한 데이터셋에서의 확장성 검증이 추가로 필요하다.
• 실험이 주로 분자 생성(De Novo, fragment-constrained)과 그래프 순회 같은 특정 구조적 태스크에 한정되어 있어, 자연어와 같은 비구조적 시퀀스에 대한 일반화 가능성은 추가 검증이 필요하다.
• 두 네트워크(generator, target rate model ϕ)를 함께 학습해야 하므로 계산 비용과 학습 복잡도가 FlexMDM 대비 증가할 가능성이 있으며 이에 대한 분석이 제한적이다.
• 학습된 생성 순서가 실제로 해석 가능하고 태스크 구조에 부합하는지에 대한 정량적/체계적 평가(단순 정성적 관찰 이상)가 추가로 필요하다.

Evaluation

총평: insertion 기반 masked diffusion model에서 생성 순서를 데이터 의존적으로 학습 가능하게 만든 이론적으로 탄탄하고 실용적인 기여이며, 구조화된 시퀀스 생성 분야에 의미 있는 발전을 제시하나 더 넓은 도메인에서의 검증이 필요하다.