LLM 사전학습 직접 구현: 교차 엔트로피 손실, 온도 조절 샘플링, GPT-2 가중치 불러오기

핵심 메시지

• 사전학습의 목표는 정답 토큰 자리에 높은 확률이 찍히도록 만드는 것이고, 그 정도를 숫자로 재는 장치가 교차 엔트로피 손실이다.
• 교차 엔트로피는 통계학에서 말하는 '음의 평균 로그 확률'과 같은 값이며, 파이토치의 cross_entropy 함수 한 줄이 소프트맥스·로그·평균·부호 반전을 한꺼번에 처리한다.
• 데이터를 학습용 90%와 검증용 10%로 나눠 두 손실을 함께 기록하면, 학습 손실만 떨어지고 검증 손실이 멈추는 과적합 신호를 그래프로 확인할 수 있다.
• 온도 조절과 top-k 샘플링은 항상 최고 확률 토큰만 고르는 방식을 대체해 생성 결과의 무작위성을 조절하며, 온도를 0으로 두면 다시 결정론적 생성으로 돌아간다.
• 직접 5분간 학습한 모델로는 쓸 만한 문장이 나오지 않으므로, 오픈AI가 공개한 GPT-2 가중치를 같은 구조에 옮겨 담아 제대로 된 모델을 확보한다.

쉽게 이해하기

이 강의는 책의 5장에 해당하며, 앞선 장에서 따로 만들어 둔 데이터 로더와 멀티헤드 어텐션, GPT 모델 구조를 한자리에 모아 실제로 사전학습을 돌려 보는 자리다. 강사는 먼저 텍스트를 토큰 ID로 바꾸고 다시 텍스트로 되돌리는 보조 함수 두 개를 만든 뒤 모델에 문장을 넣어 본다. 아직 가중치가 무작위인 상태라 결과는 앞의 입력 네 토큰만 맞고 나머지는 뜻이 통하지 않는 문자열로 나온다.

다음 단계는 그 엉망인 결과를 사람의 눈이 아니라 숫자로 평가하는 일이다. 입력과 정답은 한 토큰씩 밀린 관계이고, 모델은 각 위치마다 어휘 크기인 50,257차원의 로짓을 뱉는다. 여기에 소프트맥스를 씌워 확률로 바꾼 뒤 정답 토큰 자리의 확률만 골라 로그를 취하고 평균을 낸 다음 부호를 뒤집으면 그것이 교차 엔트로피 손실이며, 학습 전 예제에서는 10.79라는 큰 값이 나온다.

실제 학습 루프는 단출하다. 에폭을 돌면서 배치마다 기울기를 초기화하고, 손실을 구해 역전파한 뒤 AdamW 옵티마이저로 가중치를 갱신하며, 정해진 주기마다 학습·검증 손실을 기록하고 샘플 문장을 뽑아 본다. 강사는 2020년형 맥북에어의 CPU에서 돌리려고 컨텍스트 길이를 256으로 줄이고 이디스 워튼의 단편 하나만으로 10에폭을 학습시켜 약 5분 만에 끝낸다.

데이터가 워낙 작다 보니 학습 손실은 0에 가까워지는데 검증 손실은 6 근처에서 멈춰 버린다. 생성된 문장을 원문에서 검색해 보면 표현이 그대로 등장할 만큼 외운 티가 난다. 그래서 로짓을 온도로 나눠 분포를 뾰족하게 또는 평평하게 만드는 온도 조절과, 상위 k개를 뺀 나머지 로짓을 음의 무한대로 덮어 확률을 0으로 만드는 top-k 샘플링을 결합해 생성 함수를 새로 쓴다.

마지막은 모델 저장과 가중치 이식이다. 파이토치는 state_dict를 저장하는 방식을 권장하는데 여기에는 구조 정보가 없어 불러올 때 같은 아키텍처를 먼저 만들어야 한다. 이어서 오픈AI가 공개한 GPT-2 가중치를 내려받아 옮겨 담는데, 컨텍스트 길이를 1024로 되돌리고 QKV 편향을 켜는 등 원본 설정을 맞춰야 하고 하나로 합쳐진 QKV 행렬을 쪼개 대응시키는 작업이 필요하다.

주요 인사이트

• 손실 함수의 이름이 둘인 이유가 흥미롭다. 통계학 쪽에서는 음의 평균 로그 확률, 정보이론 쪽에서는 교차 엔트로피라고 부르는데 두 분야가 각자 같은 값을 찾아낸 결과이고 파이토치는 후자의 이름을 택했다.
• 작은 데이터에서 나타나는 과적합은 코드가 잘못돼서가 아니라 예정된 결과다. 강사는 보너스 자료로 공개 저작권이 풀린 프로젝트 구텐베르크 전체를 학습시킨 사례를 들며, 매번 다른 텍스트를 보면 검증 손실 정체가 사라진다고 설명한다.
• 온도 값 0은 무작위성을 끄고 argmax로 되돌리는 스위치 역할을 한다. 덕분에 결정론적 생성과 확률적 생성을 하나의 함수로 처리할 수 있다.
• top-k는 말이 안 되는 토큰이 뽑힐 가능성을 원천 차단하는 대신 창의성을 줄이는 맞교환이다. 어휘가 5만 개인 실제 모델에서는 50이나 100, 500 같은 값을 쓰며 보통 모델 카드의 권장값을 따르는 편이 낫다.
• 규모 감각도 함께 짚는다. 실제 상용 모델은 14~15조 토큰을 수천 대의 GPU로 몇 달간 학습해 수백만 달러가 들기 때문에 보통 1에폭에서 끝내지만, 교육용 예제는 데이터가 짧아 오히려 10에폭을 돌려야 학습 효과가 보인다.

자주 묻는 질문

관련 AI 소식