completion-only loss - 질문 말고 정답만 학습시키기

SFT 데이터를 학습시킬 때 질문과 정답이 한 시퀀스로 들어가는데, 처음엔 모델이 이 전체를 다 학습하는 줄 알았다. 그런데 질문은 학습하지 않고 정답 부분에만 loss를 거는 방법이 따로 있었다. 입력으로 받는 것과 loss를 거는 것의 차이를 정리해봤다.

입력으로 받는 것 ≠ loss를 거는 것#

헷갈렸던 지점이 여기였다. 질문 텍스트는 당연히 모델에게 입력으로 들어간다. 그래야 모델이 질문을 보고(attend) 정답을 풀 수 있으니까. 하지만 학습 단계에서 이 질문 텍스트를 모델이 맞게 생성했는지 채점(loss 계산)할 필요는 없다.

시험 공부에 비유하면 이렇다.

• 질문(입력): 대한민국의 수도는? → 이건 그냥 읽고 이해하면 된다
• 정답(출력): 서울 → 이걸 제대로 뱉어내는 능력만 연습하면 된다

질문 텍스트를 생성하는 능력은 굳이 학습할 필요가 없는 것이다.

DataCollatorForCompletionOnlyLM을 안 쓰면#

기본적으로 collator를 안 쓰면 모델은 시퀀스 전체에 대해 next token prediction을 학습한다. 그러면 이렇게 된다.

• 대한민국의 수도는? 이라는 문장 자체도 완성하려고 학습한다. 대 다음엔 한, 한 다음엔 민... 이런 식으로. 불필요한 에너지 낭비다.
• 그리고 서울도 맞추려고 학습한다.

쓰면#

collator를 쓰면 정답 앞부분(질문, 선지)은 그냥 힌트로만 보고 채점하지 않는다.

• 대한민국의 수도는? 부분은 조건으로만 두고 loss를 계산하지 않는다
• 오직 서울이라고 대답하는 능력만 집중적으로 채점하고 교정한다

이렇게 하면 제한된 데이터와 자원으로 정답을 맞추는 능력에 학습이 집중된다. 수능 문제처럼 긴 지문과 선지가 입력의 대부분을 차지하는 경우, 그 긴 입력을 생성하는 데 loss를 낭비하지 않는 게 중요했다.

trl 기준으로는 completion_only_loss=True 옵션으로 completion 부분만 loss에 포함시킬 수 있다. 대신 어디까지가 입력이고 어디부터가 정답인지를 자동 인식하도록 데이터 구조를 맞춰줘야 한다.

eos 토큰은 학습되어야 한다#

여기서 같이 챙겨야 했던 게 eos 토큰이다. 정답만 학습한다고 해서 정답 끝의 eos까지 빼면 안 된다. 모델이 답을 어디서 멈춰야 하는지 배우려면 eos는 정답 영역에 포함돼 학습되어야 한다. 디버깅 중에 정답 다음 토큰을 예측하는 위치가 하드코딩돼 있거나 eos가 파편화되는 문제가 있어서, 정답 토큰 다음에 eos가 제대로 학습되도록 collator 동작을 맞췄다.