모델 전체 흐름 다시 보기

모델 1편에서는 Embedding 층과 LM Head 사이에 디코더 블록을 놓고, 블록의 입력과 출력이 같은 d차원이라는 점을 살펴봤습니다. 이후에는 블록 안으로 들어가 Residual stream과 RMSNorm, Attention, MLP, RoPE, MoE를 차례로 살펴봤습니다. 이번 글에서는 각각 살펴본 구성 요소를 연결해 모델 전체 흐름을 정리하겠습니다. 벡터가 어떤 경로를 따라 변하고, 그중 어디에서 토큰 사이의 정보가 합쳐지는지 따라가 보겠습니다.

모델 전체에서 디코더 블록으로

토크나이저가 텍스트를 토큰 ID로 바꾸면, Embedding 층은 각 ID에 해당하는 d차원 벡터를 가져옵니다. 토큰 수가 T라면 이 벡터들을 행으로 쌓은 입력은 T × d 형태입니다. 이 입력이 N개의 디코더 블록을 차례로 통과하고, 마지막 RMSNorm과 LM Head를 거쳐 어휘별 점수인 logits가 됩니다.

그림 1의 오른쪽은 디코더 블록 하나의 내부입니다. 중앙의 Residual stream에서 옆 경로로 입력을 보내고, 그 경로에서 계산한 결과를 다시 더합니다. 첫 번째 경로에서는 RMSNorm을 거쳐 Attention을 적용하고, 두 번째 경로에서는 RMSNorm을 거쳐 MLP 또는 MoE를 적용합니다. RMSNorm은 각 토큰 벡터의 크기를 조절하며, Attention과 MLP는 정규화한 입력을 받아 변환을 계산합니다.

두 경로의 결과는 모두 d차원으로 돌아와 같은 토큰 위치의 Residual stream에 더해집니다. 블록을 통과해도 토큰 수 T와 모델 차원 d는 유지되지만, 벡터의 값은 달라집니다. 한 블록의 출력은 다음 블록의 입력이 되고, 이 과정이 N번 반복됩니다.

마지막 블록의 출력에도 RMSNorm을 적용한 뒤 LM Head로 전달합니다. LM Head는 각 토큰의 d차원 벡터를 Vocab size개의 점수로 변환하므로 전체 출력은 T × Vocab size가 됩니다. 디코더 블록 안에서는 각 토큰의 벡터를 변환하고, 마지막에는 그 벡터로 다음에 올 토큰의 점수를 계산하는 흐름입니다.

Attention의 연산을 하나로 연결하기

Attention의 입력에서는 Q·K·V projection으로 같은 토큰의 벡터를 세 가지 벡터로 변환합니다. Q와 K는 토큰 사이의 점수를 계산하는 데 사용하고, V는 그 점수에서 얻은 가중치에 따라 모을 정보를 담습니다. 그림 2는 앞선 글과 같은 MHA 구성으로, 토큰 수 T = 3, 모델 차원 d = 8, head 수 h = 2인 예시입니다.

Projection으로 만든 Q·K·V는 각각 3 × 8 형태입니다. 각 벡터의 여덟 성분을 두 head로 구분하면 head 하나의 차원 dₕ는 4가 되고, 각 head는 3 × 4 형태의 Q·K·V로 계산합니다. Q와 K에는 토큰 위치에 따른 RoPE를 적용합니다. 회전한 결과가 그림의 Q′와 K′이며, 차원은 그대로 유지됩니다. V는 회전하지 않고 Core Attention에 전달합니다.

Core Attention에서는 먼저 Q′K′ᵀ를 계산하고 √dₕ로 나눕니다. 토큰 세 개를 서로 비교하므로 점수 행렬은 3 × 3입니다. 여기에 Causal mask로 참조할 위치를 제한하고, 각 행에 Softmax를 적용해 가중치 P를 만듭니다. 마지막으로 PV를 계산해 각 토큰이 참조하는 Value를 가중합합니다. P는 3 × 3, V는 3 × 4이므로 head 하나의 결과는 3 × 4입니다.

두 head의 결과는 같은 토큰 위치끼리 이어 붙입니다. 각 토큰에 네 성분씩 두 묶음이 있으므로 전체 결과 H는 다시 3 × 8이 됩니다. 출력 projection은 같은 토큰의 여러 head 결과를 학습된 가중치로 조합합니다. 이렇게 얻은 3 × 8 결과가 그림 1의 첫 번째 Residual 덧셈으로 돌아갑니다.

토큰별 변환: MLP와 MoE

Attention의 결과를 Residual stream에 더한 다음에는 RMSNorm을 거쳐 MLP를 적용합니다. Dense MLP에서는 모든 토큰에 같은 MLP를 적용하고, MoE에서는 토큰마다 사용할 expert를 선택합니다. 그림 3은 토큰 하나의 입력 x를 기준으로 두 구조를 나란히 보여줍니다.

왼쪽 Gated MLP에서는 같은 d차원 입력을 Gate projection과 Up projection에 보냅니다. 두 경로는 각각 m차원 벡터를 만들고, Gate 경로에는 SiLU를 적용합니다. 두 결과를 같은 성분끼리 곱하면 Gate 경로의 값이 Up 경로의 성분을 조절합니다. Down projection은 이 m차원 결과를 다시 d차원으로 변환합니다.

오른쪽 MoE에서는 Router가 입력 x를 사용해 expert와 가중치를 정합니다. 그림에서는 E1과 E3을 선택했고, 두 expert에도 같은 x를 입력합니다. 각 expert는 왼쪽과 같은 Gated MLP이지만 서로 다른 가중치를 사용합니다. Combine에서는 두 expert의 d차원 결과에 Router가 정한 가중치를 곱해 더합니다. 여러 expert를 거치더라도 합치는 것은 같은 토큰에서 나온 결과입니다.

일부 모델은 점선으로 표시한 Shared expert도 함께 둡니다. Shared expert는 Router의 선택과 관계없이 모든 토큰에 적용하며, 그 결과를 같은 토큰의 routed expert 결과에 더합니다. Shared expert가 있는 경우에도 최종 출력은 d차원입니다. Dense MLP와 MoE 모두 이 출력을 그림 1의 두 번째 Residual 덧셈으로 전달합니다.

토큰 사이의 정보는 어디서 합쳐질까요?

이제 디코더 블록 전체를 토큰 위치별로 나누어 보겠습니다. 그림 4의 세 열은 토큰 1·2·3의 경로입니다. 지금까지 살펴본 구조에서 서로 다른 토큰의 벡터를 비교하고 정보를 결합하는 곳은 Core Attention입니다. 다른 단계에서는 같은 토큰 위치 안에서 계산이 이어집니다.

Core Attention에 들어가기 전까지 RMSNorm과 Q·K·V projection은 토큰별로 계산합니다. RoPE도 각 토큰의 Q·K와 해당 위치를 사용해 회전합니다. 이후 Core Attention에서는 Causal mask가 허용하는 범위에 따라 토큰 1은 자기 자신의 Value를, 토큰 2는 토큰 1·2의 Value를, 토큰 3은 토큰 1·2·3의 Value를 모읍니다. 그림에서 여러 토큰의 연결선이 한 위치의 결과로 모이는 부분입니다.

Core Attention 이후에도 서로 다른 토큰 위치의 결과를 합치지는 않습니다. Concat과 출력 projection은 같은 토큰의 head 결과를 조합하고, MoE의 Combine은 같은 토큰의 expert 결과를 조합합니다. Residual 덧셈도 같은 위치의 벡터끼리 더합니다. 이렇게 계산을 이어가면 블록의 끝에는 다시 각 토큰 위치에 대응하는 d차원 벡터가 남습니다.

모델 전체를 보면 Residual stream이 블록들을 지나 이어지고, Attention과 MLP 또는 MoE의 변환 결과가 여기에 더해집니다. 그 안에서 Core Attention은 토큰 사이의 정보를 결합하고, 나머지 연산은 각 위치의 벡터를 처리합니다. 하드웨어 편에서는 이 흐름을 다시 보면서, 각 계산을 어떻게 병렬로 나누고 데이터를 어디에서 읽고 쓰는지 살펴보겠습니다.