🤖 machine learning

Insertion Based Sequence Generation with Learnable Order Dynamics

이 논문은 분자 설계와 같은 구조화된 시퀀스 생성 작업에서 학습 가능한 데이터 의존적 생성 순서를 학습하여 샘플 품질을 향상하는 동시에 다루기 쉬운 훈련을 유지하는 삽입 기반 마스크 확산 모델인 LoFlexMDM을 소개한다.

원저자: Dhruvesh Patel, Benjamin Rozonoyer, Gaurav Pandey, Tahira Naseem, Ramón Fernandez Astudillo, Andrew McCallum

게시일 2026-06-09

📖 7 분 읽기🧠 심층 분석

CC BY 4.0

원저자: Dhruvesh Patel, Benjamin Rozonoyer, Gaurav Pandey, Tahira Naseem, Ramón Fernandez Astudillo, Andrew McCallum

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

LoFlexMDM — Layman's Summary

The Big Idea: Planning a Route Through a City

This paper is about insertion-based diffusion models, a way to generate structured objects whose size varies: a molecule, a graph path, a string that grows as you go. Instead of writing left to right like a typewriter, the model starts from a blank canvas. It inserts blank placeholders into gaps between what it already has, then reveals what each placeholder actually is.

Each new position goes through two steps:

Insertion — place a blank mask token in a gap.
Unmasking — reveal what that token actually is (an atom, an edge, a bracket, and so on).

Because the model can insert at any gap and reveal pieces in any order, there are many valid ways to assemble the same final object. The order it chooses matters a lot. That is the question this paper studies, and a delivery route makes the idea concrete.

Imagine you need to deliver ten packages across a large city. You must plan a route: which roads to take, in what order, to connect all the stops.

A natural strategy is to start from the delivery addresses and work inward, choosing outgoing roads that link them together. Each stop tells you which road you came from. The route builds itself step by step.

A bad strategy is to pick a road that lies somewhere between two addresses but does not yet connect to either of them. It might turn out to be useful later, but right now you are guessing. Some of those guesses look fine locally but dead-end once you try to link them up.

Why Order Matters: Not All Routes Are Equal

The difficulty is not just that some orders are slightly worse. For structured sequences, the wrong order can turn an easy task into one that essentially requires planning ahead.

The paper's star graph example is the clearest case. Picture a star-shaped road network: one central junction where several arms meet. You are given a start node and an end node on opposite sides of the graph. The task is to output the correct path connecting them.

Easy order: Start at the two endpoints and work inward along each arm toward the junction. Every step along an arm is locally obvious. You are just following the road you are already on. By the time you reach the junction, only one edge remains, and it is trivial to pick.
Hard order: Begin at the junction and pick outgoing roads at random, or hop between arms in arbitrary order. At the junction you must choose which arm to take next. That choice depends on the full path ahead. A model that only looks at what is already on the board cannot reliably make that decision.

The paper uses this task because the locally optimal order is known: endpoints first, junction last.

Molecules follow a similar pattern. Ring closures, backbone connectivity, and fragment attachment points are structural commitments — the "addresses" and major intersections of the molecule. Lock those in first and the remaining atom and bond labels become much easier. Commit to fine details before the skeleton is settled and you waste capacity on choices that may not fit the final structure.

The Problem: FlexMDM's Fixed Schedule

The previous approach, FlexMDM, keeps the two-step insert-and-unmask process but uses fixed, data-independent schedules. Every token gets the same timing rules regardless of what is being built.

Back to the delivery analogy:

FlexMDM tries to connect the locations by picking roads in no particular order. Its schedule does not depend on the destinations, so it has to be ready to build the route in every possible order, including the wasteful ones that start with disconnected roads in the middle.

That makes the planning problem much harder than it needs to be. Instead of mastering one good way to build the route, the model has to handle countless tangled starting points, most of which lead nowhere useful. Spreading its effort across all of them leaves less for the orders that actually work.

The Solution: LoFlexMDM — Learn the Right Order for Each Object

LoFlexMDM (Learnable order FlexMDM) keeps the same generator machinery but replaces the fixed schedule with learnable, example-dependent rates. The key is a two-part training setup. During training, a helper network (the "schedule network") is allowed to see each finished object and learns a good order for building it — when each position should be inserted and unmasked. The generator then learns to imitate that order. At generation time the helper is gone and the model does not see any answer; it simply follows the building habits it picked up, producing new objects in the kind of order that worked during training.

Those learned timings implicitly define a distribution over generation orders tailored to the type of object being built.

Back to the delivery route:

Instead of a universal timer, the model learns a route plan for each specific delivery job.

On a star graph, it learns to start from the arm tips and work inward, leaving the junction for last — because that is the order where each step is locally easy.
On a molecule, it learns to fix rings, separators, and attachment scaffolding before filling in atom labels and side branches — because structure constrains chemistry.

The main generator and the schedule network are trained together. The schedule network proposes an order; the generator learns to execute it; feedback from generation quality pushes both toward orders that the generator can actually complete.

How It Works (Simplified)

For each position in the final sequence:

Sample an insertion time — when the mask token appears.
Sample an unmasking time — when the mask becomes a real symbol (always after insertion).

Positions with earlier unmasking times are generated first. The schedule network outputs per-position rates that control these times. The generator network matches those target rates while learning to produce the correct tokens.

The technical trick is a parameterization that keeps this tractable: event times can be sampled in parallel, likelihoods have closed form, and the generator and schedule network can be trained jointly without simulating full trajectories step by step.

Results

Star graph traversal (the clearest proof)

On the hard star-graph setting (variable arm lengths), where the optimal inside-out order is well understood:

Model	Exact match (hard)
FlexMDM	~7%
LoFlexMDM	~88%

LoFlexMDM learns an order correlated with distance from the junction: outer edges first, junction last. FlexMDM does not.

Molecule generation

On de novo small-molecule generation (SAFE strings):

Sample quality improves by about 17.5 percentage points over FlexMDM (roughly 79.5% vs 62.0%).
Validity stays very high (up to ~99.7%).

On fragment-constrained generation, quality improves by up to 6.7 percentage points across four tasks.

The learned order is interpretable: ring closures and structural separators tend to come before atom-level details; attachment where tends to come before attachment which.

Focusing on a sharper order can trade a small amount of diversity for quality, which is expected when the model commits to structure earlier rather than exploring many equivalent assembly paths.

Summary in One Sentence

LoFlexMDM teaches an insertion-based diffusion model not just what to build, but in what order to build each specific object — because for structured sequences like star graphs and molecules, the right order makes the task easy and the wrong order makes it nearly impossible, and learning that order per example yields much better generation quality without giving up tractable training.

기술 요약: LoFlexMDM

문제 정의

기존의 삽입 기반 마스크 확산 모델(MDMs)은 토큰 삽입과 언마스킹(unmasking)을 교차하며 시퀀스를 생성합니다. 그러나 FlexMDM과 같은 현재의 접근 방식은 데이터와 독립적인 고정된 스케줄에 따라 삽입 및 언마스킹 비율을 결정합니다. 이러한 한계는 그래프나 분자와 같이 자연스럽게 가변 길이를 가지며, 절대적 인덱스보다 상대적 위치를 통해 더 잘 표현되는 의존성을 가진 구조화된 시퀀스에서 특히 문제가 됩니다.

이러한 구조적 도메인에서 최적의 생성 순서는 종종 데이터에 따라 달라집니다. 예를 들어, 별 모양(star-shaped) 그래프의 경우, 접점으로 이어지는 엣지는 국소적으로 예측 가능하지만, 접점을 빠져나가는 엣지는 앞을 내다보는 과정(lookahead)이 필요합니다. 고정된 스케줄은 모델이 가능한 모든 순서로 시퀀스를 생성하도록 강제하며, 이로 인해 확률 질량이 잠재적으로 최적이 아닌 생성 경로에 분산됩니다. 이러한 "낭비적인" 학습 과정은 모델이 토큰 그룹 간의 상관관계를 포착하는 능력을 저하시키고 생성 품질을 제한합니다.

방법론

저자들은 고정된 스케줄 대신 학습 가능한, 데이터 의존적인 삽입 및 언마스킹 비율을 사용하는 삽입 기반 마스크 확산 모델인 LoFlexMDM(Learnable Order FlexMDM)을 제안합니다. 핵심 아이디어는 생성 순서를 학습된 비율에 의해 암시적으로 정의되는 가능한 순서들에 대한 분포로 취급하는 것입니다.

주요 기술 구성 요소

Hazard Rates를 통한 학습 가능한 순서 역학(Order Dynamics):
고정된 스케le듈 대신, LoFlexMDM은 전이율이 타겟 시퀀스의 함수인 연속 시간 마르코프 체인(CTMCs)을 사용하여 생성 과정을 모델링합니다. 각 위치 $i$ 에 대해, 모델은 두 가지 시간 의존적 해저드 비율(hazard rates)을 학습합니다:
- $\lambda_{in}^i(t)$ : 위치 $i$ 에서 마스크 토큰 $[M]$ 을 삽입하는 비율 (전이 $[D] \to [M]$ ).
- $\lambda_{um}^i(t)$ : 위치 $i$ 에서 토큰을 언마스킹하는 비율 (전이 $[M] \to x_i$ ).
  이러한 비율은 깨끗한 타겟 시퀀스 $x_1$ 을 입력으로 받는 신경망 $\phi$ 에 의해 매개변수화됩니다.
Kumaraswamy CDF를 이용한 매개변수화:
학습과 샘플링의 용이성을 보장하기 위해, 저자들은 삽입 및 언마스킹 시간의 누적 분포 함수(CDFs)를 Kumaraswamy 분포를 사용하여 매개변수화합니다.
- $F(t) = 1 - (1 - t^a)^b$ .
- 이러한 선택은 역 CDF 샘플링(이벤트 시간의 병렬 샘플링 가능)을 가능하게 하며, 수치적 적분 없이 라이클리후드(likelihood)와 해저드 비율에 대한 **폐쇄형 표현(closed-form expressions)**을 제공합니다.
- 저자들은 형상 파라미터 $a$ 를 상수로 설정하고 스케일 파라미터 $b_{in}$ 및 $b_{um}$ 을 학습함으로써, 공유된 형상 함수를 유지하면서 전체 비율을 제어할 수 있도록 했습니다.
투영된 이산 흐름 매칭(Projected Discrete Flow Matching, DFM):
이 프레임워크는 정렬 정보(누락되지 않은 토큰의 위치)를 포함하는 보조 공간 $Z$ 를 사용하여 가변 길이 시퀀스로 Discrete Flow Matching을 확장합니다.
- 모델은 보조 공간에서의 타겟 비율을 매칭하도록 학습되면서, 부분 시퀀스의 공간 $X$ 에 대한 "투영된" 비율을 정의합니다.
- 이를 통해 생성기는 보조 공간에 정의된 타겟 비율을 따르면서도 가변 길이 시퀀스 $x_t$ 상에서만 동작할 수 있습니다.
결합 학습 목적 함수(Joint Training Objective):
모델은 생성기 파라미터 $\theta$ 와 타겟 비율 파라미터 $\phi$ 를 공동으로 학습합니다.
- 손실 함수: 목적 함수는 생성기의 예측된 비율과 $\phi$ 로부터 유도된 타겟 비율 사이의 Bregman 발산을 최소화합니다.
- 그래디언트 추정: 기대값이 타겟 비율 $\phi$ 에 의존하므로, 저자들은 $\phi$ 에 대한 그래디언트를 계산하기 위해 REINFORCE leave-one-out 추정기를 사용합니다. 이는 두 개의 부분 상태( $z^{(1)}_t, z^{(2)}_t$ )를 타겟 프로세스로부터 샘플링하고, 손실의 차이를 어드밴티지 신호로 사용하여 $\phi$ 를 업데이트하는 동시에, $\theta$ 에는 표준 pathwise 그래디언트를 사용합니다.
샘플링:
생성에는 $\tau$ -leaping 샘플러를 사용하여 한 타임 스텝 내에서 여러 개의 삽입 및 언마스킹 이벤트를 동시에 샘플링할 수 있어 효율성을 높입니다. 이 샘플러는 높은 예측 확신도를 가진 위치를 우선적으로 언마스킹하는 신뢰도 기반 위치 선택 기능을 통합할 수 있습니다.

주요 기여

학습 가능한 순서 역학: 고정된 스케줄을 넘어 데이터 의전적 생성 순서를 학습하는 최초의 삽입 기반 마스크 확산 모델입니다.
일반화된 DFM 프레임워크: 보조 공간에서의 투영된 비율 매칭을 통해 가변 길이 시퀀스를 처리할 수 있도록 Discrete Flow Matching을 확장했습니다.
용이한 매개변수화: 타겟 스케줄을 매개변수화하기 위해 Kumaraswamy CDF를 도입하여 폐쇄형 라이클리후드와 효율적인 병렬 샘플링을 가능하게 했습니다.
결합 학습 알고리즘: 전체 궤적을 시뮬레이션하지 않고도 생성기와 타겟 순서 역학을 공동으로 최적화하기 위해 REINFORCE leave-one-out 추정법을 사용하는 새로운 학습 절차를 제시했습니다.

실험 결과

저자들은 star graph traversal과 소분자 생성(de novo 및 fragment-constrained)이라는 두 가지 주요 작업에서 LoFlexMDM을 평가합니다.

Star Graph Traversal:
- 작업: 스타 그래프에서 두 노드 사이의 경로를 생성합니다. 최적의 전략은 양 끝단에서 중심부(junction)를 향해 안쪽으로 엣지를 생성하는 것입니다.
- 결과: LoFlexMDM은 이 최적의 순서를 학습하여, "hard" 설정(가변 팔 길이)에서 FlexMDM의 6.0% 대비 **87.9%**의 정확 일치율(exact match rate)을 달anim했습니다. 이는 순서를 학습하는 것이 액션 공간의 불확실성을 크게 줄인다는 것을 입증합니다.
소분자 생성 (Small Molecule Generation):
- De Novo 생성: LoFlexMDM은 높은 유효성(validity)과 다양성을 유지하면서 FlexMDM 대비 샘플 품질을 17.5% 개선했습니다 (79.5% vs 62.0%).
- Fragment-Constrained 생성: 모티프 확장(motif extension) 및 스캐폴드 장식(scaffold decoration)과 같은 작업에서 LoFlexMDM은 품질을 최대 **6.7%**까지 향상시켰습니다 (예: 모티프 확장에서 53.6% vs 46.9%).
- 분석: 모델은 분자를 위한 구조화된 생성 순서를 학습하여, 링 폐쇄(ring closures)와 단편 구분자(fragment separators)를 조기에 해결하고, 특정 원자와 라벨을 나중에 채워 넣습니다. 이는 FlexMDs가 생성 타이밍에서 체계적인 패턴을 보이지 않는 것과 대조적입니다.

의의 및 주장

본 논문은 타겟 생성 순서를 학습하는 것이 구조화된 데이터를 위한 삽입 기반 확산 모델의 성능을 향상시키는 데 결정적인 요소라고 주장합니다. 데이터를 샘플에 맞게 조정함으로써, LoFlexMDM은 학습의 용이성을 희생하지 않으면서 액션 공간에 대한 불확실성을 줄여 더 높은 품질의 샘플을 생성합니다.

저자들은 이 접근 방식이 고정 길이 언마스킹, 가변 길이 언마스킹, 그리고 삽입 전용 모델 모두에 일반적이고 적용 가능하다는 점을 강조합니다. 또한, 생성기와 타겟 비율 모델이 가중치를 공유하는 백본 아키텍처(shared backbone)가 확장에 유리할 수 있음에도 불구하고, 실험에서는 별도의 백본을 사용할 때 더 큰 이득이 관찰되었다고 언급했습니다. 이는 순서 역학을 학습하는 복잡성이 전용 용량(dedicated capacity)을 필요로 할 수 있음을 시사합니다. 이 연구는 가변 길이 구조적 시퀀스를 위한 더 효율적이고 고품질의 생성 모델링을 향한 길을 열어줍니다.