현대 LLM 구조 1 (RoPE · RMSNorm · SwiGLU): 기본 Transformer에서 딱 세 군데를 바꾼 이유
요즘 LLM은 기본 Transformer 뼈대를 그대로 두고 위치 표현, 정규화, FFN 세 곳만 바꿨습니다. 오늘은 RoPE(회전으로 상대 위치를 넣는 방법), RMSNorm(LayerNorm을 단순화한 것), SwiGLU(게이트가 있는 FFN)를 직접 구현해 기존 부품과 바꿔 끼워 봅니다. 절대 위치 임베딩과 RoPE를 짧은 태스크에서 비교해 차이를 숫
By 송민성
오늘의 목표
요즘 공개된 LLM(LLaMA 계열 등)은 2017년 원본 Transformer와 뼈대가 거의 같지만, 세 부품이 다릅니다. 위치를 넣는 방식(절대 위치 임베딩 → RoPE), 정규화(LayerNorm → RMSNorm), FFN의 활성화(ReLU/GELU → SwiGLU)입니다. 오늘은 이 세 가지가 각각 무엇을 어떻게 바꾸는지 설명할 수 있고, 직접 구현해서 기존 코드에 갈아 끼울 수 있는 것이 목표입니다.
개념과 직관
이 글은 Transformer 블록(Attention + FFN + 정규화 + 잔차 연결)을 이미 만들어 봤다는 전제로 씁니다. 여기서는 그 블록 안의 부품 세 개만 교체합니다.
1. RoPE: 회전으로 상대 위치를 넣기
문제. Attention은 토큰의 순서를 모릅니다. "나는 너를 좋아해"와 "너는 나를 좋아해"를 구분하려면 위치 정보를 따로 넣어야 합니다. 기본 방식인 절대 위치 임베딩(absolute positional embedding, 0번, 1번, 2번 자리마다 학습 가능한 벡터를 입력에 더하는 것)은 학습 때 본 길이를 넘어가면 대응이 어렵다는 약점이 있습니다.
RoPE(Rotary Position Embedding, 회전 위치 임베딩). 입력에 위치 벡터를 더하지 않습니다. 대신 Q와 K 벡터를 위치에 비례하는 각도만큼 회전시킵니다.
직관. 시계를 떠올리세요. 벡터를 2차원씩 짝지어 각각을 시곗바늘이라 보고, 위치 m번 토큰의 바늘은 m × (각속도)만큼 돌려 놓습니다. 두 바늘의 내적(dot product, 두 벡터가 얼마나 같은 방향인지 재는 값)은 두 바늘 사이의 각도 차이에만 의존합니다. 즉 Q의 위치가 m, K의 위치가 n이면 내적은 절대 위치 m, n이 아니라 상대 거리 (m - n)에 의존하게 됩니다. 이것이 "회전으로 상대 위치를 표현한다"의 의미입니다.
최소 수식. 2차원 짝 (x1, x2)를 각도 θ만큼 회전:
x1' = x1 * cos(θ) - x2 * sin(θ)
x2' = x1 * sin(θ) + x2 * cos(θ)위치 m의 각도는 θ_i(m) = m × base^(-2i/d)입니다. 차원 쌍 i마다 각속도가 다릅니다. 앞쪽 쌍은 빨리, 뒤쪽 쌍은 천천히 돌아서 시계의 초침·분침·시침처럼 가까운 거리와 먼 거리를 함께 표현합니다. base는 보통 10000을 쓰고, d는 head 하나의 차원입니다.
적용 위치. Q와 K에만 적용하고 V에는 적용하지 않습니다. 위치 정보는 "누구를 얼마나 볼지"를 정하는 Q·K 내적에만 필요하기 때문입니다.
2. RMSNorm: LayerNorm 단순화
LayerNorm(층 정규화, 한 토큰 벡터의 평균을 0, 분산을 1로 맞추는 것)은 평균 빼기와 분산 나누기를 합니다. RMSNorm은 평균 빼기를 생략하고 크기(RMS, root mean square, 제곱 평균의 제곱근)만 맞춥니다.
RMS(x) = sqrt( mean(x^2) + eps )
y = x / RMS(x) * gg는 차원마다 학습되는 크기 조절 파라미터이고, eps는 0으로 나누기를 막는 아주 작은 수입니다. 비유하면 LayerNorm은 "기준선을 0으로 옮기고 폭도 맞추기", RMSNorm은 "폭만 맞추기"입니다. 계산이 덜 들고, 경험적으로 학습 품질이 비슷해서 널리 쓰입니다. (왜 평균 빼기가 필요 없는지에 대한 엄밀한 설명은 오늘 다루지 않습니다. "더 단순한데 잘 된다"까지만 받아들이세요.)
3. SwiGLU: 게이트가 있는 FFN
기본 FFN은 이렇습니다.
FFN(x) = W2 · act( W1 · x )SwiGLU FFN은 길이 하나를 더 만들어 "문"으로 씁니다.
FFN(x) = W2 · ( SiLU(W1 · x) * (W3 · x) )- SiLU(x) = x * sigmoid(x). 부드러운 ReLU 모양의 활성화 함수입니다.
*는 원소별 곱셈입니다. W3·x가 각 차원을 얼마나 통과시킬지 정하는 게이트(gate, 신호를 걸러 주는 문) 역할을 합니다.
비유: 기본 FFN은 수도꼭지 하나, SwiGLU는 수도꼭지 옆에 "얼마나 열지" 정하는 밸브가 하나 더 달린 구조입니다. 행렬이 2개에서 3개로 늘어나므로, 파라미터 수를 맞추려고 hidden 크기를 보통 4d 대신 약 (8/3)d 정도로 줄입니다. 이것은 관례이고 정해진 법칙은 아닙니다.
코드로 직접 해보기
PyTorch가 설치되어 있다고 가정합니다. 각 부품을 따로 만들고 shape을 확인합니다.
1단계: RMSNorm
import torch
import torch.nn as nn
class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.eps = eps
self.g = nn.Parameter(torch.ones(dim))
def forward(self, x):
# x: (batch, seq, dim)
rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
return x / rms * self.g
x = torch.randn(2, 5, 8)
print(RMSNorm(8)(x).pow(2).mean(-1)) # 모두 1에 가까워야 함2단계: SwiGLU FFN
import torch.nn.functional as F
class SwiGLU(nn.Module):
def __init__(self, dim, hidden):
super().__init__()
self.w1 = nn.Linear(dim, hidden, bias=False)
self.w3 = nn.Linear(dim, hidden, bias=False)
self.w2 = nn.Linear(hidden, dim, bias=False)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
print(SwiGLU(8, 21)(torch.randn(2, 5, 8)).shape) # (2, 5, 8)3단계: RoPE를 Q·K에 적용
짝짓기 방식은 "앞 절반과 뒤 절반을 짝으로 삼는" 방식으로 구현합니다. 구현마다 짝짓기 방식(인접 차원끼리 vs 반반)이 다르다는 점은 아래 막히는 지점에서 다시 설명합니다.
def rope_cache(seq_len, head_dim, base=10000.0):
# 차원 쌍마다 각속도
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
pos = torch.arange(seq_len).float()
angles = torch.outer(pos, inv_freq) # (seq, head_dim/2)
return angles.cos(), angles.sin()
def apply_rope(x, cos, sin):
# x: (batch, heads, seq, head_dim)
x1, x2 = x[..., : x.size(-1) // 2], x[..., x.size(-1) // 2 :]
cos = cos[None, None, :, :] # 브로드캐스트용 차원 추가
sin = sin[None, None, :, :]
return torch.cat([x1 * cos - x2 * sin,
x1 * sin + x2 * cos], dim=-1)
B, H, T, D = 1, 2, 6, 8
q = torch.randn(B, H, T, D)
k = torch.randn(B, H, T, D)
cos, sin = rope_cache(T, D)
q_r, k_r = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
print(q_r.shape) # 같은 shape
print(q.norm(dim=-1)[0,0,:3], q_r.norm(dim=-1)[0,0,:3]) # 회전이라 길이 보존상대 위치 검증(꼭 해보세요). 같은 q 벡터, 같은 k 벡터를 위치만 옮겨 넣고 내적이 거리에만 의존하는지 봅니다.
q0 = torch.randn(1, 1, 1, D).expand(1, 1, T, D) # 모든 위치에 같은 q
k0 = torch.randn(1, 1, 1, D).expand(1, 1, T, D) # 모든 위치에 같은 k
qr, kr = apply_rope(q0, cos, sin), apply_rope(k0, cos, sin)
s = (qr[0,0] @ kr[0,0].T) # (T, T) 점수표
print(s[2, 0].item(), s[3, 1].item(), s[5, 3].item()) # 거리 2 → 거의 같은 값세 값이 (부동소수점 오차 내에서) 같으면 성공입니다.
4단계: 기존 블록에 교체
이미 만든 Transformer 블록에서 다음을 바꿉니다.
nn.LayerNorm(dim)→RMSNorm(dim)- FFN(Linear → GELU → Linear) →
SwiGLU(dim, hidden) - 입력에 더하던 위치 임베딩 제거 → Attention 안에서 Q, K를 헤드로 쪼갠 직후
apply_rope적용 (점수 계산Q @ K^T직전)
한 번에 하나씩 바꾸고 매번 학습 손실이 정상적으로 줄어드는지 확인하세요. 세 개를 동시에 바꾸면 문제가 생겼을 때 원인을 못 찾습니다.
오늘의 실험
질문: 절대 위치 임베딩과 RoPE 중 어느 쪽이 짧은 태스크에서 더 잘하는가? 학습 길이보다 긴 입력에서는 어떻게 되는가?
태스크 예시: 문자열 뒤집기, 또는 "숫자 열에서 n칸 전 숫자 복사하기"처럼 위치가 답을 좌우하는 작은 합성 태스크. 작은 모델(2~4층, dim 64~128)로 몇 분 안에 돌아가는 크기로 잡습니다.
통제 변수(고정): 데이터, 모델 크기, 학습 스텝, 학습률, 시드. 바꾸는 변수: 위치 방식 (절대 임베딩 / RoPE) 하나만.
관찰 항목:
- 학습 길이(예: 길이 32) 입력에서의 정확도와 최종 손실
- 학습 때 안 본 길이(예: 64)에서의 정확도: 이게 RoPE와 절대 임베딩의 차이가 드러나기 쉬운 지점입니다. 절대 임베딩은 본 적 없는 위치 번호가 아예 테이블에 없을 수도 있으니, 그 경우 최대 길이를 학습 길이로 설정해 "평가 불가"를 기록하세요.
- 시드를 3개 이상 돌려 평균과 편차 확인. 한 번의 결과로 결론 내지 마세요.
기록 방법: 아래 표를 노트에 채웁니다.
| 위치 방식 | 시드 | 학습 길이 정확도 | 긴 길이 정확도 | 최종 손실 |
|---|---|---|---|---|
| 절대 | 0 | |||
| RoPE | 0 |
여유가 되면 같은 방식으로 LayerNorm vs RMSNorm, GELU FFN vs SwiGLU도 각각 비교해 보세요. 작은 모델·짧은 학습에서는 차이가 작거나 노이즈에 묻힐 수 있습니다. 그 자체가 정직한 결과이니 "차이가 안 보였다"도 그대로 적으면 됩니다. 단, RoPE가 긴 길이에서 항상 완벽히 일반화한다고 기대하지 마세요. 결과가 기대와 달라도 이유를 추측해 적어 두는 것이 연구 습관입니다.
왜 이걸 배우나 (LLM / Post-training 연결)
오픈 LLM의 코드와 체크포인트를 열면 RoPE, RMSNorm, SwiGLU가 기본으로 들어 있어서, 이 부품을 읽을 수 있어야 모델 설정(config)과 가중치 이름(w1/w2/w3, norm 등)을 이해하고 수정할 수 있습니다. 이후 컨텍스트 길이를 늘리는 작업(RoPE 기반 길이 확장 기법들)이나 파인튜닝·Post-training 실험에서 "어느 층의 어떤 파라미터를 건드리는가"를 판단할 때 오늘 구조가 바로 기반이 됩니다.
초보자가 막히는 지점
- RoPE에서 shape 오류: head_dim이 홀수이면 2개씩 짝지을 수 없습니다. head_dim은 짝수여야 합니다. cos/sin의 shape은 (seq, head_dim/2)이고 x는 (batch, heads, seq, head_dim)이라, 브로드캐스트용 차원 추가(
[None, None])를 빼먹으면 에러가 납니다. - 짝짓기 방식 혼동: 인접 차원 (0,1), (2,3)을 짝짓는 구현과 앞 절반/뒤 절반을 짝짓는 구현이 있습니다. 둘 다 "회전"이지만 차원 배치가 달라서, 외부 체크포인트의 가중치를 불러올 때는 그 모델이 쓰는 방식과 반드시 맞춰야 합니다. 직접 처음부터 학습할 때는 하나로 일관되기만 하면 됩니다.
- V에도 RoPE를 적용하는 실수: Q와 K에만 적용합니다.
- RMSNorm 결과가 1이 안 나옴: g를 학습으로 바꾸고 나면 출력 크기는 달라집니다. 검증은 초기값(g=1)에서 하세요.
- SwiGLU 파라미터 수 때문에 비교가 불공정: 같은 hidden 크기로 비교하면 SwiGLU 쪽 파라미터가 1.5배 많습니다. 비교 시 hidden을 약 2/3로 줄여 파라미터 수를 맞추세요.
- 세 가지를 동시에 교체: 학습이 이상해졌을 때 원인을 못 찾습니다. 하나씩 교체하세요.
- float16에서 RMSNorm이 불안정: 제곱 평균에서 값이 넘칠 수 있으니, 필요하면
x.float()로 올려 계산한 뒤 원래 dtype으로 되돌립니다.
알고리즘 (병행 1시간)
LeetCode Medium 2문제를 DP(동적계획법, 큰 문제를 작은 문제로 나눠 결과를 저장하며 푸는 방법)로 풉니다. 문제 선택이 막막하면 "Coin Change", "Longest Increasing Subsequence"처럼 DP 입문 단계에서 자주 쓰이는 Medium 문제 중에서 고르세요.
문제마다 아래를 노트에 기록합니다.
- 접근: dp[i]가 무엇을 뜻하는지 한 문장으로 정의하고, 점화식(이전 값으로 현재 값을 구하는 식)과 초기값을 적기
- 시간·공간복잡도: 반복문 중첩으로 계산하고 Big-O로 적기
- 틀린 이유: 틀렸거나 막혔다면 원인(초기값 오류, 인덱스 범위, 상태 정의 모호함 등)을 한 줄로 남기기