길스 토머스(Giles Thomas)는 약 1억 6,300만 개 매개변수의 GPT-2 소형 모델을 학습해 RMSNorm과 LayerNorm을 비교했다. RMSNorm은 직접 구현했을 때 LayerNorm보다 약 2.5% 빨랐지만, PyTorch 내장 모듈끼리 비교하면 차이는 0.5%에 그쳤다. 실험 횟수가 적어 이 차이가 확실한지는 알 수 없으며, 직접 구현한 코드보다 PyTorch 내장 모듈을 썼을 때 학습 속도가 더 빨랐다.
RMSNorm은 벡터의 평균을 0으로 맞추는 LayerNorm의 재중심화 과정을 생략하고, 각 값을 벡터의 제곱평균제곱근(RMS)으로 나눈 뒤 학습 가능한 척도를 적용한다. 이 방식은 계산을 줄이지만, 이번 실험에서는 정규화 방식보다 직접 구현과 PyTorch 내장 구현 사이의 차이가 더 두드러졌다.
| 모델 | 학습 시간(초) | 초당 토큰 수 | 테스트 손실 |
|---|---|---|---|
| 직접 구현 LayerNorm | 9,358 | 348,364 | 3.595336 |
| PyTorch LayerNorm | 8,324 | 391,659 | 3.571158 |
| 직접 구현 RMSNorm | 9,134 | 356,899 | 3.601292 |
| PyTorch RMSNorm | 8,289 | 393,294 | 3.572383 |
실험 모델은 GPT-2 소형을 바탕으로 했으며, 가중치 공유와 QKV 행렬의 편향을 사용하지 않았다. 드롭아웃 없이 학습률을 처음 5% 구간에서 0에서 0.0014까지 올린 뒤 0.00014까지 코사인 방식으로 낮췄다. 매개변수당 약 20개 토큰을 학습했고, 테스트는 8개의 A100 GPU에서 진행했다.
직접 구현끼리 비교했을 때 RMSNorm의 학습 시간은 약 2.5% 줄었다. PyTorch 내장 모듈끼리 비교하면 개선 폭은 0.5%였다. 반면 내장 LayerNorm은 직접 구현한 LayerNorm보다 학습 시간이 11% 짧았고, 내장 RMSNorm도 직접 구현보다 10% 조금 못 미치는 시간 단축을 보였다. 토머스는 내장 모듈의 융합 커널이 속도 차이에 영향을 줬을 가능성을 들었지만, 더 많은 실험이 필요하다고 봤다.
테스트 손실은 RMSNorm 쪽이 두 비교에서 조금 높았다. 다만 모델을 각각 한 번씩만 학습했기 때문에 정규화 방식에 따른 차이라고 단정하기 어렵다고 토머스는 설명했다. PyTorch 내장 모듈의 손실이 직접 구현보다 낮게 나온 결과도 반복 실험이 없어 우연인지 체계적인 차이인지 확인되지 않았다.
RMSNorm의 계산량 절감이 대형 모델에서 얼마나 의미 있는지는 이번 실험으로 확인되지 않았다. 토머스는 모델 폭이 커질수록 다른 연산의 비용이 더 빠르게 증가할 수 있다고 추정했고, RMSNorm의 인기에는 구조적 이유나 단순성, 관행 등이 작용했을 가능성을 제시했다.
RMSNorm을 제안한 2019년 논문은 Root Mean Square Layer Normalization이다. 실험에 사용한 PyTorch의 LayerNorm 문서와 RMSNorm 문서도 참고할 수 있다.
부록: 원문 예시
class RMSNorm(nn.Module): def __init__(self, emb_dim): super().__init__() self.emb_dim = emb_dim self.scale = nn.Parameter(torch.ones(self.emb_dim)) self.eps = RMSNORM_EPS def forward(self, xs): squares = xs ** 2 mean_squares = squares.mean(dim=-1, keepdim=True) rms = torch.sqrt(self.eps + mean_squares) return (xs / rms) * self.scale
var = x.var(dim=-1, keepdim=True, unbiased=False)
def make_norm(cfg): norm_type = cfg.get("norm_type", "handrolled_layernorm") if norm_type == "handrolled_layernorm": return LayerNorm(cfg["emb_dim"]) elif norm_type == "pytorch_layernorm": return nn.LayerNorm(normalized_shape=cfg["emb_dim"], eps=LAYERNORM_EPS) elif norm_type == "handrolled_rmsnorm": return RMSNorm(cfg["emb_dim"]) elif norm_type == "pytorch_rmsnorm": return nn.RMSNorm( normalized_shape=cfg["emb_dim"], eps=RMSNORM_EPS, elementwise_affine=True ) elif norm_type == "no_norm": return nn.Identity() else: raise Exception(f"Unknown <!--CODE_BLOCK_15744--> of {norm_type!r}")
var = x.var(dim=-1, keepdim=True, correction=0)
댓글 (0)
로그인하면 이 기사에 내 생각을 남길 수 있어요