LLaMA의 5가지 특징

1. Introduction

1.1 목적 및 개요

• Transformer 모델은 다양한 자연어 처리(NLP) 작업에서 중요한 기반 아키텍처로 자리 잡고 있음.
• LLaMA는 Transformer 기반의 대규모 언어 모델(LLM)로, 성능과 효율성을 동시에 개선한 구조임.
• Vanilla Transformer와 LLaMA의 주요 아키텍처 차이를 분석하고 각 구성 요소의 설계 및 성능 개선 이유를 분석

1.2 Transformer와 LLaMA의 배경

1.2.1 Transformer 모델의 기본 구조

• Transformer는 2017년 Vaswani et al.의 논문 "Attention Is All You Need"에서 제시됨.
• Self-Attention 메커니즘과 **Multi-Head Attention (MHA)**을 통해 병렬 처리와 장기 의존성(long-term dependency) 학습이 가능함.
• 기본 구조는 인코더(Encoder)와 디코더(Decoder)로 구성되어 있으며, 주요 컴포넌트로 Layer Normalization, Feed-Forward Network (FFN), Positional Encoding 등이 있음.

1.2.2 LLaMA 모델의 개요와 발전 배경

• LLaMA는 Meta AI가 발표한 대규모 언어 모델로, GPT, PaLM 등 기존 LLM 대비 성능과 자원 효율성이 개선됨.

• Transformer의 기본 구조를 기반으로 하되, 다음과 같은 최적화 기법을 도입하여 학습 및 추론 성능을 향상시킴.

  • Pre-Normalization
  • RMSNorm
  • SwiGLU - Activation function
  • Rotary Position Embedding (RoPE)
  • Grouped-Query Attention (GQA)

1.3 목표

• Transformer 아키텍처의 다양한 변형 기법들이 존재하지만 각 기법의 성능 개선 이유와 한계를 체계적으로 비교하기가 어려움
• Transformer와 LLaMA의 아키텍처 차이를 주요 5가지 요소로 나누어 분석.
• 각 기법의 이론적 근거 및 설계 동기를 설명하고, 실험 결과를 통해 성능 및 효율성 비교를 수행.

---

2. Pre-Layer Normalization 구조 비교

Source: On Layer Normalization in the Transformer Architecture

2.1 Layer Normalization 개요 및 역할

2.1.1 개념 및 정의

• **Layer Normalization(LayerNorm)**은 입력 벡터의 각 차원에 대해 평균과 분산을 계산하여 정규화하는 기법임.

  • 입력 벡터에서 각 차원의 불균형한 분포를 완화하여 학습 안정성과 수렴 속도를 향상시킴.
  • 다양한 입력 데이터에서도 모델이 일관된 결과를 내도록 도와줌.

2.1.2 역할

• **기울기 폭발(Exploding Gradient)**과 기울기 소실(Vanishing Gradient) 문제를 완화하여 깊은 네트워크에서 안정적인 학습이 가능하도록 함.
• Transformer와 같은 복잡한 네트워크 구조에서 각 레이어의 출력을 조절하는 데 중요한 역할을 수행함.

2.2 Pre-Layer Normalization 구조의 작동 방식

2.2.1 작동 방식

• Pre-Layer Normalization(Pre-LN)은 FFN 또는 Self-Attention 연산 전에 LayerNorm을 적용하는 방식임.

  • 입력 벡터가 정규화된 상태로 각 연산에 전달됨.

2.2.2 잔차 연결과 LayerNorm의 위치 변화

• 기존 Post-LN에서는 잔차 연결 이후에 정규화가 이루어졌으나, Pre-LN에서는 잔차 연결 전 내부에서 정규화가 수행됨.
• 이로 인해 입력값의 스케일이 균일하게 유지되고, 각 레이어가 더 안정적으로 학습할 수 있음.

2.3 Pre-LN 구조가 잘 작동하는 원인

2.3.1 기울기 전파 안정성 개선

• LayerNorm이 연산 전에 적용됨으로써 기울기 소실 현상이 완화됨.

  • 각 층에서 기울기 값이 안정적으로 유지되어 깊은 네트워크에서도 학습이 원활하게 진행됨.
  • 이로 인해 초기 학습뿐만 아니라 후반부 학습에서도 기울기 흐름이 보존됨.

2.3.2 학습 속도와 안정성 향상

• 모든 레이어가 정규화된 입력을 받기 때문에 학습 과정에서 입력 데이터의 변화에 덜 민감해짐.
• 학습률 워밍업(warm-up) 단계를 생략할 수 있으며, 초기 학습 단계에서 안정적인 손실 감소가 가능함.

2.3.3 이론적 근거

• 평균장 이론(mean field theory)에 따르면, Pre-LN 구조는 각 층에서 경사가 일정하게 분포되어 경사 폭발 및 소실을 방지함.
• 모든 층에서 입력이 정규화되어, 깊은 네트워크일수록 균형 잡힌 학습 신호가 전파됨.

2.4 Pre-LN 구조의 단점

2.4.1 초기 학습 속도 저하

• 입력이 지나치게 정규화되면서 초기 학습 단계에서 모델이 새로운 패턴을 빠르게 학습하기 어려운 경우가 있음.
• 특정 입력에 대한 민감도가 줄어들어 학습 초기 손실 수렴 속도가 느려질 수 있음.

2.4.2 일부 레이어의 영향력 감소

• 각 레이어가 정규화된 입력을 받기 때문에 특정 레이어의 변화가 출력에 미치는 영향이 제한적일 수 있음.
• 이로 인해 레이어 간 상호작용이 다소 약해지는 현상이 나타날 수 있음.

2.5 실험 결과 및 성능 분석

• Pre-LN 구조는 깊은 네트워크에서도 기울기 전파가 원활하여 대규모 모델에서 안정적인 성능을 보임.
• 워밍업 스케줄이 필요하지 않기 때문에 학습 시간 단축에 기여함.
• 다양한 실험에서 Post-LN 구조 대비 더 빠른 수렴 속도와 높은 성능이 확인됨.

2.6 Pre-LN 구현

import torch
import torch.nn as nn

class TransformerLayerPreNorm(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super().__init__()
        self.layer_norm = nn.LayerNorm(hidden_dim)
        self.attention = nn.MultiheadAttention(hidden_dim, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 4),
            nn.ReLU(),
            nn.Linear(hidden_dim * 4, hidden_dim)
        )

    def forward(self, x):
        x_norm = self.layer_norm(x)
        attn_output, _ = self.attention(x_norm, x_norm, x_norm) # (output, weight)
        x = x + attn_output          # Residual connection
        x_norm = self.layer_norm(x)
        x = x + self.ffn(x_norm)     # Residual connection
        return x

---

3. RMSNorm: 효율적인 정규화 기법

Source: Root Mean Square Layer Normalization

3.1 RMSNorm 개요 및 배경

3.1.1 정의

• RMSNorm은 LayerNorm에서 Re-centering 단계를 제거하고 Re-scaling을 기반으로 정규화하는 기법.
• RMSNorm은 입력 벡터의 평균을 계산하지 않고도 스케일 조정을 통해 학습 안정성을 제공함.

3.1.2 기존 LayerNorm의 한계

• LayerNorm은 평균과 분산을 계산하는 데 연산 비용이 크며, 특히 대규모 신경망에서 속도와 메모리 사용량 문제가 발생.
• 평균과 분산 계산이 추가 메모리 대역폭을 소모하여 훈련 시간 증가와 병렬 연산 비효율을 초래함.

3.2 RMSNorm의 작동 방식

3.2.1 정규화 수식

• 여기서 g는 학습 가능한 파라미터로, 출력의 스케일을 조정하는 역할을 함.
• 평균 계산이 생략되어 Re-scaling Invarinace만 유지되며, 평균 이동에 대한 불변성은 없음.

3.2.2 작동 방식 요약

• 입력 벡터의 각 요소는 RMS 값을 기준으로 정규화됨.
• LayerNorm과 달리 입력의 평균을 고려하지 않기 때문에 연산이 단순화되어 속도와 메모리 효율이 향상됨.
• 정규화 후 학습 가능한 게인 g가 적용되어 출력 값의 크기를 조정함.

3.3 RMSNorm의 장점 및 효과

3.3.1 연산 효율성 개선

• 평균 계산 단계가 제거되어 LayerNorm 대비 최대 64%의 속도 향상을 달성함.
• 대규모 모델에서 연산 병목이 줄어들고, 메모리 대역폭 사용이 최적화됨.
• 연산 복잡도가 줄어들어 GPU/TPU 병렬 처리 성능이 향상됨.

3.3.2 학습 안정성

• 재스케일링 불변성으로 인해 입력 스케일의 변화가 학습 성능에 미치는 영향이 최소화됨.
• 초기화 조건이 비정상적이거나 입력이 극단적으로 변화하는 상황에서도 안정적인 학습이 가능함.

3.3.3 다양한 모델에서의 적용성

• RNN, Transformer, CNN 등 여러 신경망 아키텍처에서 실험적으로 성능 개선이 확인됨.
• BLEU 점수와 같은 평가 지표에서 LayerNorm과 동등한 성능을 보이면서도 더 빠른 수렴을 달성함.
• Transformer 아키텍처에서 학습 실패율을 줄이고, 테스트 정확도에서도 우수한 결과를 보임.

3.4 실험 결과 요약

• 속도: LayerNorm 대비 연산 효율이 향상되어 훈련 시간이 단축됨.
• 메모리 사용: LayerNorm보다 적은 메모리를 소모하여 대규모 모델에서 더 효율적으로 동작함.
• 성능: 다양한 실험에서 LayerNorm과 유사한 성능을 유지하며, 특히 초대형 모델에서 안정성을 보임.

3.5 결론

• RMSNorm은 LayerNorm의 핵심 기능 중 하나인 Re-scaling Invariance에 중점을 두어 설계된 기법임.
• LayerNorm 대비 연산 효율성을 높이면서도 성능 저하가 없으며, 대규모 모델에서 속도-안정성의 균형을 효과적으로 맞춤.
• 이러한 특성 덕분에 RMSNorm은 실무 환경에서 LayerNorm을 대체하는 데 유용한 선택지로 평가됨.

3.6 RMSNorm 구현

class RMSNorm(nn.Module):
    def __init__(self, hidden_dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(hidden_dim))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
        return self.scale * (x / rms)

---

4. SwiGLU: Feed-Forward Network 활성 함수

Source: GLU Variants Improve Transformer

4.1 작동 방식

• SwiGLU는 GLU(Gated Linear Units) 계열 활성 함수의 변형으로, 두 개의 선형 변환 결과 중 하나에 Swish 활성 함수를 적용하여 구성됨.

• 입력 벡터 x에 대해 두 개의 선형 변환 W와 V가 존재하며, 다음과 같은 연산이 이루어짐:

• 두 번째 선형 변환(V) 은 필터링된 정보를 강조하거나 억제하는 역할을 수행함.

• element-wise multiplication을 수행함으로써 입력에 대해 필요한 정보를 동적으로 필터링

4.2 기법의 특징 및 장점

4.2.1 더 나은 정보 흐름 관리

• Swish 함수는 미분 가능한 점진적 활성화 성질을 가지므로, 비선형성과 학습 안정성이 개선됨.
• 일반적인 ReLU나 GELU와 달리, 정보 손실이 최소화되며 더 세밀한 가중치 조정이 가능함.

4.2.2 동적 게이팅 메커니즘

• 첫 번째 선형 변환에서 Swish 활성화가 중요한 정보의 가중치를 조절함.
• 이를 통해 중요한 정보만 선택적으로 통과시켜, 피드포워드 네트워크에서 효율적인 정보 처리가 이루어짐.

4.2.3 효율성과 성능의 균형

• GLU 계열 변형들 중 SwiGLU가 여러 언어 이해 및 생성 태스크에서 가장 낮은 퍼플렉시티(perplexity)를 기록함.
• 기존 Transformer FFN 구조 대비 적은 연산량으로 비슷하거나 더 나은 성능을 발휘함.

4.3 다른 활성 함수와의 비교

4.3.1 ReLU와의 비교

• ReLU (Rectified Linear Unit):

  • 장점: 연산이 간단하여 매우 빠름.
  • 단점: 음수 값을 모두 제거하기 때문에 죽은 노드(dead neurons) 문제와 정보 손실이 발생할 수 있음.
  • ReLU는 SwiGLU에 비해 세밀한 정보 조정이 어려움.

4.3.2 GELU와의 비교

• GELU (Gaussian Error Linear Unit):

  • 장점: ReLU 대비 더 부드러운 비선형성을 제공하여 학습 안정성 향상.
  • 단점: 연산이 ReLU에 비해 다소 복잡하여 계산 비용이 증가함.
  • GELU는 SwiGLU와 함수의 형태가 비슷한 모습을 보이지만, SwiGLU가 동적 게이팅 메커니즘 덕분에 정보 선택적 필터링에서 더 나은 성능을 제공함.

4.3.3 SwiGLU의 상대적 우위

• SwiGLU는 Swish 활성 함수의 부드러운 비선형성과 GLU 구조의 게이팅 효과를 결합하여, ReLU나 GELU 대비 정보 흐름 관리에서 뛰어남.
• 특히 대규모 Transformer 모델에서 perplexity 감소, 훈련 안정성 향상, 정보 손실 최소화 효과가 검증됨.

4.4 성능 실험 결과

• SwiGLU는 T5(Text-to-Text Transfer Transformer) 모델에서 기존 ReLU 기반 FFN보다 낮은 ppl.
• 여러 언어 이해 벤치마크(GLUE, SuperGLUE 등)에서도 높은 성능을 보이며, 특히 문맥적 관계 파악에 유리함을 입증.

4.5 한계 및 적용 조건

4.5.1 추가 가중치 증가

• SwiGLU는 두 개의 선형 변환을 사용하므로, 모델 파라미터 수가 증가하여 메모리 사용량이 늘어날 수 있음.
• 메모리 제약이 있는 환경에서는 SwiGLU 사용이 비효율적일 수 있음.

4.5.2 최적화 전략

• 내부 차원을 절반으로 줄이는 방식이 권장됨.
• 이 방법은 파라미터 수 증가로 인한 메모리 문제를 완화하는 데 기여함.

4.5.3 적용 환경

• SwiGLU는 대규모 신경망(Transformer, GPT 등)에서 더 효과적임.
• 소규모 네트워크에서는 ReLU 또는 GELU가 더 적합할 수 있음.

4.6 결론

• SwiGLU는 Transformer 아키텍처의 피드포워드 네트워크에서 정보 필터링과 성능 최적화에 중요한 역할을 수행함.
• 효율적인 학습과 동적 게이팅 메커니즘을 통해, 다양한 NLP 태스크에서 높은 성능을 입증함.
• 다만, 메모리 및 연산 비용을 고려한 최적화가 필요함.

4.7 SwiGLU 구현

import torch
import torch.nn as nn

class SwiGLU(nn.Module): # Activation function 이므로  FFN 부분은 제외하고 구현
    def __init__(self, input_dim, hidden_dim):
        super(SwiGLU, self).__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.linear2 = nn.Linear(input_dim, hidden_dim)

    def forward(self, x):
        linear_out2 = self.linear2(x)
        swish_out = linear_out2 * torch.sigmoid(linear_out2) # swish: x * sigmoid(x)

        # Element-wise mul
        return self.linear1(x) * swish_out

---

5. RoPE (Rotary Position Embedding)

Source: ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING

5.1 기존 방식의 한계

5.1.1 Absolute Positional Encoding

• 작동 방식:

  • 토큰의 절대적 위치를 임베딩 벡터로 변환하여 입력 임베딩에 더하는 방식.

• 문제점:

  • 절대적 위치만 고려하여 상대적 위치 정보가 반영되지 않음.
  • 짧은 거리에 있는 토큰들 간의 관계는 잘 모델링되지만, 긴 시퀀스에서 장거리 의존성을 학습하는 데 어려움이 있음.
  • 위치 정보를 단순히 더하는 방식은 상대적 거리 변화에 따른 의존성 변화를 효과적으로 반영하지 못함.

5.2 RoPE의 작동 방식

• RoPE는 회전 변환(rotation transformation)을 통해 상대적 위치 정보를 인코딩함.

• Query와 Key벡터에 회전 행렬 ( R )을 적용하여 위치 정보를 삽입함.

• 직관적 이해

  • 벡터가 ( m * theta )만큼 회전함으로써 위치 정보를 포함함.
  • 이 회전 변환은 토큰들 간의 상대적 위치 차이를 내적에 직접 반영함.

5.3 장점 및 특징

5.3.1 상대적 위치 정보의 효과적 반영

• RoPE는 Query-Key 내적에서 상대적 거리 차이를 직접 반영하여 장거리 관계 모델링이 향상됨.
• 기존 절대적 위치 임베딩 방식의 한계를 극복하고, 상대적 위치 의존성을 강화함.

5.3.2 학습 안정성과 계산 효율성

• 회전 행렬은 정규 직교 행렬(orthogonal matrix)로 정의되어 수치적 안정성을 제공함.
• 벡터의 크기(norm)가 보존되어 학습이 안정적으로 수렴함.
• 회전 행렬이 희소(sparse) 구조를 가지므로 효율적인 행렬 연산이 가능함.

5.5 한계 및 적용 조건

1. 짧은 시퀀스에서의 효과 제한

   • RoPE는 긴 시퀀스에서 장거리 관계 모델링에 최적화되어 있음.
   • 짧은 시퀀스에서는 절대적 위치 임베딩 방식과 큰 성능 차이가 나타나지 않을 수 있음.

2. 모델 복잡도 증가

   • 회전 행렬 연산이 추가되지만, 고차원에서는 희소 행렬 최적화를 통해 계산 효율성을 유지할 수 있음.

5.6 결론

• RoPE는 Transformer 모델에서 상대적 위치 정보를 효과적으로 반영하며, 장거리 의존성을 요구하는 NLP 태스크에서 우수한 성능을 발휘함.
• 기존 절대적 위치 임베딩 방식의 한계를 극복하며, 긴 시퀀스 처리에 적합한 솔루션을 제공함.
• 상대적 위치 임베딩 기법과의 비교에서도 계산 효율성과 성능 면에서 장점을 보임.

5.7 RoPE 구현

import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=4096, base=10000):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

        # 사전 계산된 cos & sin caching
        position = torch.arange(0, max_seq_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, dim, 2, dtype=torch.float32) * -(torch.log(torch.tensor(base)) / dim))
        freqs = position * div_term
        self.register_buffer('cos_cache', torch.cos(freqs))
        self.register_buffer('sin_cache', torch.sin(freqs))

    def forward(self, x, seq_dim=1):
        # x: (batch_size, seq_len, num_heads, head_dim)
        seq_len = x.size(seq_dim)
        cos = self.cos_cache[:seq_len, :].unsqueeze(1)  # (seq_len, 1, dim)
        sin = self.sin_cache[:seq_len, :].unsqueeze(1)  # (seq_len, 1, dim)

        # 짝수 및 홀수 인덱스 분할
        x1, x2 = x[..., ::2], x[..., 1::2]

        # 회전 변환 적용
        x_rotated = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
        return x_rotated

---

6. Grouped-Query Attention (GQA)

Source: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

6.1 기존 방식의 한계

6.1.1 Multi-Head Attention (MHA)

• 작동 방식:

  • 여러 Query, Key, Value 헤드를 사용하여 다양한 정보 표현을 강화함.
  • 각 토큰의 생성 단계마다 모든 헤드의 Key-Value 데이터를 반복적으로 불러와야 함.

• 문제점:

  • Key-Value 캐시의 크기가 커져 메모리 대역폭 부담이 증가함.
  • 대규모 모델에서는 디코더 단계에서 메모리와 계산 자원 소모가 과도함.

6.1.2 Multi-Query Attention (MQA)

• 작동 방식:

  • 여러 Query 헤드를 유지하면서 단일 Key와 Value 헤드만 사용함.
  • Key-Value 데이터 로딩에 필요한 메모리 대역폭을 크게 줄임.

• 문제점:

  • 모델 품질이 저하될 가능성이 있음.
  • 일부 태스크에서 학습 안정성이 떨어짐.
  • 성능과 추론 속도의 균형을 맞추기 어려움.

6.2 GQA의 작동 방식

• **Grouped-Query Attention(GQA)**는 Query 헤드를 여러 그룹으로 나누어 각 그룹이 단일 Key와 Value 헤드를 공유하는 방식.

• GQA-G에서 G는 그룹 수를 의미하며, 다음과 같은 변환 가능성을 제공함:

  • G = 1일 때, MQA와 동일한 구조.
  • G = 전체 Query 헤드 수일 때, MHA와 동일한 구조.

변환 과정

• 기존의 Multi-Head 체크포인트에서 각 그룹의 Key와 Value 헤드를 mean pooling을 통해 결합하여 새로운 Key와 Value를 생성함.
• 추가적으로 업트레이닝(uptraining) 단계를 통해 새로운 구조에 적응시킴.

6.3 장점 및 특징

1. 효율적인 메모리 사용

   • Key-Value 캐시 크기를 줄여 메모리 대역폭을 최적화함.
   • 모델 크기가 증가할수록 메모리 대역폭 최적화 효과가 극대화됨.

2. 성능-속도 절충

   • MQA보다 높은 품질을 유지하면서도 MHA에 비해 빠른 추론 속도를 제공함.
   • 다양한 벤치마크 실험에서 MHA와 유사한 품질을 달성하면서도 MQA에 근접한 추론 속도를 보임.

3. 모델 확장성

   • 대규모 모델에서 메모리 및 연산 자원의 효율성을 극대화할 수 있음.
   • 모델 파티셔닝 시 불필요한 Key-Value 복제를 방지하여 리소스 낭비를 줄임.

6.4 성능 비교

기법	특성	메모리 대역폭 사용	품질	추론 속도
Multi-Head Attention	여러 Key-Value 헤드 유지	높음	높음	느림
Multi-Query Attention	단일 Key-Value 헤드 사용	낮음	낮음	빠름
Grouped-Query Attention	그룹 단위 Key-Value 사용	중간	높음	중간

6.5 한계 및 적용 조건

1. 초기 변환 및 추가 학습 필요

   • 기존 모델을 GQA로 변환한 후 추가 학습이 필요함.
   • 일부 작업에서는 초기 품질 변동이 발생할 수 있음.

2. 적용 범위 제한

   • GQA는 Decoder의 Self-Attention 및 Cross-Attention에만 적용 가능함.
   • Encoder의 Self-Attention에는 적용되지 않음.
   • 주로 긴 시퀀스를 생성하는 작업에 적합함.

6.6 결론

• GQA는 MHA와 MQA의 장점을 결합하여 메모리 대역폭 최적화와 모델 품질 유지를 동시에 달성함.
• 대규모 언어 모델에서 디코더 병목 문제를 해결하는 데 효과적인 솔루션으로 자리 잡고 있음.
• 앞으로 대규모 태스크에서 효율성을 더욱 향상시킬 가능성을 제시함.

6.7 GQA 구현

import torch
import torch.nn as nn
import torch.nn.functional as F

def scaled_dot_product_gqa(self, q, k, v, attn_mask=None):
    # q: (batch_size, num_heads, seq_len, hidden_dim)
    # k, v: (batch_size, num_groups, seq_len, hidden_dim)

    batch_size, num_heads, seq_len, hidden_dim = q.size()

    q = q / torch.sqrt(torch.tensor(hidden_dim))

    '''
    GQA의 특징: 그냥 내적을 할 경우 q,k의 차원이 맞지 않아서 matmul이 안됨
    torch.einsum을 활용해서 브로드캐스팅 후 연산
    '''

    attn_scores = torch.einsum("bhqd,bgkd->bhqg", query, key)

    if attn_mask is not None:
        attn_scores = attn_scores.masked_fill(attn_mask == 0, float("-inf"))

    attn_probs = F.softmax(attn_scores, dim=-1)
    attn_probs = F.dropout(attn_probs, p=self.dropout, training=self.training)

    attn_output = torch.einsum("bhqg,bgkd->bhqd", attn_probs, v)

    return attn_output, attn_probs

class GroupedQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, num_groups, dropout=0.0, bias=True):
        super(GroupedQueryAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_groups = num_groups  # 그룹 수 (GQA의 핵심 요소)
        self.dropout = dropout

        if embed_dim % num_heads != 0:
            raise ValueError("embed_dim must be divisible by num_heads")
        if num_heads % num_groups != 0:
            raise ValueError("num_heads must be divisible by num_groups")

        self.hidden_dim = embed_dim // num_heads
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.k_proj = nn.Linear(embed_dim, embed_dim // num_groups, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim // num_groups, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

    def forward(self, query, key, value, attn_mask=None):
        batch_size, seq_len, embed_dim = query.size()

        q = self.q_proj(query).view(batch_size, seq_len, self.num_heads, self.hidden_dim)
        k = self.k_proj(key).view(batch_size, seq_len, self.num_groups, self.hidden_dim)
        v = self.v_proj(value).view(batch_size, seq_len, self.num_groups, self.hidden_dim)

        q = q.transpose(1, 2)  # (batch_size, num_heads, seq_len, hidden_dim)
        k = k.transpose(1, 2)  # (batch_size, num_groups, seq_len, hidden_dim)
        v = v.transpose(1, 2)  # (batch_size, num_groups, seq_len, hidden_dim)

        attn_output, attn_output_weights = self.scaled_dot_product_gqa(q, k, v, attn_mask)

        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        attn_output = self.out_proj(attn_output)

        return attn_output, attn_output_weights

---

7. 결론

7.1 결론 요약

1. Pre-Norm vs Post-Norm

   • Pre-Norm이 안정적인 학습과 더 빠른 수렴을 제공하며, 깊은 모델에서도 효과적으로 작동함을 확인함.

2. RMSNorm의 활용

   • LayerNorm 대비 연산 효율성이 높아 메모리와 속도 면에서 개선을 제공하며, 대규모 모델에서 더 나은 확장성을 보임.

3. SwiGLU 활성화 함수

   • 기존 ReLU와 GELU 대비, SwiGLU가 정보 흐름을 효과적으로 제어하여 모델의 표현력과 성능을 크게 향상시킴.

4. Rotary Position Embedding (RoPE)

   • 절대적 위치 임베딩의 한계를 극복하고, 상대적 위치 정보를 효율적으로 반영하여 장거리 관계 모델링 성능을 향상시킴.

5. Grouped-Query Attention (GQA)

   • 메모리 대역폭 문제를 해결하고, MHA와 MQA 간의 성능-속도 절충안을 제시함으로써 대규모 모델의 추론 효율성을 크게 개선함.

7.2 주요 개선점 및 성과

• 모델 효율성: 각 개선 기법은 대규모 언어 모델에서 연산 효율성과 메모리 사용 최적화를 제공하여 병목 현상을 완화함.
• 성능 개선: RoPE, SwiGLU 및 GQA는 모델의 성능과 안정성을 동시에 높이며, 다양한 자연어 처리(NLP) 태스크에서 일관된 성과를 보임.
• 확장성: RMSNorm과 GQA 같은 기법들은 대규모 모델 확장 시 리소스 활용을 최적화하여 실질적인 배포 환경에서 유리하게 작용함.

7.3 결론

Vanila Transformer와 LLaMA 아키텍처의 차별적 설계 요소와 개선 기법들을 체계적으로 분석. 각 기법은 모델의 성능과 효율성을 향상시키며, 대규모 언어 모델의 상용화 및 확장에 기여할 수 있는 강력한 솔루션으로 자리 잡고 있음. 향후 모델들 (DeepSeek, Qwen) 역시 위에서 말한 기법들을 기준으로 나오고 있음. 여기서 어떤 부분들이 또 개선이 될지가 흥미로운 관전 포인트.