기술 요약: DREAMON – 고정된 크기의 캔버스를 넘어선 코드 인필링을 위한 확산 언어 모델

1. 문제 정의

확산 언어 모델(Diffusion Language Models, DLMs)은 비결정적이고 임의의 순서로 생성을 가능하게 하여, 코드 인필링 작업(접두사와 접미사 사이의 콘텐츠 생성)에 이론적으로 이상적인 유연한 생성 방식을 제공하며 자기회귀(autoregressive) 모델의 유망한 대안으로 꼽힙니다. 그러나 현재의 DLM들은 생성 시 사전에 지정된 고정 길이의 마스크 시퀀스가 필요하다는 결정적인 한계점을 가지고 있습니다.

이러한 제약은 "고정된 크기의 캔버스(fixed-size canvas)"라는 병목 현상을 야기합니다. 코드 인필링에서 이상적인 완성 길이는 예제마다 크게 다릅니다. 사전 정의된 마스크 길이가 실제 정답(ground-truth) 솔루션의 길이와 일치하지 않을 경우 성능이 급격히 저하됩니다. 논문에 따르면, Dream-Coder-7B는 마스크 길이와 타겟 길이가 불일치할 때 HumanEval-Infilling 벤치마크에서 평균 **38%**의 성능 하락을 겪습니다. 구체적으로는 다음과 같습니다:

• 마스크가 너무 적은 경우: 모델이 의미 있는 코드를 생성할 충분한 공간을 확보하지 못해 불완전한 솔루션이 생성됩니다.
• 마스크가 너무 많은 경우: 모델이 남는 공간을 채우기 위해 불필요하거나 잘못된 코드 스니펫을 과잉 생성하는 경향이 있습니다.

기존의 해결책들은 특수한 프롬프팅이나 자연스러운 문맥 구조를 해치는 우회 방법을 사용해 왔으나, DLM은 길이를 동적으로 조절할 수 있는 네이티브 메커니즘이 부족했습니다.

2. 방법론: DREAMON

저자들은 아키텍처의 수정 없이 이산 확산 모델(discrete diffusion models)에서 동적이고 가변적인 길이의 생성을 가능하게 하는 프레임워크인 DREAMON을 제안합니다. 핵심 혁신은 확산 과정에 통합된 두 가지 길이 제어 상태(length-control states)의 도입입니다.

2.1 핵심 메커니즘: 특수 상태(Special States)

DREAMON은 어휘 집합(vocabulary)을 두 개의 특수 센티널 토큰으로 확장합니다:

• [expand]: 이 토큰이 예측되면, 이는 결정론적으로 두 개의 [mask] 토큰으로 확장되어 시퀀스 길이를 효과적으로 늘리고 생성을 위한 더 많은 공간을 제공합니다.
• [delete]: 이 토큰이 예측되면, 시퀀스에서 해당 위치를 제거하여 길이를 효과적으로 축소합니다.

2.2 학습 절차

학습 과정은 **데이터 증강(data augmentation)**과 **손실 재가중치(loss re-weighting)**를 포함합니다:

1. 증강된 시퀀스 구축 ($z_0$): 원래의 입력 시퀀스($x_0$)로부터 저자들은 증강된 시퀀스 $z_0$를 구축합니다. 무작위 토큰 구간(spans)이 [expand] 토큰으로 병합되고, [delete] 토큰이 삽입됩니다. 이는 전방 확산 과정(forward diffusion process) 중에 길이 조절의 필요성을 시뮬레이션합니다.
2. 전방 과정 (Forward Process): $z_0$에 표준 마스크 확산이 적용됩니다. 결정적으로, [expand]와 [delete] 상태 모두 전방 노이징 과정 중에 [mask]로 매핑되어 학습 신호에 기여하도록 합니다.
3. 손실 가중치 부여 (Loss Weighting): [delete] 타겟이 [expand] 타겟보다 더 빈번하게 발생하는 심각한 불균형이 존재합니다 (하나의 [delete]는 하나의 토큰을 대체하지만, 하나의 [expand]는 여러 토큰의 병합을 나타내기 때문). [delete] 예측에 모델이 과적합되는 것을 방기하기 위해 손실 가중치 체계가 도입되었습니다. [delete] 예측의 손실 기여도를 낮추어, 시퀀스 내 모든 [delete] 예측의 총 가중치가 단일 표준 [mask] 예측의 가중치와 동일하도록 설정합니다.

2.3 추론 절차

추론 단계에서 모델은 표준 디노이징을 수행하되 동적인 길이 제어를 적용합니다:

• 적응형 궤적 (Adaptive Trajectory): 모델은 모든 마스크된 위치를 예측합니다. 만약 [expand]가 예측되면, 해당 위치는 두 개의 [mask] 토큰으로 교체됩니다. 만약 [delete]가 예측되면, 해당 토큰은 제거됩니다.
• 삭제 브로드캐스팅 (Deletion Broadcasting): 초기 마스크가 타겟보다 훨씬 클 때 효율성을 높이기 위해 저자들은 휴리스틱을 도입합니다: 만약 모델이 [delete]를 예측하면, 그 오른쪽의 모든 후속 토큰들이 (모두 [mask]인 경우) 즉시 제거됩니다. 이는 토큰 단위의 삭제 과정을 길이 예측 동작으로 변환하여 추론 단계를 크게 줄여줍니다.
• 제약 조건: 무한한 성장을 방지하기 위해 최대 출력 길이($L_{max}$)가 강제됩니다.

3. 주요 기여

1. 동적 길이 제어: DREAMON은 아키텍처 변경 없이 오직 두 개의 특수 토큰([expand] 및 [delete])만을 사용하여 마스크 확산 모델에서 네이티브하고 가변적인 길이 생성을 가능하게 한 첫 번째 프레임워크입니다.
2. 최소한의 학습 오버헤드: 이 방법은 학습 목표(보조 상태 및 손실 가중치)에 대한 최소한의 수정만을 요구하며, 기존의 사전 학습된 DLM(예: Dream-Coder-7B, DiffuCoder-7B)에 무시할 만한 추가 연산량(사전 학습 비용의 0.15%)으로 적용할 수 있습니다.
3. 성능 회복: 이 접근 방식은 확산 모델이 자기회귀 모델과 대등한 성능을 달acer 할 수 있게 하여, 고정 길이 병목 현상을 효과적으로 제거합니다.

4. 실험 결과

저자들은 HumanEval-Infilling (단일 행 및 다중 행) 및 SantaCoder-FIM (Python 서브셋)에서 DREAMON을 평가했습니다.

• 베이스라인 비교: 표준 확산 모델(Dream-7B, DiffuCoder-7B, DreamCoder-7B)은 고정 길이 제약으로 인해 자기회귀 베이스라인(예: Qwen2.5-Coder-7B, Deepseek-Coder-6.7B)보다 크게 뒤처졌습니다.
• DREAMON 성능:
  • DREAMON은 확산 베이스라인 대비 **평균 절대 성능 향상 26.4%**를 달성했습니다.
  • HumanEval-Infilling에서 DREAMON이 강화된 모델(예: DiffuCoder-7B + DREAMON)은 92.2 Pass@1 (단일 행) 및 63.1 Pass@1 (다중 행)에 도달하여, Qwen2.5-Coder-7B와 같은 최상위 자기회귀 모델을 능가하거나 대등한 수준에 도달했습니다.
  • DREAMON 변형 모델들은 오라클 수준에 근접한 성능(초기 마스크 길이가 정답과 일치할 때의 성능)을 달성하여, 모델이 초기 마스크 크기에 관계없이 올바른 길이에 성공적으로 적응할 수 있음을 입증했습니다.
• 절제 연구 (Ablation Studies):
  • [expand] 또는 [delete] 중 하나라도 제거하면 성능이 크게 하락하여, 두 메커니즘이 모두 필수적임을 확인했습니다.
  • **손실 균형(Loss balancing)**은 매우 중요했습니다. 이 기능이 없으면 삭제에 대한 과적합으로 인해 성능이 평균 약 7% 하락했습니다.
  • 삭제 브로드캐스팅은 품질 저하 없이 전체 추론 단계를 약 50% 감소시켰습니다 (예: 마스크 길이가 64일 때 122.8단계에서 52.4단계로 감소).

5. 의의 및 주장

본 논문은 DREAMON이 근본적인 장벽을 제거한다고 주장합니다. 동적인 길이 조정을 가능하게 함으로써, DREAMON은 특히 코드 인필링과 같이 가변 길이 생성 작업에서 DLM의 유연성과 적용 가능성을 크게 진전시켰습니다.

저자들은 자신들의 연구가 출력 길이가 알려지지 않은 실제 시나리오에서 DLM을 자기회귀 모델과 경쟁할 수 있게 만드는 단계라고 위치시킵니다. 현재 설계는 단일 확장 계수(1-to-2)를 사용하지만, 향-후 연구에서는 효율성을 더욱 최적화하기 위해 더 풍부한 확장 어휘 집합이나 명시적인 길이 예측 헤드를 탐색할 수 있다고 언급합니다. 이 연구는 이러한 이득이 기존 트랜스포머의 아키텍처 변경 없이 달성되었으며, 이를 통해 원래의 확산 프레임워크가 가진 단순성을 보존했음을 강조합니다.