DPO 학습 계산 그래프를 확률적으로 가지치기해 1.4배 빠르게 훈련한다

무엇인가

이 논문은 DPO 훈련 한 스텝의 비용을 출력 점수화, 시퀀스 실행, 역전파라는 세 단계로 나누고, 각 단계에서 어떤 계산을 생략할 수 있는지 다룬다. DPO는 선택·거부 응답의 정책-참조 로그비를 평균해 마진을 만들고 로지스틱 손실을 최소화하는데, 큰 어휘에서는 T×V 로짓을 물질화하는 출력 점수화가 메모리를 크게 차지하고, 시퀀스 실행의 어텐션·활성값 저장과 양쪽 응답 역전파도 비용을 만든다. 기존의 토큰 선택이나 메모리 효율 커널은 감독이나 커널을 바꾸지만, 저자들은 감독을 줄이는 것만으로는 계산이 반드시 줄지 않는다고 보고 계산 그래프 자체를 확률적으로 잘라 자원 절감과 목적 함수 변화를 함께 분석하는 틀을 제안한다.

어떻게 동작하나

방법은 세 가지 독립적인 샘플링을 해당 연산 전에 수행한다. 위치 가지치기는 응답 위치의 균일 부분집합을 골라 그 위치에서만 언어모델 헤드를 적용해 어휘 투영과 로짓, 그 역전파를 건너뛰고, 시퀀스 가지치기는 응답 접두 길이를 ρ~U(ρ_min,1)로 뽑아 잘라 정책과 참조 모델 실행 전에 접미사 어텐션·활성값·재계산을 제거한다. 분기 가지치기는 선택 응답과 거부 응답 중 한쪽 역전파 그래프만 남기고 반대쪽은 no_grad로 실행해 마진 계산에 필요한 보상은 유지한다. 세 선택은 실행 순서대로 조합되며 기본값은 α=0.6, ρ_b~U(0.5,1), q_w=q_l=0.5다.

무엇과 다른가

이론적으로 위치 가지치기는 균일 샘플링이라 평균 보상과 평균 기울기를 보존하지만, 비선형 로지스틱 손실 때문에 분산에 비례하는 손실 갭이 생긴다. 논문은 이 갭을 β²/8·Var(Δ_U)로 묶고, 최적점 손실 차이가 대략 (1-α)/(α(T0-1))에 비례해 응답 길이가 길수록 작아지며 α=1에서 사라진다고 주장한다. 시퀀스 가지치기는 접두 포함과 길이 정규화 때문에 앞쪽 토큰에 더 큰 시간 가중치를 주고, 평균 마진 자체를 이동시킬 수 있어 위치 가지치기와 다르다. 분기 가지치기는 선택 확률 q_b로 나눈 역확률 가중 손실을 써서, 위치와 접두를 고정한 조건에서 기울기가 편향되지 않음을 보인다. 다만 분기 선택은 목적 함수 편향 없이 기울기 분산을 더하고, 강볼록·매끄러움 가정 아래 투영 SGD의 평균제곱 오차가 O(1/N)으로 줄어든다는 분석을 제시한다.

어떻게 쓰나

실험은 Qwen3-4B, Qwen3-8B, Mistral-7B-Instruct-v0.3과 HelpSteer3-Preference를 사용했고, LoRA 랭크 64, FlashAttention-2, 그래디언트 체크포인팅, AdamW, DPO β=1, 전역 배치 64, 8개의 Ascend 910B 32GiB 환경에서 수행됐다. 기준선은 원래 패딩 스택을 쓰는 Full과 선택·거부 응답을 함께 묶어 패딩을 줄인 Full-concat이다. Full-concat은 17.343초·23,627MB로 Full의 18.204초·18,571MB보다 1.049배 빨랐지만 메모리는 27% 더 썼다. 중요한 귀속 실험에서 전체 감독을 유지한 채 선택 위치에서만 헤드를 실행하는 구현이 이미 8,905MB를 절감했고, 위치 샘플링 자체가 추가로 줄인 것은 642MB였다. 즉 출력 메모리 절감의 93.3%가 구현에서, 6.7%가 감독 삭제에서 나왔다.

전제와 한계

주요 결과는 분기 가지치기의 시간 이득이 가장 뚜렷하다는 것이다. Full-concat 대비 분기 가지치기는 스텝 시간을 29.9% 줄여 1.426배 빨라졌고 피크 HBM은 26.7% 낮췄다. 위치+분기 조합은 1.443배 속도와 48.6% 낮은 HBM을 보였지만, 절감량 11,471MB가 각각의 절감량 합 15,855MB보다 작은 비가산성은 원문에서도 원인이 미해결로 남아 있다. 시퀀스 가지치기는 Full 입력 토큰의 0.85~0.94만 실행하고 피크 HBM을 8~11% 줄였으나, 스텝 시간은 18.430~18.482초로 Full 대비 최대 2.1% 감소에 그쳤다. 8개 구성 요인 실험에서는 분기 가지치기가 포함된 네 구성이 126~128분에 끝나 분기가 없는 185~188분보다 31~32% 빨랐다.

정확도와 개방형 평가는 공짜가 아니다. 두 Qwen 시드와 Mistral에서 Full 대비 분기 가지치기는 전체 학습 시간을 29.9~32.7% 줄였지만 평균 정확도는 0.25~0.45%p 낮아졌고, Mistral 프로필에서는 1.507배 속도와 1,347MB(6.4%) 메모리 절감을 보였다. 위치 가지치기에서 α를 0.3에서 0.9로 올리면 평균 정확도는 57.90에서 58.32로 올랐지만 피크 HBM은 11,211~11,580MB로 거의 변하지 않았다. 개방형 점수는 분기 가지치기가 MT-Bench와 WildBench 점추정을 올리고 Arena-Hard는 Full에 가깝게 유지한 반면, 위치나 시퀀스가 포함된 구성은 Arena-Hard에서 Full보다 낮았다. 저자들은 이 개방형 점수가 단일 시드의 서술적 요약이며 우월성을 입증하지 않는다고 명시한다.

개발자 관점에서 이 논문의 실무적 메시지는 먼저 출력 헤드 구현을 최적화하고, 그다음 시간이 병목이면 분기 가지치기를, 메모리가 병목이면 위치·시퀀스 가지치기를 조합해 보라는 것이다. 특히 위치·시퀀스 가지치기의 메모리 절감 상당 부분은 감독을 실제로 삭제해서가 아니라 선택 위치에서만 헤드를 실행하는 구현에서 나오므로, 무작정 토큰 감독을 줄이는 것보다 선택적 헤드 실행을 먼저 적용하는 편이 목적 함수 변화를 줄일 수 있다. 분기 가지치는 조건부 기울기 편향은 없지만 분산이 커지고, 위치·시퀀스 가지치기는 목적 함수 자체를 바꾸므로 평균 정확도와 벤치마크별 개방형 점수를 자체 검증해야 한다. 또한 이득은 하드웨어, 배치 기하, 참조 모델 실행 방식에 따라 달라질 수 있다.

저자가 밝힌 한계는 평가 범위가 하나의 8가속기 환경에서 온라인 참조를 쓰는 파라미터 효율적 텍스트 선호 학습에 한정된다는 점이다. 분기 가지치기는 두 Qwen 시드와 Mistral로 교차 검증됐지만 위치·시퀀스 및 그 조합은 Qwen에서만 평가됐고, 개방형 점수는 단일 시드 서술 통계라 훈련 시드 변동성을 특성화하지 못한다. 위치+분기의 비가산성과 시퀀스 토큰 감소가 시간으로 잘 이어지지 않는 물리적 원인은 미해결로 남아 있고, 시퀀스 가지치기의 이론적 최적점 한계는 ρ_min=1/2에서 손실 범위 전체와 같아져 사실상 공허하다. 따라서 이 논문은 DPO 계산 그래프를 어디서 잘라야 자원이 줄어드는지에 대한 통제된 측정과 조건부 보장을 제공하지만, 모든 DPO 워크로드에 대한 보편적 속도 향상이나 정확도 보존을 주장하지는 않는다.