코드 편집을 확산 편집 경로로 예측해 추론 지연을 줄인다

무엇인가

이 논문이 푸는 문제는 AI 코딩 IDE의 다음 편집 예측이다. 확산 대규모 언어 모델(dLLM)은 여러 토큰을 병렬로 갱신할 수 있어 매력적이지만, 대부분은 마스크 확산으로 시퀀스를 처음부터 재생성하도록 최적화되어 있어 편집 예측과 맞지 않는다. IDE에서는 사용자의 최근 편집 이력, 현재 파일, 커서 주변 문맥을 보고 다음 코드 변경을 실시간으로 추론해야 하는데, 이런 변경은 보통 희소해서 소스 토큰 대부분을 그대로 보존한다. 자가회귀(AR) 모델은 바뀌지 않은 조각까지 토큰 단위로 직렬 출력해야 하고 이미 생성한 내용을 되돌릴 수 없으며 여러 위치의 편집을 병렬 처리하기 어렵다. 마스크 확산은 작은 변경 하나에도 전체 시퀀스를 디노이징해야 한다.

어떻게 동작하나

제안 방법인 Diff Diffusion은 디노이징 자체를 구조화된 편집 연산으로 표현한다. Dream-Coder에서 초기화하고 Transformer 구조는 유지하되, 디노이징 목표를 전체 시퀀스 재구성에서 편집 경로 예측으로 바꿨다. 현재 상태의 각 토큰에 편집 레이블을 부여하는데, token_id를 그대로 내보내면 유지, del_token_id는 삭제, insert_token_id는 현재 위치 앞에 마스크 토큰 삽입, replace_token_id는 현재 토큰을 마스크로 치환, 마스크 위치에서는 코드 토큰 id를 예측해 채운다. 삭제 표시된 토큰은 제거되고 요청된 마스크가 삽입·치환된 뒤, 다시 편집 레이블을 할당하는 과정을 반복한다. 이 방식은 마스크 확산의 병렬 갱신 이점을 생성이 아니라 편집으로 확장하며, 삽입·삭제 연산이 있으므로 가변 길이 출력이 과정에 내재한다. 학습에서는 스케줄러를 조정해 초기 상태·목표 상태·중간 확산 상태를 서로 다른 확률로 샘플링하고, 중간 상태는 초기 코드와 목표 코드의 LCS 정렬로 구성한다. 편집 예측 데이터는 대부분의 토큰이 unchanged이고 삽입·삭제·치환이 드물게 나타나는 심한 불균형을 보이므로, 손실에서 insert·delete·replace 레이블에 가중치 λ_edit를 주는 가중 교차엔트로피를 쓴다. 추론 시에는 ApplyDiff로 x_t에서 x_{t-1}을 반복 생성하며 편집 내용이 더 이상 변하지 않을 때까지 진행한다.

무엇과 다른가

학습은 3단계다. 1단계는 CommitPack의 다국어 부분집합으로 커밋 diff를 대규모 지속 사전학습해 diff 기반 모델링 형식과 출력 공간을 익힌다. 2단계는 원본 스니펫, 주변 문맥, 편집 이력, 목표 스니펫으로 구성된 대화형 IDE 입출력 형식에 맞춰 SFT를 수행하며, C3PO 학습셋과 고품질 편집 예측 데이터를 결합한다. 3단계는 모델이 스스로 만들어낸 중간 상태에서도 신뢰할 수 있는 편집 방향을 학습하도록 확장한 후처리로, 이산 흐름 매칭에서 착안해 중간 상태에서 예측한 diff를 목표 상태로 향하는 이산 수정 방향으로 본다. Diff-DPO는 정답 스니펫을 양성 상태, 모델이 샘플링한 오류 편집을 음성 상태로 두고 전이 로그확률 차이를 비교해 정답 편집 방향을 선호하도록 하며, 참조 정책과의 격차로 2단계 정책에 업데이트를 고정한다. AGRPO는 하나의 예시에서 G개의 완전한 편집 궤적을 샘플링하고 최종 상태가 정답인지에 반복 페널티를 더한 종료 보상을 그룹 상대 어드밴티지로 변환한 뒤, 자가회귀 토큰 로그확률이 아니라 디노이징 전이에 PPO식 클리핑 목적함수와 KL 항을 적용한다. 일부 확산 스텝만 샘플링해 저비용 몬테카를로 추정을 한다.

어떻게 쓰나

추론 엔진은 블록 단위 병렬 편집 디코딩과 근사 KV 캐시로 구성된다. 각 확산 스텝에서 same, delete, insert, replace, mask-to-token 동작을 병렬로 예측하고, 마스크가 아닌 위치는 유지 또는 특수 편집 동작으로 제한해 치환은 반드시 replace와 후속 마스크 채우기를 거치게 한다. 재작성 범위를 국소 블록으로 나눠 신뢰도 높은 편집을 병렬 적용하고, 원본 코드의 변경되지 않은 토큰을 좌우 앵커로 삼아 편집 경계를 제약한다. 삽입 길이는 적응적이어서 신뢰도가 높은 삽입 영역에서는 한 번의 insert 동작이 여러 토큰을 연속 방출한다. KV 캐시는 두 단계로 재사용하는데, 하나는 서비스 요청 내 안정적 프리픽스와 수렴한 블록의 KV 상태 재사용이고, 다른 하나는 긴 시퀀스 편집에서 편집 지점 뒤의 긴 불변 서픽스를 PIE 방식으로 RoPE 위치 보정 후 재사용하는 것이다.

전제와 한계

실험은 C3PO 편집 완성 벤치마크, IDE 편집 예측 평가셋, 온라인 A/B 검증으로 구성된다. AR 베이스라인은 확산 전환 이전에 프로덕션에서 쓰던 편집 예측 모델로 Seed-Coder에서 학습했고, 동일한 편집 이력·파일 문맥·편집 범위를 받는다. C3PO에서 Diff Diffusion은 결과 정확도 81.33%로 AR 베이스라인에 근접하면서 평균 디코딩 스텝을 14.62에서 3.90으로 줄였고, 샘플당 추론 시간은 145ms로 AR의 403ms 대비 2.78배 빨랐다. IDE 편집 예측 평가는 Act, Stay, Multi-line, Roll-back, Delete 다섯 지표로 나뉜다. 지속 사전학습만 거친 베이스 모델이 다섯 지표 모두에서 Mercury-Coder를 앞섰고 Multi-line에서는 이미 AR과 동등했으며 Roll-back과 Delete는 더 낮았다. SFT 후 Act 49.18%, Multi-line 73.68%가 됐고, DPO는 Roll-back을 낮게 유지하며 Stay를 개선했고, AGRPO는 Act 50.00%, Stay 96.40%, Delete 0.5%로 가장 강한 Act와 Stay를 기록했다.

지연 실험에서 140토큰 재작성 범위는 AR의 126 스텝을 약 5번의 편집 스텝으로 압축해 총 시간이 3964ms에서 약 415~420ms로 줄었다. 이 짧은 케이스에서는 PIE가 도움이 되지 않았는데 두 변형 모두 메모리 바운드라 서픽스 캐시 관리와 위치 보정 비용이 절약분을 상회했기 때문이다. 1024토큰 장문 재작성에서는 PIE를 통한 근사 서픽스 캐시 재사용이 총 시간을 784ms에서 676ms로 추가 단축했다. 프로덕션 온라인 A/B에서 Diff Diffusion은 수용률과 CPO를 AR과 비슷하게 유지하면서 사용자당 수용 라인을 9.6% 늘리고 P50 모델 추론 지연을 39.0% 줄였다. AR과 Diff Diffusion 모두 텐서 병렬 4, FP8 양자화로 서빙됐고 AR 쪽은 추가로 NGram 추측 샘플링을 켰다는 점도 조건으로 명시된다.

개발자 입장에서 이 논문의 의미는 코드 편집을 전체 재생성이 아니라 국소 변환으로 다루는 모델링 관점과, 그 관점이 실제 지연 이득으로 이어지는 경로를 함께 제시한다는 데 있다. IDE 편집 제안, 코드 리뷰 자동 수정, 에이전트의 파일 패치 생성처럼 변경이 희소하고 응답 시간이 중요한 작업에 직접 적용 가능한 형태다. 도입을 검토한다면 백본이 Dream-Coder라는 점, 1단계가 CommitPack 커밋 diff에 의존한다는 점, 지연 수치가 블록 단위 병렬 디코딩과 근사 KV 캐시를 포함한 전용 추론 엔진 조건에서 측정됐다는 점, 그리고 AR 베이스라인이 NGram 추측 샘플링까지 켠 프로덕션 모델이라는 비교 조건을 확인해야 한다.

저자들이 밝힌 한계는 두 가지다. 편집 경로 확산에 대한 더 깊은 이론적 분석을 아직 하지 못했고, 평가가 지연에 민감한 IDE 편집 예측에 집중되어 있다. 향후 계획으로는 지시 기반 편집, 저장소 수준 변경, 계획과 도구 사용 및 사람 검토를 결합한 에이전트 워크플로로 프레임워크를 확장하는 것을 제시한다.