비동기 온폴리시 증류(AsyncOPD) 해설: GPU 놀리지 않고 학습 속도 높이는 법

핵심 메시지

• 온폴리시 증류는 학생 모델이 스스로 문장을 만들고 교사 모델이 그것을 채점하는 방식인데, 순차적으로 돌리면 채점과 역전파가 진행되는 동안 생성용 GPU가 놀게 된다.
• 논문이 제안한 방식은 생성을 멈추지 않고 계속 돌리면서 채점과 학습을 겹쳐 실행하는 완전 비동기 구조다. 강화학습 쪽의 파이프라인 방식을 증류로 옮겨온 셈이다.
• 그 결과 수학 과제 실험에서 처리량이 대략 1.5~2배 이상 개선됐고, 벤치마크 성능은 편차 범위 안에서 비슷하게 유지됐다.
• 문제는 학생 모델이 계속 바뀌면서 미리 저장해둔 상위 후보 토큰이 어긋나는 캐시 미스다. 논문은 각 단계에서 다음 토큰을 여러 번 뽑아두고 나중에 중요도 샘플링으로 보정하는 방식으로 이를 우회한다.
• 참석자들은 훨씬 단순한 몇 단계 지연 방식만으로도 의미 있는 속도 향상이 나온다는 점을 들어, 완전 비동기의 복잡도가 늘 정당화되는지에 대해서는 의견이 갈렸다.

쉽게 이해하기

허깅페이스가 진행하는 논문 스터디에서 비동기 온폴리시 증류를 다뤘다. 증류는 큰 교사 모델의 능력을 작은 학생 모델에 옮기는 기법이고, 온폴리시 증류는 학생이 직접 만들어낸 문장을 교사가 평가해 학습에 쓰는 방식이다. 강화학습과 비슷해 보이지만 결정적인 차이가 있다. 강화학습은 궤적에 점수 하나만 붙이면 되지만, 증류는 학생과 교사 양쪽의 확률값을 함께 다뤄야 한다.

기존의 동기식 구조는 흐름이 단순하다. 학생이 문장을 생성하고, 교사가 채점하고, 역전파로 가중치를 갱신하고, 갱신된 가중치를 동기화한 뒤 다시 생성한다. 문제는 채점과 역전파가 도는 동안 생성 쪽 GPU가 그냥 멈춰 있다는 점이다. 일부 라이브러리는 몇 단계 정도 뒤처진 정책으로 미리 생성을 시작해 이 공백을 메우는 방식을 쓰지만, 얼마나 뒤처져도 되는지에는 한계가 있다.

논문이 제시한 방식은 여기서 한발 더 나간다. 가중치 동기화를 기다리지 않고 계속 생성하며, 서로 다른 시점의 정책에서 나온 문장들이 섞인 상태로 채점과 역전파를 쉬지 않고 돌린다. 발표자는 이를 강화학습에서 쓰이는 파이프라인 구조를 증류에 그대로 옮긴 것이라고 요약했다. 수학 과제로 측정한 처리량은 설정에 따라 1.5배에서 2배 남짓, 일부 구간에서는 그 이상까지 올라갔고 벤치마크 정확도는 대체로 유지됐다.

가장 까다로운 부분은 확률값을 저장하는 방식이었다. 전체 어휘에 대한 확률을 모두 계산하면 정확하지만, 교사 모델이 별도 서버에 있는 경우 통신 비용이 급격히 커진다. 그래서 상위 몇 개 후보만 저장하는데, 여기서 문제가 생긴다. 교사는 고정돼 있으므로 교사 기준으로 계산하는 방향에서는 문제가 없지만, 학생 기준으로 계산하는 역방향에서는 학생 정책이 바뀌면서 정작 필요한 토큰의 확률값이 저장돼 있지 않은 상황이 벌어진다.

해법은 각 단계에서 다음 토큰을 한 번만 고르지 않고 여러 번 뽑아 그 확률값들을 함께 저장해두는 것이다. 이때 문장을 여러 갈래로 뻗어 끝까지 생성하는 것이 아니라, 같은 지점에서 다음 토큰만 여러 번 표본으로 뽑아 기록한다. 나중에 손실을 계산할 때 저장해둔 값과 현재 정책의 확률을 비교해 중요도 샘플링으로 보정한다. 참석자들은 이 방식이 사실상 추가 비용이 거의 없다는 점을 확인하고서야 설계가 과하지 않다는 데 동의했다.

주요 인사이트

• 속도를 얻는 대가는 '얼마나 뒤처진 정책까지 허용할 것인가'라는 조절 손잡이로 나타난다. 표본을 많이 뽑을수록 더 큰 지연을 견딜 수 있지만, 지연이 아주 커지면 정확도는 결국 떨어진다. 공짜 점심은 없다는 뜻이다.
• 정방향과 역방향 계산은 목적이 다르다. 교사 쪽에 가중치를 두는 방향은 교사의 분포를 넓게 흡수하려 하고, 학생 쪽에 가중치를 두는 방향은 특정 능력을 좁고 정확하게 겨냥한다. 수학처럼 영역이 분명한 능력을 옮길 때 후자가 유용하다는 설명이 나왔다.
• 교사 모델을 어디에 두느냐가 설계를 좌우한다. 교사가 학습용 GPU에 함께 올라간다면 전체 어휘 확률을 쓸 수 있지만, 별도 서버에 있으면 통신량이 토큰 수에 비례해 늘어나기 때문에 상위 후보만 주고받는 절충이 불가피하다.
• 가장 실용적인 결론은 오히려 절충안이었다. 완전 비동기가 아니라 두 단계 정도만 뒤처지게 하는 단순한 구조로도 눈에 띄는 속도 향상이 나왔고, 이 경우 복잡한 보정 장치가 필요 없다는 점이 논의에서 반복해서 언급됐다.

자주 묻는 질문

관련 AI 소식