디코더 블록의 기본 흐름: Residual과 RMSNorm

지난 글에서는 LLM의 전체 구조를 임베딩 층, 디코더 블록, LM Head로 나누어 살펴봤습니다. 디코더 블록에는 T개 토큰의 d차원 벡터가 들어가고, 같은 T × d 형태의 벡터가 나온다는 점도 확인했습니다. 이번 글에서는 하나의 상자로 두었던 디코더 블록의 내부를 살펴보겠습니다. 먼저 벡터가 이어지는 기본 경로와, 그 경로에서 갈라졌다가 다시 합쳐지는 계산을 구분해 보겠습니다.

Residual 흐름을 중심으로 보는 디코더 블록

디코더 블록은 크게 Residual stream, RMSNorm, Attention, MLP로 나누어 볼 수 있습니다.

디코더 블록의 구조를 그릴 때는 보통 아래 그림의 왼쪽처럼 Attention과 MLP를 가운데에 놓고, 입력을 우회시켜 다시 더하는 선을 바깥쪽에 그립니다. 이번에는 오른쪽처럼 입력 벡터가 이어지는 경로를 가운데에 두고, Attention과 MLP를 옆 경로에 놓아보겠습니다. 이 가운데 경로를 Residual stream이라고 합니다. 각 분기에서 입력은 가운데 경로로 그대로 전달되고, 옆 경로에서는 그 입력으로 변환 결과를 계산합니다. 두 경로가 만나는 지점에서 변환 결과를 더하고, 더한 벡터가 다시 다음 계산으로 이어집니다.

오른쪽 그림을 위에서부터 따라가면 첫 번째 옆 경로에서는 RMSNorm과 Attention을 거친 결과가 더해지고, 두 번째 옆 경로에서는 RMSNorm과 MLP를 거친 결과가 더해집니다. Attention과 MLP가 중심 변환을 담당하고, 그 결과가 Residual stream에 더해지는 구조입니다. Attention과 MLP는 위치만 확인할 수 있도록 옅은 상자로 표시했습니다. 이번 글에서는 Residual stream과 옆 경로에 있는 RMSNorm부터 살펴보겠습니다.

Residual: 입력 벡터에 변환 결과 더하기

Residual 연결에서 하는 계산은 그대로 전달된 입력에 변환 결과를 더하는 것입니다. 입력을 x, 옆 경로에서 계산한 결과를 Δ라고 하면 출력은 x + Δ입니다. 벡터의 덧셈은 같은 위치의 성분끼리 이루어집니다. 첫 번째 성분은 첫 번째 성분끼리, 두 번째 성분은 두 번째 성분끼리 더하는 방식입니다. 따라서 입력과 더할 결과는 같은 차원이어야 합니다. 토큰 하나의 입력이 d차원이라면 변환 결과도 d차원이고, 둘을 더한 출력도 d차원입니다.

그림에서는 토큰 하나의 벡터를 확대했습니다. 입력이 [1, 2, −1, 0]이고 변환 결과가 [0.2, −0.5, 0.3, 0.1]이라면, 출력은 [1.2, 1.5, −0.7, 0.1]이 됩니다. 예를 들어 두 번째 성분은 2 + (−0.5) = 1.5로 계산됩니다. 입력은 덧셈 직전까지 같은 값으로 전달되지만, 합류한 뒤에는 더한 결과가 다음 경로를 따라갑니다. 이처럼 차원은 유지하면서 각 성분의 값이 달라집니다. 전체 토큰을 함께 보면 두 T × d 배열을 위치별로 더하는 계산입니다. 그림의 수치는 설명용 예시이며, Attention은 변환 결과를 계산할 때 다른 토큰의 벡터도 참조할 수 있습니다.

RMSNorm: 토큰 하나의 d개 성분으로 크기 조절하기

옆 경로로 들어간 벡터는 Attention이나 MLP에 전달되기 전에 RMSNorm을 거칩니다. RMS는 Root Mean Square의 약자로, 각 성분을 제곱한 뒤 평균을 구하고 제곱근을 취한 값입니다. 한 토큰의 벡터에 d개의 성분이 있다면, 이 d개 성분으로 RMS 하나를 구합니다. 그런 다음 각 성분을 같은 RMS로 나누어 벡터의 크기를 조절합니다.

RMSNorm에는 여기에 한 단계가 더 있습니다. RMS로 나눈 결과에 학습 가능한 가중치 벡터 γ(감마)를 성분별로 곱합니다. 입력이 d차원이면 γ도 d개의 값을 가지며, 각 값은 학습 과정에서 조정됩니다. 아래 그림에서는 먼저 RMS로 나누는 변화를 보기 위해 γ를 [1, 1, 1, 1]로 설정했습니다. 1을 곱하면 값이 그대로이므로, 이 예시의 출력은 RMS로 나눈 결과와 같습니다.

입력이 [1, −1, 3, 3]이면 제곱한 값의 합은 1 + 1 + 9 + 9 = 20입니다. 성분 수 d가 4이므로 제곱의 평균은 5이고, RMS는 √5 ≈ 2.24가 됩니다. 각 성분을 √5로 나누면 약 [0.45, −0.45, 1.34, 1.34]가 나옵니다. 그림의 막대는 좌우가 같은 눈금이므로 각 값의 크기가 어떻게 달라졌는지 비교할 수 있습니다. 이 계산을 간단히 보여주기 위해 작은 값 ε(엡실론)는 생략했습니다. 실제 식에서는 분모가 0이 되는 것을 막기 위해 제곱근 안에 ε를 더합니다.

여기서 RMS를 구하는 범위는 토큰 하나의 d개 성분입니다. T개 토큰이 있으면 각 토큰의 벡터에서 RMS를 따로 구합니다. 같은 RMSNorm 층은 모든 토큰에 같은 γ를 사용하지만, 나누는 RMS 값은 각 토큰의 벡터에 따라 달라집니다. 성분의 개수는 바뀌지 않으므로 입력과 출력은 모두 T × d 형태입니다. 그림 1에서는 이렇게 정규화한 값을 Attention이나 MLP에 전달하고, 가운데 경로에는 정규화하기 전의 값을 유지해 두었다가 변환 결과를 더합니다.

전체 구조에서 다시 보는 RMSNorm

지난 글의 전체 구조에 RMSNorm의 위치를 추가해 보겠습니다. 디코더 블록 안에서 살펴본 RMSNorm 외에도, 마지막 디코더 블록의 출력과 LM Head 사이에 RMSNorm이 한 번 더 놓입니다. 마지막 블록에서 나온 각 토큰의 d차원 벡터를 정규화한 뒤, LM Head가 어휘별 점수로 변환합니다.

이제 디코더 블록을 Residual stream과 그 옆의 변환 경로로 나누어 볼 수 있습니다. Residual stream은 각 분기의 입력을 그대로 전달하고, 합류 지점에서 변환 결과를 더해 다음 단계로 보냅니다. 각 옆 경로에서는 RMSNorm을 거쳐 Attention이나 MLP의 변환 결과를 계산합니다. 다음 글에서는 옅은 상자로 두었던 Attention과 MLP의 역할을 비교하고, MLP 내부에서 어떤 계산이 이루어지는지 살펴보겠습니다.