[논문 리뷰] Energy-Based Models for Code Generation under Compilability Constraints

Tomasz Korbak, Hady Elsahar외 2명|arXiv (Cornell University)|2021. 06. 09.

Topic Modeling참고 문헌 43인용 수 5

한 줄 요약

이 논문은 자동회귀 모델이 문법적으로 유효한 코드만 생성하도록 제약을 두어 컴파일 가능성을 강제하는 에너지 기반 모델(EBM)을 제안한다. KL-적응형 분포 정책 강화 학습(KL-Adaptive Distributional Policy Gradient)을 사용하여 생성 모델을 EBM에 근접하게 미세조정함으로써, 다양성이나 복잡성에 손상 없이 컴파일 가능성 비율을 크게 향상시킨다. 이는 전통적인 강화학습 미세조정 방식이 치명적인 기억 상실(catastrophic forgetting)을 유발하는 것과는 대조된다.

ABSTRACT

Neural language models can be successfully trained on source code, leading to applications such as code completion. However, their versatile autoregressive self-supervision objective overlooks important global sequence-level features that are present in the data such as syntactic correctness or compilability. In this work, we pose the problem of learning to generate compilable code as constraint satisfaction. We define an Energy-Based Model (EBM) representing a pre-trained generative model with an imposed constraint of generating only compilable sequences. We then use the KL-Adaptive Distributional Policy Gradient algorithm (Khalifa et al., 2021) to train a generative model approximating the EBM. We conduct experiments showing that our proposed approach is able to improve compilability rates without sacrificing diversity and complexity of the generated samples.

연구 동기 및 목표

• 자기회귀 신경 코드 모델이 시각적 제약과 전역적 시퀀스 제약의 부족으로 인해 문법적으로 올바른 컴파일 가능한 코드를 생성하지 못하는 한계를 해결하기 위해.
• 컴파일 가능성 제약 조건 하에서 코드 생성을 제약 만족 문제로 재정의하여 체계적인 분포 기반 접근을 가능하게 하기 위해.
• 원래 모델의 다양성과 복잡성을 유지하면서 제약 조건이 부여된 EBM 분포를 근사하는 생성 모델을 훈련하기 위해.
• 문장 수준의 보상에 대해 강화학습 미세조정에서 흔히 발생하는 치명적인 기억 상실 문제를 해결하기 위해.
• 생성된 코드의 유창성이나 구조적 복잡성의 열화 없이 컴파일 가능성만 향상시키는 방법을 제공하기 위해.

제안 방법

• 컴파일 가능한 코드 시퀀스에만 낮은 에너지(높은 확률)를 할당하는 에너지 기반 모델(EBM)을 정의하여, 컴파일 가능성 제약 조건을 전역적 특성으로 표현하기 위해.
• 자기회귀 생성 모델이 EBM 분포를 근사하도록 KL-적응형 분포 정책 강화 학습(KL-DPG) 알고리즘을 사용하여 훈련하기 위해.
• 원래 미리 훈련된 언어 모델을 사전으로 활용하여, 미세조정된 정책와 원래 모델 간의 KL 발산을 최소화하기 위해.
• EBM 제약 조건을 미분 가능 대체 함수를 통해 훈련 중에 적용하여, 정책 강화 방법을 사용한 엔드 투 엔드 최적화를 가능하게 하기 위해.
• 컴파일 성공 여부 피드백을 보상 신호로 활용하여, 정책가 유효하고 컴파일 가능한 코드 시퀀스만 생성하도록 이끌기 위해.
• EBM를 강화학습 프레임워크에 통합하여 제약 조건 이행과 분포 충실도 사이의 균형을 이루기 위해.

실험 결과

연구 질문

• RQ1 에너지 기반 모델이 코드 생성에서 컴파일 가능성을 전역적 제약 조건으로 효과적으로 표현할 수 있는가?
• RQ2 KL-적응형 분포 정책 강화 학습 알고리즘이 다양성을 유지하면서 EBM를 근사하는 생성 모델을 성공적으로 훈련시킬 수 있는가?
• RQ3 기존의 RL 미세조정 방식과 비교해 볼 때, EBM를 통한 컴파일 가능성 강제화가 기억 상실 위험을 줄이는가?
• RQ4 제안된 방법은 생성된 코드의 복잡성이나 유창성의 열화 없이 컴파일 가능성 비율을 얼마나 향상시키는가?
• RQ5 다양한 미세조정 전략(예: $R(x)=b(x)$, $R(x)=P(x)$, 또는 KL-DPG)이 컴파일 에러 분포에 어떤 영향을 미치는가?

주요 결과

• 제안된 EBM 기반 방법은 표준 자기회귀 모델 및 강화학습 미세조정 기준선 대비 컴파일 가능성 비율을 크게 향상시킨다.
• 이 방법은 생성된 코드의 높은 다양성과 복잡성을 유지하며, 기존의 강화학습 미세조정 방식에서 관찰되는 심각한 열화 현상은 피한다.
• 표준 RL 방법보다 KL-DPG 미세조정이 원래 모델과 더 나은 분포 일치를 달성하여 치명적인 기억 상실을 감소시킨다.
• 전역적 제약 조건을 강제함으로써 장거리 문법적 의존성에서의 컴파일 에러를 효과적으로 줄인다.
• 정성적 분석 결과, EBM 기반 방법은 기준선 대비 더 적고 국소화된 컴파일 에러를 생성한다.
• 이 방법은 히ュ리스틱 보상이 적용된 표준 RL과 단지 컴파일 정확도를 최대화하는 정책 강화 방법 모두를 뛰어넘는 성능을 보인다.