Diffusion On Syntax Trees For Program Synthesis (Tree Diffusion)

 

Abstract

Syntax tree에 순차적으로 추가된 noise를 반전시켜 역 그래픽 작업을 수행할 수 있는 Tree Diffusion 제안

 

[Project Page]

[Github]

[arXiv](2024/05/30 version v1)

 

 

사실 역 그래픽 작업을 수행할 일이 도대체 어디에 있겠는가?

이 논문에서 Tree Diffusion이 할 수 있는 일 보다 최단거리의 tree를 찾아서 훈련에 사용하는 아이디어가 더 마음에 들었다.

 

 

 

Method

그래픽 작업에서 복잡한 모양을 표현하기 위해 CFG(Context-Free Grammar)를 사용한다.

CFG가 뭐임? 아래 그림과 그 아래 그림을 보면 무슨 느낌인지 알 수 있다.

 

우리의 목표는 x₀을 보고 프로그램 z_T에서 시작하여 z₀를 도출하는 것이다.

 

---

Sampling Small Mutations

z_t에 노이즈를 추가하여 z_t+1을 얻는다.

 

생성되는 z_t+1은 변화의 규모가 일정 범위 σ_min < σ(z) ≤ σ_max 안에 있으면서 문법적으로 유효해야 하기 때문에 어떤 방법을 사용하는데, CFG를 이용한 노이즈 추가 방법, 변화의 크기를 측정하는 방법에 대해서는 지엽적이고 논문의 핵심 내용은 아니므로 생략.

---

Reverse Mutation Paths

기존의 확산 모델 방법론처럼 마르코프 체인의 궤적을 역전시키는 모델을 훈련할 수 있다. 

하지만, 단순히 마지막 돌연변이를 원래대로 반전시키는 훈련은 언뜻 보기에 거친 신호를 생성할 수 있다.

 

예를 들어 원본인 Red를 Green으로 변형하고 다시 Red를 도출하는 목표를 가진다고 해보자.

 

기존의 방법을 사용하면 Green을 Blue로, Blue를 Red로 변환하는 방법을 학습한다. 그러나 Green에서 Red로 바로 변환하는 더 나은 방법을 학습할 수도 있는 것 아닌가?

 

트리의 최적 경로를 계산하는 어떤 알고리즘을 통해 트리 쌍의 edit path를 측정할 수 있는데, 더 나은 예를 위해 Green 뒤에 Yellow를 추가하고 생각해 보겠다.

 

목표 트리 (R→B)와 돌연변이 트리(R→B→G→Y) 간의 edit path를 계산하고 거리가 임계값 σ_small 이하인 경우 그 변이를 목표 트리의 변이로 간주하여 Green을 건너뛸 수 있다. (R→B→Y)

 

그림으로 간단히 설명하면,

왼쪽과 같이 start에서 변이를 거듭하여 end에 도착한 뒤 이를 다시 복구하려고 할 때, 

굳이 궤적을 따라 어렵게 올 것이 아니라 더 효율적인 방법이 있으면 건너뛰겠다는 것이다.

 

이러한 방식으로 트리 간의 edit path를 측정하면서 더 나은 훈련 트리를 찾는다.

---

Value Network & Search

위 과정의 데이터를 바탕으로 '이미지만 보고 두 이미지 간의 edit path를 예측할 수 있는 가치 함수'를 훈련할 수 있으며,

 

이를 통해 beam search 추론의 각 반복에서 유망한 노드(목표 이미지와의 edit path가 가장 짧은)만을 확장할 수 있다.

---

Architecture

정석적인 아키텍처 사용: Image encoder, tokenizer, decoder-only transformer

 

편집 시작을 의미하는 <EDIT> 토큰, 편집 지점을 가리키는 <POS x> 토큰을 사용한다.

모델은 편집 위치와 대체 텍스트를 자기회귀적으로 예측한다.

예측 logits을 마스킹하여 문법적으로 유효한 토큰만을 생성할 수 있도록 한다.

σ_small = 2, 노이즈 단계 수 [1, 5].

 

 

 

Experiments

Project page에서 직접 체험(?)해 볼 수 있다.

https://tree-diffusion.github.io/

 

Tree Diffusion model만 사용했을 때는 빨강, 가치 함수를 이용한 beam search를 추가한 버전은 파랑.