Unlocking Prompt Infilling Capability for Diffusion Language Models
이 논문은 현재 확산 언어 모델 (dLM) 의 프롬프트 인필링 능력을 제한하는 아키텍처적 한계가 아닌, 응답만 마스킹하는 기존 SFT 관행에 기인함을 밝히고, 프롬프트와 응답을 함께 마스킹하는 전 시퀀스 마스킹 방식을 도입하여 모델이 직접 인필링된 프롬프트를 생성하도록 함으로써 기존 수동 템플릿을 능가하는 성능을 달성했다고 요약할 수 있습니다.
새로운 AI 모델 (확산 언어 모델) 은 마치 모든 재료를 한 번에 보고 요리를 완성하는 천재 요리사와 같습니다. 기존 AI(자동 생성 모델) 가 "앞에서부터 하나씩 재료를 추가해가며" 요리를 한다면, 이 새로운 AI 는 "완성된 요리를 보고 어떤 재료가 빠졌는지"를 양쪽에서 동시에 추론할 수 있는 능력이 있습니다.
하지만, 문제는 교육 방식에 있었습니다. 지금까지 이 AI 를 가르칠 때, 레시피 (질문/프롬프트) 는 그대로 두고 오직 요리 결과물 (답변) 만 가려서 "이 가려진 부분을 채워줘"라고만 시켰습니다.
결과: AI 는 "완성된 레시피를 보고 요리를 만드는 법"은 잘 배웠지만, "요리 결과물을 보고 레시피를 다시 쓰는 법"은 전혀 배우지 못했습니다.
현실: AI 가 "이 질문을 이렇게 바꾸면 더 좋은 답이 나올 텐데"라고 스스로 레시피를 수정하는 능력 (프롬프트 인필링) 이 잠겨 있었던 것입니다.
2. 해결책: "레시피와 요리 모두 가려서" 가르치기
저자들은 이 잠겨 있는 능력을 해제하기 위해 아주 간단한 훈련 방식을 바꿨습니다.
새로운 훈련 방식: 질문 (레시피) 과 답변 (요리) 모두를 가리고, AI 가 빈칸을 스스로 채우게 합니다.
효과: 이제 AI 는 "완성된 요리를 보고, 원래 레시피가 어떻게 되어야 했는지"를 추론할 수 있게 되었습니다. 마치 요리사가 "이 요리를 맛있게 만들려면 레시피의 '소금' 부분을 '설탕'으로 바꿔야겠다"라고 스스로 깨닫는 것과 같습니다.
3. 실제 효과: "스스로 만든 레시피"가 더 맛있다
이제 AI 가 스스로 레시피를 채워 넣는 (인필링) 능력을 테스트해 보니 놀라운 결과가 나왔습니다.
사람이 만든 레시피보다 낫다: 인간 전문가가 수동으로 고안한 복잡한 질문 (프롬프트) 보다, AI 가 몇 가지 예시를 보고 스스로 채운 질문이 더 정확한 답을 이끌어냈습니다.
다른 모델에게도 통한다: A 모델이 스스로 고안한 레시피를 B 모델에게 주면, B 모델도 더 잘 요리할 수 있었습니다. (레시피가 모델 간에 호환된다는 뜻입니다.)
기존 기술과 잘 섞인다: 이미 다른 방법으로 최적화된 레시피에 이 기술을 적용하면, 요리 실력이 한 단계 더 업그레이드되었습니다.
💡 핵심 요약
이 논문의 결론은 매우 명확합니다.
"AI 가 레시피를 스스로 고칠 수 없는 것은 머리 (아키텍처) 가 부족해서가 아니라, 교육 방식 (훈련 데이터) 이 잘못되었기 때문입니다."
우리가 AI 에게 질문과 답변을 모두 가리는 훈련을 시키기만 하면, AI 는 스스로 더 좋은 질문을 만들어내는 놀라운 능력을 발휘할 수 있습니다. 이는 AI 개발자들에게 "아키텍처를 복잡하게 바꿀 필요 없이, 훈련 방식을 조금만 바꾸면 성능을 획기적으로 높일 수 있다"는 중요한 교훈을 줍니다.
한 줄 요약: "AI 가 스스로 더 좋은 질문을 만들 수 있게 하려면, 질문과 답변을 동시에 가려서 훈련시켜야 합니다. 그래야 AI 는 레시피를 스스로 고쳐서 더 맛있는 요리를 할 수 있습니다."
1. 문제 정의 (Problem)
**마스크된 확산 언어 모델 (Masked Diffusion Language Models, dLMs)**은 기존 autoregressive(자기회귀) 모델과 달리 양방향 (bidirectional) 디노이징을 통해 시퀀스 내 임의의 위치를 채울 수 있는 잠재력을 가지고 있습니다. 그러나 현재 널리 사용되는 dLM(예: LLaDA, Dream) 은 **Supervised Fine-Tuning (SFT) 단계에서 '응답만 마스킹 (Response-only Masking)'**하는 관행을 따르고 있습니다.
현재의 한계: SFT 과정에서 프롬프트 (입력) 는 완전히 노출된 상태로 유지되고, 오직 응답 (출력) 만이 마스킹되어 학습됩니다.
결과: 모델은 "깨끗한 프롬프트가 주어졌을 때 응답을 생성"하는 법은 배우지만, "마스킹된 프롬프트를 채워 넣는 (Prompt Infilling)" 능력을 전혀 학습하지 못합니다.
핵심 통찰: 이는 모델의 아키텍처적 한계가 아니라, 학습 데이터의 마스킹 전략 (Training Signal) 부재로 인한 문제입니다.
2. 방법론 (Methodology)
저자들은 이 격차를 해소하기 위해 SFT 단계의 마스킹 전략을 단순하게 변경하는 **전체 시퀀스 마스킹 (Full-Sequence Masking, FS)**을 제안합니다.
가. 전체 시퀀스 마스킹 SFT (Full-Sequence Masking SFT)
방식: 기존에는 응답만 마스킹했으나, 제안된 방법은 프롬프트 (P) 와 응답 (R) 을 포함한 전체 시퀀스 X=(P,R)을 마스킹합니다.
학습 과정: 각 학습 예제에서 마스킹 비율 t를 샘플링하고, 시퀀스 내 임의의 토큰들을 마스킹합니다. 이는 프롬프트 전용, 응답 전용, 그리고 전체 시퀀스 마스킹 등 다양한 패턴을 모델이 경험하도록 합니다.
손실 함수: 모든 마스킹된 토큰에 대해 교차 엔트로피 손실 (Cross-Entropy Loss) 을 계산하여 모델이 프롬프트와 응답 모두를 복원 (디노이징) 하도록 학습시킵니다.
나. 선택적 정제 단계 (Optional Refinement)
전체 시퀀스 마스킹 학습 후, 일부 데이터셋에서 성능을 더 높이기 위해 응답 전용 마스킹 (Response-Only, RO) 단계로 추가 파인튜닝을 수행할 수 있습니다. 이는 표준 생성 작업을 위한 특화 능력을 유지하면서도 프롬프트 인필링 능력을 보존합니다.
다. 추론 시 프롬프트 인필링 (Prompt Infilling at Inference)
프로세스:
Few-shot 예시 (질문과 정답) 와 마스킹된 프롬프트 템플릿을 입력으로 받습니다.
학습된 dLM 이 마스킹된 프롬프트 부분을 디노이징하여 최적의 프롬프트를 생성합니다.
생성된 여러 프롬프트 후보를 검증 (Validation) 하여 가장 성능이 좋은 프롬프트를 선택합니다.
선택된 프롬프트를 고정하여 모든 테스트 입력에 적용합니다.
3. 주요 기여 (Key Contributions)
학습 - 추론 격차의 규명: 현재 SFT 관행인 '응답 전용 마스킹'이 프롬프트 인필링 능력을 차단하는 근본적인 원인임을 규명했습니다.
학습 전략의 혁신: 전체 시퀀스 마스킹을 도입하여 아키텍처적으로 가능했으나 학습되지 않았던 프롬프트 인필링 능력을 활성화했습니다.
실용적 적용 및 검증:
생성된 인필링 프롬프트가 수동으로 설계된 템플릿보다 성능이 우수하거나 동급임을 입증했습니다.
한 모델에서 학습된 인필링 프롬프트가 다른 모델 (Transferred) 에도 효과적으로 전이됨을 확인했습니다.
기존 자동 프롬프트 최적화 방법 (COPRO, GEPA 등) 과 결합 시 시너지 효과를 내어 추가적인 성능 향상을 이끌어냈습니다.
4. 실험 결과 (Results)
저자들은 GSM8K(수학 문제), SummEval(요약 평가), BigGen-Bench(다양한 LLM-as-a-Judge 태스크), HoVer(팩트 검증) 등 다양한 벤치마크에서 실험을 수행했습니다.
GSM8K (공식 체크포인트 한계 확인): 공개된 LLaDA 및 Dream 모델 (RO 학습) 로 프롬프트 인필링을 시도한 결과, 성능이 급격히 하락했습니다 (예: LLaDA 73.1% → 64.8%). 이는 학습 데이터가 부족할 때 인필링이 실패함을 보여줍니다.
SummEval & BigGen-Bench (FS 의 효과):
LLaDA: 전체 시퀀스 마스킹 (FS) 을 적용한 모델은 Judge 템플릿 인필링 시 Spearman 상관관계가 0.490으로 크게 향상되었으며, 추가 정제 (FS+RO) 를 거치면 0.535까지 도달했습니다. 이는 공개 체크포인트 (0.087) 나 응답 전용 마스킹 (0.477) 보다 월등히 높습니다.
BigGen-Bench: FS 모델이 가장 낮은 퍼플렉시티 (Perplexity) 와 높은 일관성을 보이며, FS+RO 보다 FS 단독이 더 강력한 일반화 성능을 보였습니다.
프롬프트 전이 (Prompt Transfer): FS+RO 모델이 생성한 인필링 프롬프트를 다른 모델 (RO 모델 등) 에 적용했을 때, 해당 모델의 성능이 크게 향상되었습니다. 이는 인필링된 프롬프트가 모델 간에 공유 가능한 최적화된 지식임을 의미합니다.
효율성: GSM8K 에서 인필링된 프롬프트는 16-shot In-Context Learning (ICL) 보다 **더 적은 토큰 수 (약 315 토큰 vs 3333 토큰)**로 더 높은 정확도 (76.4% vs 73.1%) 를 달성했습니다.
기존 최적화 방법과의 결합: COPRO 및 GEPA 와 같은 자동 프롬프트 최적화 기법에 인필링을 결합하면, 특히 COPRO 의 경우 분산이 크게 줄어들고 성능이 추가 향상되었습니다.
5. 의의 및 결론 (Significance)
이 논문은 확산 언어 모델 (dLM) 의 프롬프트 인필링 실패가 아키텍처의 한계가 아니라 학습 데이터의 마스킹 전략 부재에서 기인함을 증명했습니다.
패러다임 전환: "아키텍처적 혁신"을 찾기 전에, 기존 모델의 잠재력을 끌어낼 수 있는 **적절한 학습 신호 (Training Signal)**의 중요성을 강조합니다.
실용적 가치: 수동 프롬프트 엔지니어링의 노동을 줄이고, Few-shot 예시로부터 자동으로 최적화된 프롬프트를 생성하여 다양한 LLM 에 적용할 수 있는 새로운 프레임워크를 제시합니다.
향후 방향: 프롬프트 인필링은 기존 자동 최적화 방법과 상호 보완적이며, 확산 모델의 양방향 특성을 활용한 효율적인 프롬프트 최적화 도구로 자리 잡을 수 있음을 시사합니다.
요약하자면, 이 연구는 단순한 학습 전략 변경 (전체 시퀀스 마스킹) 만으로 확산 언어 모델이 가진 강력한 인필링 능력을 해방시켜, 더 짧고 효과적인 프롬프트를 자동으로 생성하고 다양한 모델 간에 전이할 수 있음을 입증했습니다.