GPT-2 위치 임베딩을 RoPE로 바꾸니 10억 토큰 뒤 loss가 0.122 낮았고 학습은 5% 느렸습니다

GPT-2 위치 임베딩을 RoPE로 바꾸니 10억 토큰 뒤 loss가 0.122 낮았고 학습은 5% 느렸습니다

GPT-2(124M)에서 학습되는 위치 임베딩을 빼고 그 자리에 RoPE(회전 위치 임베딩)를 넣었습니다. 다른 것은 하나도 바꾸지 않았습니다. 두 조건을 세 번씩, 같은 10억 토큰을 같은 순서로, 같은 학습률로 학습했습니다. RoPE 쪽 평균 검증 loss는 3.552, 원래 GPT-2 쪽은 3.674로 끝났습니다.

차이는 0.122 nats입니다. 앞 글에서는 시드만 바꿨을 때 결과가 얼마나 흔들리는지 재고, 이 연재에서 차이로 인정할 기준을 정했습니다. 이번 실행으로 계산한 기준은 0.012이니, RoPE는 그 기준을 열 배쯤 넘었습니다.

이 연재는 GPT-2에서 출발해 이후 학습 레시피들이 바꾼 것을 하나씩 더해 보고 있고, 이번이 두 번째 변경입니다. modded-nanogpt 스피드런에서는 RoPE가 기록 2번에 학습률 조정과 함께 들어갔고, 두 변경의 효과가 한데 묶여 있습니다. 이 글은 RoPE 하나만 바꿨고, 학습률은 GPT-2 값인 6e-4를 그대로 썼습니다.

요약

• RoPE 쪽이 0.122 nats 낮게 끝났습니다(3.552와 3.674, 각 세 번의 평균). 이 비교의 기준은 미리 정한 규칙으로 0.012가 나왔으니, 차이로 셉니다.
• 모든 평가 지점에서 RoPE 세 실행이 모두 GPT-2 세 실행보다 낮았습니다. 가장 나쁜 RoPE 실행이 3.554, 가장 좋은 GPT-2 실행이 3.668로 끝났습니다.
• 차이는 초반에 가장 컸고 끝날 때까지 줄고 있었습니다. 1.3억 토큰에서 0.443, 5.2억 토큰에서 0.161, 10억 토큰에서 0.122였습니다. 더 오래 학습하면 어디서 멈출지는 이번 실행으로는 알 수 없습니다.
• 이 코드에서 RoPE는 학습 처리량이 5.3% 낮았습니다(A100 두 장에서 초당 34.2만 토큰과 36.1만 토큰). 그래도 RoPE 실행은 6.6억 토큰과 7.9억 토큰 사이 어딘가에서 이미 GPT-2 실행이 끝날 때의 loss보다 낮아졌습니다.

바꾼 것과 그대로 둔 것

모델은 앞 글과 같은 GPT-2 small입니다. 12층, 폭 768, 64차원 헤드 12개, 문맥 1,024토큰, 입력 임베딩과 출력층이 행렬을 같이 쓰는 구조이고, GELU를 GPT-2와 같은 tanh 근사로 바꾼 nanoGPT 코드입니다. 다른 점은 설정 하나뿐입니다.

나머지는 모두 같습니다. FineWeb-Edu 샤드를 같은 순서로 읽고, 524,288토큰짜리 스텝을 1,907번(10억 토큰) 돌렸습니다. 옵티마이저는 AdamW, 최대 학습률 6e-4, 워밍업 100스텝 뒤 코사인으로 10%까지 줄였고, bfloat16으로 A100 80GB 두 장에서 학습했습니다. 실행 기록을 보면 짝마다 다른 설정은 이 하나뿐이고, 여섯 실행 모두 첫 배치가 같았습니다.

GPT-2 쪽 세 실행은 새로 돌리지 않았습니다. 계획대로 앞 글의 시드 세 실행을 그대로 썼습니다. RoPE 실행에도 같은 시드 0, 1, 2를 썼지만 처음 가중치까지 같아지지는 않습니다. 위치 표가 빠지면서 초기화 과정에서 난수를 꺼내 쓰는 순서가 달라지기 때문입니다.

RoPE 실행을 시작하기 전에, 구현이 RoPE의 핵심 성질을 지키는지 CPU에서 확인했습니다. 쿼리와 키의 어텐션 점수는 두 위치가 얼마나 떨어져 있는지에만 달려 있어야 합니다. 일곱 칸 떨어진 쿼리·키 쌍 세 개의 내적이 모두 −14.552로 같았고, 회전해도 벡터 길이는 그대로였습니다.

실행 전에 정한 판정 규칙

규칙은 앞 글에서 정한 것을 씁니다. 두 조건의 평균 최종 loss 차이가 2 × σ × √(1/n₁ + 1/n₂)보다 클 때만 차이로 셉니다. σ는 두 조건의 시드 간 표준편차를 합쳐서 구합니다. 이 비교의 질문과 실행, 판정 규칙은 10시 11분에 적어 커밋했고, 첫 RoPE 실행은 12시 50분에 시작했습니다. 규칙을 계산하는 스크립트는 그보다 늦은 13시 55분에 커밋했습니다. 첫 RoPE 실행이 끝난 뒤라서 결과를 보지 않고 썼다고는 말하지 않겠습니다. 스크립트는 커밋해 둔 규칙을 그대로 계산할 뿐이고, 규칙에 아무것도 더하지 않았습니다.

합친 σ는 0.0073이고, 기준은 2 × 0.0073 × √(2/3) = 0.0119입니다. 차이는 −0.1224입니다.

앞 글에서 시드 3개일 때의 기준을 0.0163으로 적었는데, 이번에는 그보다 낮게 나왔습니다. RoPE 세 실행이 우연히 서로 더 가깝게 끝나서(표준편차 0.0025) 합친 σ가 내려갔기 때문입니다. 실행 세 번으로는 표준편차를 정확히 잡을 수 없으니, RoPE가 시드에 덜 흔들린다고 읽지는 않겠습니다. 판정에는 영향이 없습니다. 차이가 어느 기준과 비교해도 일곱 배를 넘습니다.

학습 전체에서 본 차이

(각 세 실행의 평균입니다. 두 조건의 실행은 한 번도 겹치지 않았습니다. 모든 지점에서 가장 높은 RoPE 실행이 가장 낮은 GPT-2 실행보다 낮았습니다.)

표에서 두 가지를 읽을 수 있습니다. 첫째, 차이는 첫 평가 지점인 1.3억 토큰에서 이미 나타납니다. 계획의 두 번째 질문이었던 "언제부터 벌어지는가"의 답은 학습 후반이 아니라 처음부터입니다. 둘째, 차이는 꾸준히 줄어드는데, 처음에는 빠르게 줄고 나중에는 느리게 줄었습니다. 첫 지점에서 중간 지점(5.2억 토큰)까지 0.28 줄었고, 나머지 절반 동안에는 0.04 줄었습니다.

이렇게 줄어드는 모습을 설명할 수 있는 해석이 하나 있습니다. 학습되는 위치 표는 난수(정규분포, 표준편차 0.02)에서 출발해 1,024개 위치가 각각 무엇을 뜻하는지 데이터에서 배워야 합니다. RoPE는 배울 것 없이 첫 스텝부터 어텐션에 토큰 사이 거리를 알려 줍니다. 이 설명이 맞다면 GPT-2 쪽 위치 표가 따라잡는 중이고, 더 길게 학습하면 차이는 더 줄어들 것입니다. 다만 이것을 시험하지는 않았습니다. 마지막 평가 지점에서도 차이는 줄고 있었고(마지막 8,200만 토큰 동안 0.0024), 이번 실행만으로는 차이가 없어질지, 어딘가에서 멈출지 말할 수 없습니다.

비용

시간이 어디서 더 드는지 프로파일링하지는 않았습니다. 회전은 12개 층 모두에서 쿼리와 키에 더해지는 원소별 연산이고, 이 코드에서는 일반 PyTorch로 써서 torch.compile에 맡겼습니다. 연산을 합친 전용 커널을 쓰면 비용이 더 적을 가능성이 큽니다.

이 비용을 넣어도 비교 결과는 바뀌지 않습니다. RoPE 실행은 7.9억 토큰 지점에서 이미 GPT-2 실행의 최종 loss(3.674)보다 낮았고(3.609), 6.6억 토큰 지점에서는 아직 높았습니다(3.681). RoPE의 처리량으로 1,500스텝을 학습하는 데 38분이 걸리고, GPT-2 쪽이 1,907스텝 전체를 학습하는 데는 46분이 걸립니다. 둘 다 기록된 처리량으로 계산한 값이고, 평가와 컴파일 시간은 뺐습니다.

한계

• 학습률은 하나만 썼습니다. 두 조건 모두 GPT-2 값인 6e-4입니다. 스피드런은 RoPE와 학습률을 함께 바꿨습니다. 조건마다 학습률을 맞춰 고르면 두 숫자가 모두 움직일 수 있고, 그 둘을 아직 분리하지 않았습니다. 학습률은 5편에서 여러 값으로 바꿔 재 볼 계획입니다.
• 예산 하나, 모델 크기 하나입니다. GPT-2 small을 10억 토큰 학습했고, 끝날 때도 차이는 줄고 있었습니다.
• 긴 문맥은 재지 않았습니다. RoPE를 고르는 이유로 학습 길이보다 긴 문맥으로 늘리기 쉽다는 점이 자주 꼽히는데, 여기서는 학습과 평가 모두 1,024토큰입니다.
• RoPE 설정도 하나뿐입니다. 기저 10,000, 모든 차원에 적용했습니다. 일부 차원에만 쓰는 방식이나 다른 기저는 해 보지 않았습니다.
• 처리량 숫자는 이 구현의 숫자입니다. 다른 커널을 쓰면 다른 숫자가 나옵니다.

다음 글

다음 글에서는 LayerNorm을 RMSNorm으로 바꾸고, 그 위에 QK norm을 더합니다. 이번과 같은 GPT-2 세 실행과 비교합니다. 실행은 아직 진행 중입니다.

계획서, 스크립트, 여섯 실행의 원시 로그는 아래 재현 패키지에 있고 로그인 없이 받을 수 있습니다. README에 GPU 없이 그 로그에서 이 글의 숫자를 모두 다시 계산하는 명령이 있습니다.