Diffusion On Syntax Trees For Program Synthesis

https://arxiv.org/abs/2405.20519

Diffusion On Syntax Trees For Program Synthesis

 

대형 언어 모델(LLM)은 코드를 하나의 토큰씩 생성합니다. 이들의 자기회귀 생성 과정은 프로그램의 출력을 관찰하는 피드백이 부족합니다. LLM을 직접 수정 제안을 할 수 있도록 훈련시키는 것은 풍부한 수정 데이터의 희소성 때문에 어려울 수 있습니다. 이러한 문제를 해결하기 위해, 우리는 임의의 문맥 자유 문법의 구문 트리에서 작동하는 신경 확산 모델을 제안합니다. 이미지 확산 모델과 유사하게, 우리의 방법도 구문 트리에 적용된 "노이즈"를 반전시킵니다. 코드를 순차적으로 생성하는 대신, 우리는 구문적 유효성을 유지하면서 반복적으로 코드를 수정합니다. 이는 이 신경 모델을 검색과 결합하기 쉽게 만듭니다. 우리는 이 접근 방식을 역 그래픽 작업에 적용하여 모델이 이미지를 해당 이미지를 생성하는 프로그램으로 변환하는 방법을 학습하도록 합니다. 검색과 결합된 우리의 모델은 그래픽 프로그램을 작성하고, 실행 결과를 확인하며, 요구된 사양을 충족하기 위해 디버깅할 수 있습니다. 추가적으로, 우리의 시스템이 손으로 그린 스케치를 위한 그래픽 프로그램을 작성하는 방법을 보여줍니다. 비디오 결과는 https://tree-diffusion.github.io에서 확인할 수 있습니다.

 

1. 서론

대형 언어 모델(LLM)은 코드 생성에서 눈부신 발전을 이루었지만, 자기회귀적 특성은 근본적인 도전 과제를 제시합니다: 이들은 프로그램의 런타임 출력을 사용하지 않고 하나의 토큰씩 코드를 생성합니다. 이로 인해 오류를 수정하기 어렵고, 모델은 프로그램의 출력을 보고 이에 따라 조정하는 피드백 루프가 부족합니다. LLM이 기존 코드를 수정하는 제안을 할 수 있도록 훈련할 수 있지만 [6, 42, 17], 이를 위한 충분한 훈련 데이터를 확보하기는 어렵습니다.

이 논문에서는 구문 트리에서 직접 작동하는 신경 확산 모델을 사용한 새로운 프로그램 합성 접근 방식을 소개합니다. 확산 모델은 이전에 이미지 생성에서 큰 성공을 거두었습니다 [14, 22, 31]. 확산을 활용하여, 우리는 모델이 구문적 유효성을 보장하면서 프로그램을 반복적으로 정제하도록 학습시킵니다. 중요한 것은, 우리의 접근 방식이 모델이 각 단계에서 프로그램의 출력을 관찰할 수 있도록 하여 효과적으로 디버깅 과정을 가능하게 한다는 점입니다.

AlphaZero [29]와 같은 시스템의 정신을 이어받아, 확산의 반복적 특성은 검색 기반 프로그램 합성에 자연스럽게 적합합니다. 확산 모델과 함께 가치 모델을 훈련시켜 디노이징 과정을 원하는 출출을 달성할 가능성이 높은 프로그램으로 유도할 수 있습니다. 이를 통해 프로그램 공간을 효율적으로 탐색하고, 생성 과정의 각 단계에서 더 정보에 입각한 결정을 내릴 수 있습니다.

우리는 이미지를 그리는 도메인 특화 언어를 가정하여 역 그래픽 작업에 우리의 접근 방식을 구현합니다. 역 그래픽 작업은 코드의 작은 변경이 렌더링된 이미지에 의미 있는 변화를 일으키기 때문에 우리의 접근 방식에 자연스럽게 적합합니다. 예를 들어, 이미지에서 잘못된 위치에 있는 모양은 프로그램 공간에서 쉽게 확인하고 수정할 수 있습니다.

우리의 주요 기여는 (a) 구문 트리에 확산을 사용한 프로그램 합성에 대한 새로운 접근 방식과 (b) 이전 방법을 크게 능가하는 역 그래픽 작업에 대한 우리의 접근 방식의 구현입니다.

 

2.배경 및 관련 연구

신경 프로그램 합성

신경 프로그램 합성은 입력-출력 예제에서 프로그램을 생성하는 신경망을 연구하는 중요한 분야입니다. Parisotto et al. [23]와 같은 초기 연구는 이 접근 방식의 실현 가능성을 보여주었습니다. 현대 언어 모델은 프로그램 합성에 직접 적용될 수 있지만, 신경망과 검색 전략을 결합하면 더 나은 결과와 보장을 얻을 수 있습니다. 이 패러다임에서 신경망은 제안 분포를 제공하거나 후보 프로그램을 평가하여 검색 과정을 안내합니다. 이러한 하이브리드 방법의 예로는 Balog et al. [2], Ellis et al. [12], Devlin et al. [9]이 있습니다. 우리의 작업과의 주요 차이점은 이러한 방법들이 부분 프로그램의 방대한 공간을 탐색하면서 프로그램을 점진적으로 구성한다는 것입니다. 반면, 우리의 접근 방식은 프로그램을 수정하는 데 중점을 두어, 프로그램을 처음부터 성장시키고 프로그램 실행에 따라 수정할 수 있습니다.

신경 확산

신경 확산 모델은 이미지와 같은 고차원 데이터를 모델링하는 데 인상적인 결과를 보여주었습니다 [14, 22, 31]. 신경 확산 모델은 데이터 분포에서 샘플(예: 실제 이미지)을 취하여 노이즈를 추가하여 데이터를 점진적으로 손상시키고, 신경망을 훈련시켜 노이즈를 점진적으로 제거합니다. 새로운 샘플을 생성하기 위해서는 무작위 노이즈에서 시작하여 입력을 디노이즈하기 위해 신경망을 반복적으로 적용합니다.

이산 데이터에 대한 확산

최근 연구는 그래프와 같은 이산 및 구조화된 데이터에 확산을 확장하고 있습니다 [35]. 이러한 연구는 분자 설계 [15, 27, 8] 등의 분야에 적용됩니다. Lou et al. [20]는 언어 모델링을 위해 새로운 스코어 매칭 목적을 사용하는 이산 확산 모델을 제안했습니다. 구조화된 데이터에 대한 생성 모델링의 또 다른 유망한 연구는 생성 흐름 네트워크(GFlowNets) [3]입니다. 이 모델은 신경 모델이 구조화된 데이터를 하나의 원자씩 구성합니다.

코드 생성에 대한 확산

Singh et al. [30]는 코드 생성을 위해 확산 모델을 사용합니다. 그러나 이 접근 방식은 먼저 텍스트를 연속 잠재 공간에 임베딩하고, 그 공간에서 연속 확산 모델을 훈련한 다음, 마지막에 이를 임베딩 해제하는 것입니다. 이는 잠재 표현의 중간 단계가 실제 코드와 대응하도록 훈련되지 않음을 의미합니다. 임베딩 토큰은 마지막 몇 단계에서 가장 가까운 임베딩에 고정됩니다. 신경 모델을 사용한 직접적인 코드 편집도 연구되었습니다. Chakraborty et al. [6]는 실제 코드 패치 데이터셋으로 훈련된 그래프 신경망을 사용하여 코드 편집을 수행합니다. 유사하게, Zhang et al. [42]는 언어 모델을 훈련시켜 [MASK] 토큰을 수정하거나 삽입하거나 기존 토큰을 삭제하여 코드를 편집합니다. 이들은 실제 댓글과 패치에 모델을 추가로 미세 조정합니다. 이러한 방법과 달리, 우리의 접근 방식은 광범위한 코드 편집 데이터셋의 필요성을 피하고, 우리의 사전 훈련 작업을 통해 구문적 유효성을 본질적으로 보장합니다.

역 그래픽을 위한 프로그램 합성

우리는 Sharma et al. [28], Ellis et al. [10, 11]의 이전 작업에서 영감을 받았습니다. 이들은 CSG2D 언어를 사용합니다. Sharma et al. [28]은 이미지를 프로그램으로 변환하기 위해 컨볼루션 인코더와 순환 모델을 제안합니다. Ellis et al. [11]은 중간 프로그램 실행 출력을 읽고 평가하고 출력하는 루프(REPL)에서 신경 모델에 제공하는 방법을 제안합니다. 우리의 방법과 달리, 부분 그래픽 프로그램을 실행할 수 있는 능력은 그들의 작업에서 중요한 요구 사항입니다. 우리의 시스템은 완전한 프로그램을 대상으로 하며, 맞춤형 부분 컴파일러가 필요하지 않습니다. 그들의 연구에서 언급했듯이, 그들의 정책은 또한 취약합니다. 정책이 객체를 제안하면 그 제안을 취소할 수 없습니다. 따라서 이러한 시스템은 실수에 덜 취약하게 만들기 위해 순차적 몬테카를로(SMC) 샘플러에서 많은 수의 입자가 필요합니다.

Figure 1: 우리 시스템이 복구한 프로그램의 예시. 첫 번째 줄은 손으로 그린 아이콘 스케치(왼쪽), 복구된 프로그램(중간), 복구된 프로그램의 컴파일 결과(오른쪽)를 보여줍니다. 상위 두 줄은 구성적 솔리드 기하학 언어(CSG2D-Sketch)에 해당합니다. 마지막 줄은 도형과 색상의 계층적 프로그램을 반전시키는 TinySVG 환경의 출력 예시입니다. 비디오 예시는 https://tree-diffusion.github.io에서 확인할 수 있습니다.

 

 

3. 방법

우리 방법의 주요 아이디어는 구문 트리에 대해 이미지 확산 모델과 유사한 형태의 디노이징 확산 모델을 개발하는 것입니다.

Ellis et al. [11]의 예시 과제를 고려해보면, 이는 이미지에서 구성적 솔리드 기하학(CSG2D) 프로그램을 생성하는 것입니다. CSG2D에서는 원과 사각형 같은 단순한 기본 도형을 덧셈과 뺄셈 같은 불리언 연산을 사용해 결합하여 더 복잡한 도형을 만들 수 있습니다. 이때 문맥 자유 문법(CFG)을 사용합니다.

이 방법을 통해 우리는 구문 트리를 기반으로 프로그램을 생성하고, 이 과정에서 구문적 유효성을 유지하면서 프로그램을 점진적으로 수정합니다. 이는 이미지를 프로그램으로 변환할 때 프로그램의 중간 결과를 디버깅하고 수정할 수 있는 능력을 제공합니다. 우리의 접근 방식은 이미지에서 추출된 프로그램이 구문 트리 형태로 표현되며, 이 구문 트리에 노이즈를 추가하고, 신경망을 사용하여 노이즈를 점진적으로 제거하면서 프로그램을 정제합니다.

 

 

3.1 작은 변이 샘플링

3.2 정책

3.2.1 순방향 과정

3.2.2 변이 경로 반전

3.3 가치 네트워크 및 검색

3.4 아키텍처

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

-----------------------------------------------------------------------------------------------------------------------------------------------------------------

 

4. 실험

4.1 환경

우리는 네 가지 도메인 특화 그래픽 언어에 대해 실험을 수행합니다. 완전한 문법 명세는 부록 B에 제공됩니다.

• CSG2D: 2D 구성적 솔리드 기하학 언어로, 기본 도형을 추가하고 빼서 더 복잡한 형태를 만듭니다. 이는 우리의 기준 방법 [11, 28]에서 탐구한 것입니다. 우리는 또한 CSG2D-Sketch를 만들어, Wood et al. [39]의 알고리즘을 사용해 손으로 그린 스케치를 시뮬레이션하는 관찰 모델을 추가했습니다.
• TinySVG: 색상이 있는 기본 도형과 수평 및 수직 정렬을 위한 Arrange 명령, 도형 오프셋을 위한 Move 명령을 특징으로 하는 언어입니다. 그림 1은 예제 프로그램을 보여줍니다. CSG2D의 구성적 특성과 달리, TinySVG는 계층적입니다. 즉, 하위 표현식을 결합해 고수준 조작을 위한 복합 객체로 만들 수 있습니다. 우리는 또한 Move 명령이 없는 TinySVG의 단순화 버전인 Rainbow를 만들어 계산 요구 사항이 줄어드는 절제 연구를 수행했습니다.
• 구현: 우리는 Lark [19]와 Iceberg [16] 파이썬 라이브러리를 사용하여 이러한 언어를 구현했으며, 우리의 트리-확산 구현은 어떤 문맥 자유 문법 및 관찰 모델에도 적응할 수 있도록 설계되었습니다.

4.2 기준 방법

우리는 Ellis et al. [11]과 CSGNet [28]의 두 가지 이전 작업을 기준 방법으로 사용합니다.

• CSGNet: Sharma et al. [28]은 입력 이미지에서 프로그램 문을 생성하기 위해 컨볼루션 및 순환 신경망을 사용했습니다. 공정한 비교를 위해, 우리는 현대적 자기회귀 코드 생성 접근 방식을 나타내는 동일한 비전-언어 트랜스포머 아키텍처를 사용하여 CSGNet을 다시 구현했습니다. 우리는 거절 샘플링을 사용하여, 일치하는 프로그램이 발견될 때까지 반복적으로 프로그램을 생성합니다.

그림 4: CSG2D 및 TinySVG 언어에서 기준 방법과 비교한 우리의 접근 방식의 성능을 보여줍니다. 우리는 테스트 세트에서 n = 256개의 이미지를 사용하여 솔루션을 찾기 위해 확장된 노드 수를 측정합니다. 자기회귀 기준 방법은 거절 샘플링으로 쿼리되었습니다. 우리의 정책은 이전 방법보다 우수하며, 검색과 결합된 우리의 정책은 성능을 더욱 향상시킵니다. 오류 막대는 5개의 랜덤 시드에 걸친 표준 편차를 나타냅니다.

 

REPL 흐름

Ellis et al. [11]은 모든 기본 요소가 배치될 때까지 한 번에 하나씩 프로그램을 구축하는 방법을 제안했습니다. 그들은 정책 네트워크에 REPL 접근 권한, 즉 코드를 실행하고 출력을 확인할 수 있는 능력을 제공합니다. 특히, 현재 이미지는 현재 부분 프로그램에서 렌더링됩니다. 따라서 사용자 정의 부분 컴파일러가 필요합니다. CSG2D의 경우 이는 구성적 언어이기 때문에 간단합니다. 단순히 지금까지 배치된 도형들을 렌더링합니다. TinySVG의 경우, 부분 컴파일러를 어떻게 작성해야 할지 즉시 명확하지 않습니다. 이는 렌더링이 하향식으로 발생하기 때문입니다. 기본 요소들이 배치되고, 그 배치가 다시 배치됩니다(그림 1 참조). 따라서 우리는 이 기준 방법을 CSG2D에서만 사용합니다. 생성 흐름 네트워크(Generative Flow Networks) [3]와의 유사성 때문에, 우리는 수정된 방법을 "REPL Flow"라고 부릅니다.

테스트 작업

TinySVG의 경우 무작위로 생성된 표현식과 해당 이미지로 구성된 홀드아웃 테스트 세트를 사용했습니다. CSG2D 작업의 경우, 모든 방법이 인-분포 홀드아웃 테스트 세트에서 최대 성능에 도달한 것을 확인했습니다. Ellis et al. [11]에서는 저자들이 더 많은 객체가 포함된 더 어려운 테스트 세트를 만들었습니다. 그러나 CSG2D와 같은 환경에서 단순히 더 많은 객체를 추가하면, 장면의 큰 부분을 빼는 큰 객체를 샘플링할 가능성이 높아져 더 간단한 최종 장면이 됩니다. 대신, 어려운 테스트 세트를 생성하기 위해, 우리는 LZ4 [7, 37] 압축 알고리즘을 사용하여 비압축성의 95번째 백분위수 이상의 이미지를 필터링했습니다.

 

평가

.

그림 5: 시스템의 여러 설계 결정 변경 효과. 우리는 Rainbow 환경에서 더 작은 모델을 훈련시킵니다. 모델에게 n = 256개의 테스트 문제를 해결하도록 합니다. (a)에서, 'No Reverse Path'의 경우, 명시적인 역 경로를 계산하지 않고, 노이징 과정의 마지막 단계만을 목표로 사용하여 모델을 훈련시킵니다. 'No Current Image'의 경우, 모델이 수정 중인 프로그램의 컴파일된 출력 이미지를 볼 수 없도록 훈련시킵니다. 'No Noising'의 경우, 우리의 노이징 과정을 사용하지 않고, 두 개의 무작위 표현을 생성하고 그 사이의 경로를 목표로 사용합니다. (b)에서, 우리는 순방향 확산(ρ = 0.0)과 순수 무작위 초기화(ρ = 1.0) 사이의 훈련 혼합 효과를 추가로 조사합니다. 오류 막대는 5개의 랜덤 시드에 걸친 표준 편차를 나타냅니다.

 

4.3 절제 연구

우리의 설계 결정의 영향을 이해하기 위해, 간소화된 Rainbow 환경에서 더 작은 트랜스포머 모델을 사용하여 절제 연구를 수행했습니다.

먼저, 정책 네트워크 입력에서 현재 이미지를 제거하는(No REPL) 효과를 조사했습니다. 그림 5(a)에 나타난 바와 같이, 이는 성능에 큰 저해를 일으켜, Ellis et al. [11]이 관찰한 REPL 유사 인터페이스의 중요성을 확인시켜 주었습니다.

다음으로, 우리의 역 변이 경로 알고리즘의 필요성을 조사했습니다. 마지막 변이 단계만을 학습하는 것은 유효한 경로를 제공하지만, 비최적의 중간 상태를 대상으로 하여 노이즈를 도입할 수 있습니다. 그림 5(a)는 역 변이 경로를 이용하는 것이 성능을 크게 향상시키며, 특히 적은 단계로 해결책을 찾는 데 유리함을 보여줍니다. 그러나 두 방법 모두 결국 유사한 성능 수준에 도달하여, 노이즈가 많은 경로가 덜 효율적이더라도 해결책에 도달할 수 있음을 시사합니다.

마지막으로, 트리 편집 경로 알고리즘을 고려할 때, 점진적인 노이즈 과정이 중요한지 탐구했습니다. 두 무작위 표현식을 직접 샘플링하고 경로를 계산하여 네트워크가 이를 모방하도록 학습시킬 수는 없을까요? 그림 5(b)에 나타난 바와 같이, 우리는 순방향 확산(ρ = 0.0)과 순수 무작위 초기화(ρ = 1.0) 사이의 훈련 데이터 구성을 다양하게 하였습니다. 소량(ρ = 0.2)의 순수 무작위 초기화와 순방향 확산을 결합한 것이 최고의 결과를 낳았습니다. 이는 순방향 확산이 목표 지점 주변의 더 풍부한 훈련 분포를 제공하는 반면, 무작위 초기화는 모델이 프로그램 공간을 더 광범위하게 탐색하도록 가르침을 시사합니다. 순방향 확산에서의 세밀한 수정이 우리의 평가에서 정확한 픽셀 매치를 달성하는 데 유익하다는 점을 증명합니다.

5. 결론

이 연구에서, 우리는 구문 트리를 기반으로 프로그램 합성을 수행하는 신경 확산 모델을 제안했습니다. 우리는 주어진 이미지를 렌더링할 프로그램을 찾는 역 그래픽 작업에 우리의 접근 방식을 구현했습니다. 이전 연구와 달리, 우리의 모델은 프로그램을 구성하고, 그 출력을 확인하며, 피드백 루프에서 실수를 수정할 수 있도록 프로그램을 수정할 수 있습니다. 우리는 이러한 역 그래픽 작업에서 우리의 접근 방식이 기준 방법보다 뛰어남을 정량적으로 보여주었습니다. 또한, 절제 실험을 통해 주요 설계 결정의 영향을 연구했습니다.

 

Figure 6: 두 역 그래픽 언어인 CSG2D (상단 두 줄)와 TinySVG (하단 두 줄)에서 우리의 방법과 기준 방법의 정성적 예시입니다. 가장 왼쪽 열은 테스트 세트에서 실제 렌더링된 프로그램을 보여줍니다. 다음 열은 다양한 방법으로 렌더링된 프로그램을 보여줍니다. 우리의 방법은 실제 프로그램과 더 밀접하게 일치하도록 세밀하게 조정할 수 있습니다.

 

Figure 7: CSG2D-Sketch 언어에서 입력 스케치에 대해 복구된 프로그램의 예시입니다. 입력 스케치는 손으로 그린 스케치를 시뮬레이션하는 우리의 관찰 모델에서 가져온 것입니다(상단 행). 출력 프로그램이 렌더링된 결과(하단 행)는 기본 도형을 추가하고 빼서 입력 스케치와 일치할 수 있습니다. 이러한 스케치에 대한 비디오 결과는 https://tree-diffusion.github.io에서 확인할 수 있습니다.

 

한계점

이 작업에는 몇 가지 중요한 한계가 있습니다. 첫째, 우리는 변수 바인딩, 루프, 문자열, 연속 매개변수 등이 없는 표현식에서 작동합니다. 우리의 접근 방식이 이를 지원하도록 확장될 수 있다고 생각하지만, 더 많은 작업과 신중한 설계가 필요합니다. 현재의 대형 언어 모델은 많은 도메인에서 복잡한 프로그램을 작성할 수 있지만, 우리는 매우 좁은 작업에 집중하고 있습니다. 또한, 역 그래픽 작업은 이미지 출력에 정보성 있는 변화를 만드는 작은 변이에 특히 적합할 수 있습니다.

미래 작업

미래에는 프로그램에 대한 대규모 인터넷 데이터를 활용하여 우리의 시스템을 훈련시키고, 구문 트리에 작은 변이를 가하여 이를 반전시키는 방법을 배우고자 합니다. 우리는 또한 역 그래픽 이외의 도메인에서 이 접근 방식을 연구하고 싶습니다. 추가적으로, 이 접근 방식을 이산 구문 구조와 연속 부동 소수점 상수를 모두 처리하도록 확장하고자 합니다.

영향

구현의 범위가 좁기 때문에, 기계 보조 프로그래밍의 미래 연구 방향에 정보를 제공하는 것 외에는 직접적인 사회적 영향은 없다고 생각합니다. 이 작업의 미래 방향이 특히 역 그래픽 분야에서 아티스트, CAD 모델러, 프로그래머가 아이디어를 정확한 프로그램으로 신속하게 변환할 수 있는 도구를 제공하는 데 도움이 되기를 바랍니다.

감사 및 자금 지원 공개

논의, 피드백 및 기술 지원에 대해 Kathy Jang, David Wu, Cam Allen, Sam Toyer, Eli Bronstein, Koushik Sen, Pieter Abbeel에게 감사드립니다. Shreyas는 부분적으로 Schmidt Futures의 AI2050 프로그램(Grant G-22-63471)의 지원을 받았습니다. Erik은 Future of Life Institute와 Open Philanthropy의 펠로십으로 지원받았습니다.

 

2405.20519v1.pdf

5.75MB