어텐션만으로는 충분하지 않았습니다

2025년 4월 29일 · 태그: ai, deep learning, transformers, python · 원문 · 한국어 번역(원문에서 변경됨)

어텐션을 처음 제안한 Attention Is All You Need 논문이 나온 이후로 수많은 최신 기법이 개발되었습니다. 수년에 걸쳐 개발된 기법 중 가장 중요한 몇 가지를 살펴보고, 그 기본 아이디어를 최대한 간결하게 구현해 봅시다. 대부분의 예제에는 PyTorch 프레임워크를 사용하겠습니다. 이 예제들은 대부분 핵심(Core) 아이디어를 매우 단순화한 스케치입니다. 전체 구현을 원하신다면 원 논문이나 PyTorch, Jax 같은 프레임워크의 프로덕션 코드를 읽어 보십시오.

그룹 쿼리 어텐션
멀티헤드 잠재 어텐션
Flash 어텐션
링 어텐션
사전 정규화
RMSNorm
SwiGLU
회전 위치 임베딩
전문가 혼합
학습률 워밍업
코사인 스케줄
AdamW 옵티마이저
다중 토큰 예측
추측 디코딩

그룹 쿼리 어텐션

특별한 순서 없이 시작해 보면, **그룹 쿼리 어텐션(Grouped Query Attention)**은 추론 중 KV 캐시의 메모리 사용량을 줄이는 기법입니다. 그룹 쿼리 어텐션은 표준 멀티헤드 어텐션 메커니즘에 대한 아키텍처 수준의 최적화입니다. GQA의 핵심(Core) 아이디어는 MHA의 계산 병목과 메모리 사용량이 K와 V 프로젝션의 크기, 그리고 그에 대응하는 캐시의 크기에 크게 좌우된다는 관찰에 기반합니다. GQA는 하나의 K 및 V 투영 세트를 여러 Q 헤드가 공유하도록 하여 이 비용을 줄이자고 제안합니다. GQA는 (MHA처럼) Q, K, V 각각에 대해 $N_{h}$ 개의 서로 다른 헤드를 두는 대신, $N_{h}$ 개의 쿼리 헤드와 $N_{k v}$ 개의 키/값 헤드만을 사용합니다. 여기서 $N_{k v} < N_{h}$ 이며, $N_{h}$ 는 일반적으로 $N_{k v}$ 의 배수입니다. 이 $N_{h}$ 개의 쿼리 헤드는 $N_{k v}$ 개의 그룹으로 나뉘며, $N_{h} / N_{k v}$ 개의 쿼리 헤드로 이루어진 각 그룹은 동일한 키 및 값 헤드에 어텐션을 수행합니다. 이 구조는 K와 V 프로젝션 행렬의 매개변수 수를 크게 줄이며, 더 중요하게는 자기회귀 디코딩 중에 필요한 K/V 캐시의 크기를 줄여 줍니다.

입력 시퀀스 표현을 $X \in R^{L \times d_{model}}$ 이라 합시다. 여기서 $L$ 은 시퀀스 길이이고 $d_{model}$ 은 모델의 임베딩 차원입니다. GQA는 먼저 서로 다른 선형 변환을 사용하여 $X$ 를 쿼리, 키, 값으로 투영합니다: $Q = X W_{Q}$ , $K = X W_{K}$ , $V = X W_{V}$ . 여기서 $W_{Q} \in R^{d_{model} \times (N_{h} d_{k})}$ , $W_{K} \in R^{d_{model} \times (N_{k v} d_{k})}$ , $W_{V} \in R^{d_{model} \times (N_{k v} d_{k})}$ 이며, $d_{k}$ 는 모델의 각 헤드의 차원(head_dim)입니다. 이들은 $N_{h}$ 개의 쿼리 헤드 $Q_{i} \in R^{L \times d_{k}}$ ( $i = 1\dots N_{h}$ )와 $N_{k v}$ 개의 키/값 헤드 $K_{j}, V_{j} \in R^{L \times d_{k}}$ ( $j = 1\dots N_{k v}$ )로 재구성됩니다. GQA의 핵심 단계는 공유입니다. $i$ 번째 쿼리 헤드에 대응하는 키 헤드와 값 헤드는 $K_{⌈ i / g ⌉}$ 와 $V_{⌈ i / g ⌉}$ 이며, 여기서 $g = N_{h} / N_{k v}$ 는 그룹 크기(KV 헤드당 쿼리 수)입니다. $i$ 번째 헤드의 어텐션 출력은 다음과 같이 계산됩니다:

$Attention (Q_{i}, K_{⌈ i / g ⌉}, V_{⌈ i / g ⌉}) = softmax (\frac{Q _{i} K _{⌈ i / g ⌉}^{T}}{d _{k}}) V_{⌈ i / g ⌉}$

구현에서는 예제 코드의 repeat_interleave 단계에서 볼 수 있듯이, $N_{k v}$ 개의 키/값 헤드를 계산한 다음 어텐션 점수를 위한 배치 행렬 곱셈 전에 이를 $g$ 번 반복하거나 인터리빙하여 $N_{h}$ 개의 쿼리 헤드에 맞춥니다. 마지막으로 모든 $N_{h}$ 개 헤드의 출력을 연결하고 출력 투영 $W_{O}$ 에 통과시킵니다.

GQA는 주로 모델 성능을 크게 손상시키지 않으면서 추론 속도를 높이고 메모리 요구량을 줄이는 기법으로 사용됩니다. 자기회귀 생성 중에는 컨텍스트 시퀀스에 대해 이전에 계산된 키와 값이 캐시되어 이후의 토큰 예측에 재사용됩니다. 이 K/V 캐시의 크기는 K/V 헤드의 수(GQA에서는 $N_{k v}$ , MHA에서는 $N_{h}$ )에 정비례합니다. GQA는 $N_{k v}$ 를 줄임으로써 각 디코딩 단계에서 K/V 캐시를 불러오는 데 필요한 메모리 대역폭을 대폭 줄이며, 이 메모리 대역폭이 주요 성능 병목입니다.

K/V 투영이 공유되므로 MHA에 비해 모델의 표현 능력이 약간 떨어질 수 있지만, 실험 결과에 따르면 GQA는 MHA 품질의 대부분을 유지하면서도 상당한 속도 향상과 메모리 절감을 제공하는 유리한 절충점을 달성하며, 이 때문에 대형 모델을 효율적으로 배포하는 데 널리 사용됩니다. $N_{k v} = 1$ 인 멀티 쿼리 어텐션(MQA)은 GQA의 극단적인 형태입니다.

class GroupQueryAttention(nn.Module):
    def __init__(self, dim, num_heads, num_kv_heads=None, head_dim=64, dropout=0.0):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads if num_kv_heads else num_heads
        self.head_dim = head_dim

        # Ensure num_heads is divisible by num_kv_heads
        assert self.num_heads % self.num_kv_heads == 0, "num_heads must be divisible by num_kv_heads"

        # Number of queries per key-value head
        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

        # Projections
        self.q_proj = nn.Linear(dim, num_heads * head_dim)
        self.k_proj = nn.Linear(dim, self.num_kv_heads * head_dim)
        self.v_proj = nn.Linear(dim, self.num_kv_heads * head_dim)
        self.o_proj = nn.Linear(num_heads * head_dim, dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Project to queries, keys, values
        q = self.q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)

        # Transpose for attention computation
        q = q.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        k = k.transpose(1, 2)  # [batch_size, num_kv_heads, seq_len, head_dim]
        v = v.transpose(1, 2)  # [batch_size, num_kv_heads, seq_len, head_dim]

        # Repeat k,v for each query head in the group
        k = k.repeat_interleave(self.num_queries_per_kv, dim=1)
        v = v.repeat_interleave(self.num_queries_per_kv, dim=1)

        # Scaled dot-product attention
        scale = 1.0 / math.sqrt(self.head_dim)
        attn = torch.matmul(q, k.transpose(2, 3)) * scale

        # Apply mask if provided
        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)

        # Softmax and dropout
        attn = torch.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        # Apply attention to values
        out = torch.matmul(attn, v)  # [batch_size, num_heads, seq_len, head_dim]
        out = out.transpose(1, 2).reshape(batch_size, seq_len, self.num_heads * self.head_dim)

        # Output projection
        out = self.o_proj(out)

        return out

멀티헤드 잠재 어텐션

멀티헤드 잠재 어텐션은 입력 시퀀스 요소들 사이에서 중간 병목 역할을 하는 학습 가능한 “잠재” 벡터 집합을 도입합니다. 핵심(Core) 아이디어는 표준 셀프 어텐션 메커니즘에 내재된 이차 계산 비용 $O (L^{2})$ 을 완화하는 것입니다. 여기서 $L$ 은 시퀀스 길이입니다. 모든 입력 요소가 다른 모든 요소에 직접 어텐션하도록 하는 대신, 입력은 먼저 고정된 수( $N_{latents}$ )의 잠재 단위에 어텐션하고, 이후 이 잠재 벡터들이 다시 입력(또는 그 변형)에 어텐션합니다. 이는 긴 입력 시퀀스 내부의 직접적인 상호작용을 효과적으로 분리하고, 이를 훨씬 작은 잠재 벡터 집합이 관여하는 두 번의 교차 어텐션 단계로 대체합니다. 이 접근 방식은 입력 시퀀스의 핵심 정보가 이러한 잠재 표현으로 효과적으로 요약되거나 압축될 수 있다고 가정하며, 이를 통해 특히 $N_{latents} ≪ L$ 일 때 표현력을 유지하면서도 계산량을 크게 줄입니다.

이 메커니즘은 일반적으로 멀티헤드 프레임워크 안에서 이루어지는 두 가지 주요 어텐션 계산 단계로 구성됩니다. 입력 시퀀스를 $X \in R^{L \times d}$ , 학습 가능한 잠재 배열을 $L \in R^{N_{latents} \times d}$ 라고 합시다. $X$ 와 $L$ 은 모두 공유되거나 별도인 투영 행렬을 사용하여 Q, K, V로 투영됩니다. 여러 헤드로 분할된 입력 투영을 $Q_{X}, K_{X}, V_{X}$ 로, 잠재 투영을 $Q_{L}, K_{L}, V_{L}$ 로 표기합시다. 첫 번째 교차 어텐션 단계는 잠재 벡터가 입력에 어떻게 어텐션하는지를 계산합니다. 잠재 쿼리 $Q_{L}$ 이 입력 키 $K_{X}$ 에 어텐션하고 입력 값 $V_{X}$ 로부터 정보를 집계합니다. 잠재 벡터에 대한 어텐션 출력은 다음과 같습니다.

$H_{L} = Attention (Q_{L}, K_{X}, V_{X}) = softmax (\frac{Q _{L} K _{X}^{T}}{d _{k}}) V_{X}$

여기서 $d_{k}$ 는 헤드 차원입니다. 두 번째 교차 어텐션 단계에서는 입력 쿼리 $Q_{X}$ 가 잠재 벡터로부터 도출된 키(예: $K_{L}$ )에 어텐션하고, 잠재 벡터와 연관된 값(이는 $V_{L}$ 일 수도 있고, 예제 코드에 구현된 것처럼 갱신된 잠재 표현 $H_{L}$ 일 수도 있습니다)으로부터 정보를 집계합니다. 그러면 최종 출력 $O$ 는 다음과 같습니다.

$O = Attention (Q_{X}, K_{L}, H_{L}) = softmax (\frac{Q _{X} K _{L}^{T}}{d _{k}}) H_{L}$

이러한 연산은 각 헤드마다 독립적으로 수행되며, 그 결과를 연결한 뒤 최종 선형 투영에 통과시킵니다.

멀티헤드 잠재 어텐션은 주로 표준 셀프 어텐션으로는 계산이 불가능할 만큼 매우 긴 시퀀스나 고차원 입력을 처리하도록 설계된 아키텍처에서 사용됩니다. 고정된 개수의 잠재 벡터( $N_{latents}$ )를 사용하면 계산 복잡도가 $O (L^{2})$ 에서 $O (L \cdot N_{latents})$ 로 줄어들어 훨씬 더 큰 입력으로 확장할 수 있습니다. 학습 가능한 잠재 벡터 $L$ 은 (코드의 self.latents = nn.Parameter(...)에서 볼 수 있듯이 무작위로 초기화되고 역전파를 통해 갱신되며) 학습 과정에서 작업과 관련된 압축 표현 또는 메모리 뱅크 역할을 하도록 적응합니다. 이 방식은 정보 병목을 유발하여 전체 셀프 어텐션에 비해 세밀한 국소 상호작용이 제한될 수 있지만, 전역 문맥을 효율적으로 포착하는 데 탁월하며 다양한 모달리티에서 효과가 입증되었습니다. 그 덕분에 이전에는 시퀀스 길이 제약 때문에 다루기 어려웠던 도메인에도 트랜스포머 계열 아키텍처를 적용할 수 있게 되었습니다.

class MultiHeadLatentAttention(nn.Module):
    def __init__(self, dim, num_heads, num_latents=64, head_dim=64, dropout=0.0):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.num_latents = num_latents
        self.head_dim = head_dim

        # Projections
        self.q_proj = nn.Linear(dim, num_heads * head_dim)
        self.k_proj = nn.Linear(dim, num_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_heads * head_dim)
        self.o_proj = nn.Linear(num_heads * head_dim, dim)

        # Latent vectors (learned)
        self.latents = nn.Parameter(torch.randn(1, num_latents, dim))

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Get latents for this batch
        latents = self.latents.expand(batch_size, -1, -1)

        # Project inputs to queries, keys, values
        q_x = self.q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k_x = self.k_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        v_x = self.v_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)

        # Project latents to queries, keys, values
        q_latents = self.q_proj(latents).reshape(batch_size, self.num_latents, self.num_heads, self.head_dim)
        k_latents = self.k_proj(latents).reshape(batch_size, self.num_latents, self.num_heads, self.head_dim)
        v_latents = self.v_proj(latents).reshape(batch_size, self.num_latents, self.num_heads, self.head_dim)

        # Transpose for attention computation
        q_x = q_x.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        k_x = k_x.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        v_x = v_x.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]

        q_latents = q_latents.transpose(1, 2)  # [batch_size, num_heads, num_latents, head_dim]
        k_latents = k_latents.transpose(1, 2)  # [batch_size, num_heads, num_latents, head_dim]
        v_latents = v_latents.transpose(1, 2)  # [batch_size, num_heads, num_latents, head_dim]

        # Scale factor for attention
        scale = 1.0 / math.sqrt(self.head_dim)

        # Compute latent-to-input attention
        attn_latent_to_input = torch.matmul(q_latents, k_x.transpose(2, 3)) * scale

        # Apply mask if provided
        if mask is not None:
            # Expand mask for the latent queries
            latent_mask = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
            attn_latent_to_input = attn_latent_to_input.masked_fill(latent_mask == 0, -1e9)

        # Softmax and dropout
        attn_latent_to_input = torch.softmax(attn_latent_to_input, dim=-1)
        attn_latent_to_input = self.dropout(attn_latent_to_input)

        # Apply attention weights to input values
        latent_output = torch.matmul(attn_latent_to_input, v_x)  # [batch_size, num_heads, num_latents, head_dim]

        # Compute input-to-latent attention
        attn_input_to_latent = torch.matmul(q_x, k_latents.transpose(2, 3)) * scale

        # Softmax and dropout
        attn_input_to_latent = torch.softmax(attn_input_to_latent, dim=-1)
        attn_input_to_latent = self.dropout(attn_input_to_latent)

        # Updated latent values are used as values for input-to-latent attention
        output = torch.matmul(attn_input_to_latent, latent_output)  # [batch_size, num_heads, seq_len, head_dim]

        # Reshape and apply output projection
        output = output.transpose(1, 2).reshape(batch_size, seq_len, self.num_heads * self.head_dim)
        output = self.o_proj(output)

        return output

Flash 어텐션

Flash 어텐션(특히 최신 구현인 FlashAttention-3)은 Transformer 내의 표준 셀프 어텐션 메커니즘에 내재된 심각한 메모리 병목 문제를, 특히 긴 시퀀스에 대해 해결합니다. 기존 방식은 전체 어텐션 점수 행렬 $S = Q K^{T}$ 를 계산하며, 여기서 $Q, K \in R^{N \times d}$ 는 길이 $N$ 인 시퀀스의 쿼리 행렬과 키 행렬입니다. 이를 위해서는 $N \times N$ 행렬 $S$ 를 저장해야 하므로 시퀀스 길이에 대해 $O (N^{2})$ 의 메모리 복잡도가 발생합니다. 이는 $N$ 이 클 때 감당할 수 없는 수준이 됩니다. Flash 어텐션은 GPU의 느린 고대역폭 메모리에 전체 $S$ 행렬을 실체화하고 저장하는 것을 피함으로써 이 문제를 극복합니다. 대신 타일링과 재계산 기법을 활용하여, 훨씬 빠른 온칩 SRAM에 들어가는 더 작은 블록 단위로 어텐션 계산을 처리합니다.

핵심(Core) 메커니즘은 Q, K, V 행렬을 블록으로 나누는 것입니다. Flash 어텐션은 K와 V의 블록을 반복적으로 SRAM에 로드하고, Q의 각 블록에 대해 역시 SRAM에 있는 현재 K 블록과의 어텐션 점수를 계산합니다. 결정적으로, 온라인 소프트맥스 알고리즘을 사용합니다. 모든 키에 대한 소프트맥스 분모 전체를 한 번에 계산하는 대신, K/V 블록을 순회하면서 누적 통계(수치 안정성을 위해 지금까지 본 최대 점수와, 정규화를 위한 지수화된 점수의 누적 합)를 유지합니다. 이를 통해 완전한 $N \times N$ 행렬을 전혀 필요로 하지 않고도 올바르게 스케일링된 어텐션 출력을 블록 단위로 계산할 수 있습니다. 중간 결과를 주로 빠른 SRAM 안에 유지하고 고대역폭 메모리와의 데이터 전송을 최소화함으로써, Flash 어텐션은 시퀀스 길이에 따른 메모리 사용량을 $O (N^{2})$ 에서 $O (N)$ (Q, K, V 자체의 저장이 대부분을 차지함)으로 크게 줄이고, 개선된 메모리 접근 패턴 덕분에 상당한 속도 향상을 달성합니다.

실제로 FlashAttention-3 구현은 다양한 하드웨어 구성에서 효율적으로 동작하도록 설계된, 고도로 최적화된 CUDA 커널들의 모음입니다. 하지만 PyTorch로 작성한 최소한의 장난감 구현은 다음과 같습니다:

class FlashAttention(nn.Module):
    def __init__(self, dim, num_heads, head_dim=64, dropout=0.0, block_size=1024):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.block_size = block_size

        # Projections
        self.q_proj = nn.Linear(dim, num_heads * head_dim)
        self.k_proj = nn.Linear(dim, num_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_heads * head_dim)
        self.o_proj = nn.Linear(num_heads * head_dim, dim)

        self.dropout = nn.Dropout(dropout)

    def _flash_attention_forward(self, q, k, v, mask=None):
        # This is a simplified approximation of Flash Attention
        # In practice, FlashAttention uses custom CUDA kernels for tiled attention

        batch_size, num_heads, seq_len, head_dim = q.shape
        scale = 1.0 / math.sqrt(head_dim)

        # Initialize output and attention statistics
        output = torch.zeros_like(q)
        normalizer = torch.zeros((batch_size, num_heads, seq_len, 1), device=q.device)

        # Process blocks of keys and values
        for block_start in range(0, seq_len, self.block_size):
            block_end = min(block_start + self.block_size, seq_len)

            # Extract key and value blocks
            k_block = k[:, :, block_start:block_end]
            v_block = v[:, :, block_start:block_end]

            # Compute attention scores for this block
            attn_scores = torch.matmul(q, k_block.transpose(2, 3)) * scale

            # Apply mask if provided
            if mask is not None:
                block_mask = mask[:, :, :, block_start:block_end]
                attn_scores = attn_scores.masked_fill(block_mask == 0, -1e9)

            # Apply softmax and dropout
            attn_probs = torch.softmax(attn_scores, dim=-1)
            attn_probs = self.dropout(attn_probs)

            # Update output with the attention results for this block
            output += torch.matmul(attn_probs, v_block)
            normalizer += attn_probs.sum(dim=-1, keepdim=True)

        # Normalize the output
        output = output / (normalizer + 1e-6)

        return output

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Project queries, keys, values
        q = self.q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)

        # Transpose for attention computation
        q = q.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        k = k.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        v = v.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]

        # Compute flash attention
        output = self._flash_attention_forward(q, k, v, mask)

        # Reshape and apply output projection
        output = output.transpose(1, 2).reshape(batch_size, seq_len, self.num_heads * self.head_dim)
        output = self.o_proj(output)

        return output

실제로 플래시 어텐션 저장소 Dao-AILab/flash-attention에는 다양한 하드웨어 구성에 맞춘 여러 가지 구현이 있습니다. flash_attn_qkvpacked_func 함수는 FlashAttention-3 구현을 사용하는 방법을 보여 주는 최소한의 예제입니다. 이 함수는 패킹된 QKV 텐서를 입력으로 받아 어텐션 출력을 반환합니다.

import torch
from flash_attn import flash_attn_qkvpacked_func

# Minimal configuration
BATCH_SIZE, SEQ_LEN, NUM_HEADS, HEAD_DIM = 2, 64, 4, 32
CAUSAL = False
DTYPE = torch.float16
DEVICE = "cuda"

# Create dummy packed QKV tensor
# Shape: (batch_size, seq_len, 3, num_heads, head_dim)
qkv = torch.randn(
    BATCH_SIZE,
    SEQ_LEN,
    3,
    NUM_HEADS,
    HEAD_DIM,
    dtype=DTYPE,
    device=DEVICE,
)

print(f"Input qkv shape: {qkv.shape}")

# Call FlashAttention packed QKV function
output = flash_attn_qkvpacked_func(
    qkv,
    dropout_p=0.0, # Set dropout probability (0.0 for no dropout)
    causal=CAUSAL,
    softmax_scale=None # Use default scaling (1 / sqrt(head_dim))
)

# Output shape: (batch_size, seq_len, num_heads, head_dim)
print(f"Output shape: {output.shape}")
print("FlashAttention call successful.")

링 어텐션

Ring Attention은 여러 GPU에서 셀프 어텐션을 블록 단위로 계산하여, 단일 장치로 처리하기에는 너무 긴 시퀀스의 학습과 추론을 가능하게 합니다. 이 방식은 표준 셀프 어텐션 메커니즘에 내재된 심각한 메모리 병목을 해결하며, 특히 전체 어텐션 점수 행렬의 이차 메모리 복잡도 $O (N^{2})$ 를 감당할 수 없게 되는 매우 긴 시퀀스를 처리할 때 유용합니다.

핵심(Core) 아이디어는 개념적으로 링 토폴로지로 배열된 GPU와 같은 여러 처리 장치에 계산을 분산하는 것입니다. 이 방식에서는 어떤 단일 장치도 K와 V 텐서 전체를 보유할 필요가 없습니다. 대신 이 텐서들은 시퀀스 길이 차원을 따라 샤딩되거나 청크로 나뉘므로, 장치당 최대 메모리 요구량이 크게 줄어들고 개별 가속기의 메모리 용량을 초과했을 시퀀스에 대해서도 어텐션 계산이 가능해집니다.

실제 분산 구현에서 각 장치는 처음에 입력 시퀀스 중 자신이 맡은 구간에 해당하는 Q, K, V 샤드를 보유합니다. 어텐션 계산은 이 링 전체에서 동기화된 단계별로 진행됩니다. 각 단계에서 장치는 자신의 로컬 Q 샤드와 현재 보유한 K 샤드를 사용하여 부분 어텐션 점수를 계산합니다. 핵심 요소는 그다음에 이어지는 통신입니다. K와 V 샤드가 링의 다음 장치로 전달됩니다. 이 순환은 모든 Q 샤드가 모든 K/V 샤드와 상호작용할 때까지 반복됩니다. 이 과정 전반에 걸쳐 각 장치는 부분 출력(가중된 V 벡터)과 정규화 인자(소프트맥스 분모)를 누적합니다. 어텐션 출력을 최종적으로 완성하려면 일반적으로 모든 장치에 걸친 집합 연산을 수행하여 시퀀스의 각 구간에 대해 이러한 부분 결과를 올바르게 결합해야 합니다.

아래의 Python 예제는 단일 장치에서 링 어텐션(Ring Attention) 로직을 시뮬레이션하여, 실제 멀티 GPU 하드웨어 없이도 그 기본 원리를 보여 줍니다. _simulate_ring_attention 함수는 가상의 샤드들을 순회하면서 분산 과정을 모방합니다. 각 반복에서 이 함수는 K와 V 텐서의 슬라이스(k_shard, v_shard)를 선택하여 특정 단계에서 한 장치가 처리할 데이터를 나타냅니다. 그런 다음 전체 Q 텐서(실제 분산 환경을 단순화한 것입니다)와 현재 k_shard 사이의 어텐션 점수를 계산합니다. 이 시뮬레이션은 반복 전반에 걸쳐 가중된 값과 소프트맥스 정규화 인자를 누적함으로써 링 방식의 본질을 효과적으로 포착하며, 이는 분산 환경에서 최종 정규화 단계가 출력을 만들어 내기 전에 부분 결과가 결합되는 방식을 그대로 반영합니다. 이 시뮬레이션은 계산 흐름을 보여 주지만, 당연히 실제 다중 장치 링 어텐션(Ring Attention) 구현이 제공하는 병렬성이나 메모리 절감 효과는 제공하지 않습니다.

class RingAttention(nn.Module):
    def __init__(self, dim, num_heads, head_dim=64, dropout=0.0, num_shards=4):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.num_shards = num_shards

        # Projections
        self.q_proj = nn.Linear(dim, num_heads * head_dim)
        self.k_proj = nn.Linear(dim, num_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_heads * head_dim)
        self.o_proj = nn.Linear(num_heads * head_dim, dim)

        self.dropout = nn.Dropout(dropout)

    def _simulate_ring_attention(self, q, k, v, mask=None):
        # This simulates ring attention without actual multi-GPU support
        batch_size, num_heads, seq_len, head_dim = q.shape
        scale = 1.0 / math.sqrt(head_dim)

        # Compute shard sizes
        shard_size = (seq_len + self.num_shards - 1) // self.num_shards

        # Initialize outputs
        output = torch.zeros_like(q)
        normalizer = torch.zeros((batch_size, num_heads, seq_len, 1), device=q.device)

        # Simulate sharded processing
        for shard_idx in range(self.num_shards):
            start_idx = shard_idx * shard_size
            end_idx = min(start_idx + shard_size, seq_len)

            # Process this shard's keys and values
            if start_idx < seq_len:
                k_shard = k[:, :, start_idx:end_idx]
                v_shard = v[:, :, start_idx:end_idx]

                # Compute attention scores
                attn_scores = torch.matmul(q, k_shard.transpose(2, 3)) * scale

                # Apply mask if provided
                if mask is not None:
                    shard_mask = mask[:, :, :, start_idx:end_idx]
                    attn_scores = attn_scores.masked_fill(shard_mask == 0, -1e9)

                # Apply softmax and dropout (accumulated over shards)
                attn_probs = torch.softmax(attn_scores, dim=-1)
                attn_probs = self.dropout(attn_probs)

                # Update output and normalizer
                output += torch.matmul(attn_probs, v_shard)
                normalizer += attn_probs.sum(dim=-1, keepdim=True)

        # Normalize the output
        output = output / (normalizer + 1e-6)

        return output

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        # Project queries, keys, values
        q = self.q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim)

        # Transpose for attention computation
        q = q.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        k = k.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]
        v = v.transpose(1, 2)  # [batch_size, num_heads, seq_len, head_dim]

        # Compute ring attention
        output = self._simulate_ring_attention(q, k, v, mask)

        # Reshape and apply output projection
        output = output.transpose(1, 2).reshape(batch_size, seq_len, self.num_heads * self.head_dim)
        output = self.o_proj(output)

        return output

사전 정규화

사전 정규화(흔히 Pre-LN이라고 부릅니다)는 잔차 블록의 아키텍처 설계에 일어난 변화였습니다. 기존의 사후 정규화 방식처럼 정규화 레이어를 주요 연산(셀프 어텐션이나 피드포워드 네트워크 등) 뒤에 적용하는 대신, 사전 정규화는 이를 앞에 적용합니다. 사소해 보이는 이 변화는 학습 동역학에 중대한 영향을 미칩니다. 사전 정규화는 입력이 계산 집약적인 하위 레이어에 들어가기 전에 이를 정규화함으로써 네트워크를 통해 흐르는 활성화와 그래디언트를 안정화하는 데 도움을 줍니다. 이러한 안정화 효과는 매우 깊은 네트워크에서 특히 두드러지며, 그래디언트 소실이나 폭주 같은 문제를 완화하고 흔히 더 높은 학습률과 더 빠르고 안정적인 수렴을 가능하게 합니다.

잔차 블록 내의 일반적인 구현은 PreNorm 클래스에서 보여 주듯이 $x + f (norm (x))$ 구조를 따릅니다. 여기서 $x$ 는 블록의 입력이고, $norm (\cdot)$ 은 레이어 정규화(LN)나 RMS 정규화(RMSNorm) 같은 정규화 함수를 나타내며, $f (\cdot)$ 는 주 변환 함수(예: 멀티헤드 어텐션 또는 위치별 피드포워드 네트워크)를 나타냅니다. 입력 $x$ 는 먼저 정규화 레이어(예: self.norm(x))를 통과합니다. 그런 다음 정규화된 출력은 함수 fn에 의해 처리됩니다. 결정적으로, 이 함수의 출력은 잔차 연결을 통해 정규화되지 않은 원래 입력 $x$ 에 다시 더해집니다. 이 구조는 항등 연결(+ x)을 통한 깨끗한 그래디언트 경로를 보장하므로, 정규화 레이어가 잔차 경로 자체에 놓이는 사후 정규화에 비해 학습 안정성을 한층 더 높여 줍니다.

class PreNorm(nn.Module):
    def __init__(self, dim, fn, norm_type='layer'):
        super().__init__()
        self.fn = fn

        if norm_type == 'layer':
            self.norm = nn.LayerNorm(dim)
        elif norm_type == 'rms':
            self.norm = RMSNorm(dim)
        else:
            raise ValueError(f"Unknown normalization type: {norm_type}")

    def forward(self, x, *args, **kwargs):
        # Apply normalization first, then the function
        return self.fn(self.norm(x), *args, **kwargs) + x

RMSNorm

RMSNorm(제곱평균제곱근 정규화)은 널리 사용되는 LayerNorm을 단순화한 것으로, 비슷한 성능을 유지하고 흔히 학습 안정성을 높이면서도 계산 오버헤드를 줄이도록 설계되었습니다. 평균을 빼서 활성화를 중심화한 다음 표준편차로 스케일링하는 LayerNorm과 달리, RMSNorm은 평균 중심화 단계를 완전히 생략합니다. 이러한 단순화의 동기는 LayerNorm의 재중심화 연산이 계산 비용의 상당 부분을 차지하며, 이를 제거해도 모델 성능이 크게 저하되지 않는 경우가 많고 때로는 오히려 도움이 되기도 한다는 경험적 관찰에서 비롯되었습니다. RMSNorm은 오직 입력을 제곱평균제곱근 크기에 따라 다시 스케일링하는 방식으로만 동작합니다.

RMSNorm의 핵심(Core) 메커니즘은 한 레이어 내의 활성화 값을 특성(또는 지정된 차원) 전체에 걸쳐 계산한 제곱평균제곱근 값으로 나누어 정규화하는 것입니다. 입력 벡터 $x = (x_{1}, \dots, x_{n})$ 에 대해 RMS 값은 $RMS (x) = \frac{1}{n} \sum_{i = 1}^{n} x_{i}^{2}$ 로 계산됩니다. 그러면 정규화된 출력 $\overset{x}{ˉ}_{i}$ 는 $\overset{x}{ˉ}_{i} = \frac{x _{i}}{RMS ( x ) + ϵ}$ 이 되며, 여기서 $ϵ$ 은 수치적 안정성을 위한 작은 상수입니다. LayerNorm과 마찬가지로 RMSNorm은 일반적으로 학습 가능한 스케일링 매개변수 $g$ 를 포함하며(때로는 편향 $b$ 도 포함하지만, 원래 공식은 단순화 원칙에 더 충실하기 위해 이를 생략하는 경우가 많습니다), 그 결과 최종 출력은 $y_{i} = g_{i} \overset{x}{ˉ}_{i}$ 가 됩니다. 평균 계산을 생략함으로써 RMSNorm은 계산량(설정에 따라 GPU에서 LayerNorm보다 7-64% 빠른 것으로 추정됩니다)과 메모리 사용량을 줄여 주므로, 특히 효율성이 무엇보다 중요한 대규모 모델에서 매력적인 대안이 됩니다.

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-8, elementwise_affine=True):
        super().__init__()
        self.eps = eps
        self.elementwise_affine = elementwise_affine

        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_parameter('weight', None)

    def forward(self, x):
        # Calculate root mean square along the last dimension
        rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)

        # Normalize by RMS
        x_normalized = x / rms

        # Apply scaling if using learnable parameters
        if self.elementwise_affine:
            x_normalized = x_normalized * self.weight

        return x_normalized

SwiGLU

SwiGLU는 게이트 선형 유닛(Gated Linear Unit, GLU) 계열에서 파생된 활성화 함수로, 신경망의 성능을 향상시키도록 특별히 설계되었습니다. GLU 변형들의 핵심(Core) 개념은 네트워크를 통과하는 정보의 흐름을 적응적으로 제어하는 게이팅 메커니즘을 도입하는 것입니다. 표준 피드포워드 계층은 일반적으로 입력의 선형 변환에 단일 비선형성을 적용합니다. 반면 GLU 기반 활성화 함수는 선형 레이어의 출력을 두 부분으로 나누며, 한 부분은 비선형성을 거친 뒤 “게이트” 역할을 하여 요소별 곱셈을 통해 다른 부분을 조절합니다. SwiGLU는 게이팅 부분에 적용되는 비선형성으로 Swish라고도 알려진 Sigmoid-weighted Linear Unit(SiLU)( $SiLU (x) = x \cdot σ (x)$ , 여기서 $σ$ 는 시그모이드 함수입니다)를 사용한다는 점에서 차별화됩니다. 이러한 선택은 ReLU와 같은 다른 활성화 함수나 게이트에 시그모이드 또는 다른 함수를 사용하는 표준 GLU 변형에 비해 다양한 트랜스포머 기반 모델에서 상당한 성능 향상을 가져온다는 것이 실험적으로 입증되었습니다.

피드포워드 블록 내에서 SwiGLU의 작동 메커니즘은 일반적으로 두 개의 별도 선형 변환을 사용해 입력 $x$ 를 사영하여 $W x + b$ 와 $V x + c$ 를 얻는 과정을 포함합니다. 그런 다음 SwiGLU 활성화는 $SwiGLU (x, W, V, b, c) = SiLU (W x + b) ⊙ (V x + c)$ 로 계산되며, 여기서 $⊙$ 는 요소별 곱셈을 나타냅니다. 사실상 입력은 두 개의 병렬 선형 경로를 통해 처리됩니다. 한 경로는 SiLU 활성화를 거쳐 게이트 값을 형성하고, 이 값이 두 번째 선형 경로의 출력을 스케일링합니다. 이러한 게이팅을 통해 네트워크는 입력 맥락에 따라 어떤 특성을 앞으로 전달할지 동적으로 제어할 수 있으며, 그 결과 더 단순한 활성화 함수에 비해 표현력이 높아지고 그래디언트 흐름이 개선됩니다. PaLM과 LLaMA 같은 모델에서의 성공은 언어 데이터 내의 복잡한 패턴을 포착하는 데 있어 SwiGLU의 효과를 잘 보여 줍니다.

class SwiGLU(nn.Module):
    def __init__(self, dim_in, dim_hidden=None, dim_out=None, bias=True):
        super().__init__()
        dim_hidden = dim_hidden or 4 * dim_in
        dim_out = dim_out or dim_in

        # Linear transformations for gating
        self.w1 = nn.Linear(dim_in, dim_hidden, bias=bias)
        self.w2 = nn.Linear(dim_in, dim_hidden, bias=bias)

        # Output projection
        self.w3 = nn.Linear(dim_hidden, dim_out, bias=bias)

    def forward(self, x):
        # SwiGLU applies SiLU activation to one branch and gates it with the other
        hidden1 = self.w1(x)
        hidden2 = self.w2(x)

        # SiLU (Swish) activation: x * sigmoid(x)
        hidden1_act = hidden1 * torch.sigmoid(hidden1)

        # Element-wise product for gating
        hidden = hidden1_act * hidden2

        # Output projection
        return self.w3(hidden)

회전 위치 임베딩

회전 위치 임베딩(Rotary Positional Embedding, RoPE)은 트랜스포머 모델의 셀프 어텐션 메커니즘에 위치 정보를 직접 통합하는 우아한 방법을 제시하며, 특히 상대적 위치 의존성을 효과적으로 포착하도록 설계되었습니다. 기존 접근 방식은 흔히 입력 임베딩에 절대 위치 인코딩을 더하거나 어텐션 점수 계산 내에 복잡한 상대 위치 편향 항을 사용하는 방식에 의존합니다. RoPE는 위치 인코딩을 쿼리 벡터와 키 벡터의 내적을 계산하기 전에 이들 벡터에 적용하는 회전 연산으로 본다는 점에서 다른 길을 택합니다. 핵심 통찰은 위치 $m$ 에 해당하는 Q 벡터와 위치 $n$ 에 해당하는 K 벡터를 각각 $m$ 과 $n$ 에 비례하는 각도만큼 회전시키면, 그 결과로 얻는 내적이 본질적으로 상대적 위치 $m - n$ 과 벡터 자체의 내용에만 의존하게 되어 벡터의 크기를 바꾸지 않고도 상대적 거리를 깔끔하게 인코딩한다는 것입니다.

핵심(Core) 수학적 아이디어는 복소수 곱셈의 성질, 또는 이와 동등하게 2D 회전 행렬의 성질을 활용합니다. 임베딩 차원의 쌍을 복소수로 표현한다고 상상해 보십시오. RoPE는 위치에 따라 달라지는 회전 행렬 $R_{m}$ 과 $R_{n}$ 을 사용하여 위치 $m$ 의 쿼리 벡터 $q_{m}$ 과 위치 $n$ 의 키 벡터 $k_{n}$ 에 각각 회전을 적용합니다. 이 행렬들은 사실상 차원 쌍으로 생성되는 2D 부분 공간에서 벡터를 회전시킵니다. 각 2D 부분 공간의 회전 각도는 절대 위치( $m$ 또는 $n$ )에 주파수 항 $θ_{i}$ 를 곱하여 결정되며, 이 주파수 항은 사인파 임베딩과 유사하게 차원이 높을수록 감소합니다. 결정적인 성질은 회전된 벡터 간의 내적 $(R_{m} q_{m})^{T} (R_{n} k_{n})$ 이 원래 벡터 $q_{m}, k_{n}$ 과 상대적 거리에 해당하는 회전 행렬 $R_{m - n}$ 에만 의존하도록 단순화된다는 것이며, 이를 통해 상대적 위치 정보가 어텐션 점수 계산에 직접 임베딩됩니다.

실제로 이 회전은 명시적인 행렬 곱셈 없이 효율적으로 구현됩니다. apply_rotary_pos_emb 함수에서 볼 수 있듯이, 임베딩 차원은 보통 두 개씩 쌍으로 나뉩니다. 각 쌍 $(x_{i}, x_{i + 1})$ 에 대해, 위치 $m$ 과 주파수 $θ_{j}$ (RotaryEmbedding 클래스의 inv_freq에서 도출됨)에 해당하는 회전이 삼각함수를 사용해 적용됩니다. 즉, 새로운 쌍 $(x_{i}^{'}, x_{i + 1}^{'})$ 은 $(x_{i} cos (m θ_{j}) - x_{i + 1} sin (m θ_{j}), x_{i + 1} cos (m θ_{j}) + x_{i} sin (m θ_{j}))$ 가 됩니다. RotaryEmbedding 클래스는 시퀀스 길이와 역주파수 대역(inv_freq)을 바탕으로 필요한 코사인 값과 사인 값(cos, sin)을 미리 계산합니다. 역주파수 대역은 기준값(base)과 임베딩 차원(dim)에서 도출됩니다. 이렇게 미리 계산된 값은 모든 위치 $m$ 과 관련 주파수 $θ_{j}$ 에 대한 $cos (m θ_{j})$ 와 $sin (m θ_{j})$ 를 나타냅니다.

RoPE를 적용하려면 주어진 시퀀스 길이에 대해 이 cos 및 sin 임베딩을 생성한 다음, 초기 선형 사영 이후이면서 어텐션 점수 계산 이전에 이를 사용해 Q 벡터와 K 벡터를 헤드 차원에 걸쳐 쌍 단위로 변환합니다. 이 방법에는 여러 장점이 있습니다. 상대 위치를 자연스럽게 인코딩하고, 단어 임베딩에 위치 정보를 직접 더하지 않으며(그 결과 더 많은 의미 정보를 보존할 수 있습니다), 학습 중에 본 것보다 긴 시퀀스 길이에도 잘 일반화되는 등 실험적으로 뛰어난 성능을 보여 왔습니다. RoPE는 회전을 통해 위치 정보를 통합함으로써 문맥을 고려한 시퀀스 모델링을 위한 계산 효율적이고 효과적인 메커니즘을 제공합니다.

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, base=10000, interleaved=False):
        super().__init__()
        self.dim = dim
        self.base = base
        self.interleaved = interleaved

        # Generate inverse frequency bands
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, seq_len, device=None):
        # Get device from buffer if not specified
        if device is None:
            device = self.inv_freq.device

        # Generate position indices
        positions = torch.arange(seq_len, device=device).float()

        # Compute sinusoidal patterns
        freqs = torch.outer(positions, self.inv_freq)

        # Get sine and cosine embeddings
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = torch.cos(emb)[:, :self.dim]
        sin = torch.sin(emb)[:, :self.dim]

        return cos, sin

def apply_rotary_pos_emb(q, k, cos, sin, interleaved=False):
    # Apply rotary embeddings to queries and keys
    batch_size, num_heads, seq_len, head_dim = q.shape
    cos = cos.reshape(1, 1, seq_len, cos.shape[-1])  # [1, 1, seq_len, dim/2]
    sin = sin.reshape(1, 1, seq_len, sin.shape[-1])  # [1, 1, seq_len, dim/2]

    # Split queries and keys for rotation
    half_dim = head_dim // 2
    q1, q2 = q[..., :half_dim], q[..., half_dim:]
    k1, k2 = k[..., :half_dim], k[..., half_dim:]

    # Apply rotation using half-dim rotary embeddings
    q_rotated = torch.cat([
        q1 * cos - q2 * sin,
        q2 * cos + q1 * sin
    ], dim=-1)

    k_rotated = torch.cat([
        k1 * cos - k2 * sin,
        k2 * cos + k1 * sin
    ], dim=-1)

    return q_rotated, k_rotated

전문가 혼합(Mixture of Experts)

전문가 혼합(Mixture of Experts, 줄여서 MoE)은 추론이나 학습 시 계산 비용을 비례적으로 크게 늘리지 않으면서 신경망의 매개변수 수, 그리고 그에 따른 잠재적 용량을 크게 늘리도록 설계된 모델 아키텍처입니다. 핵심(Core) 아이디어는 트랜스포머의 피드포워드 네트워크 블록처럼 계산 비용이 큰 구성 요소를 여러 개의 더 작은 “전문가(expert)” 네트워크로 대체하는 것입니다. 중요한 점은 모든 전문가가 모든 입력 토큰을 처리하지는 않는다는 것입니다. 대신 가벼운 “라우터(router)” 또는 “게이팅(gating)” 네트워크가 각 입력 토큰의 특징을 바탕으로 해당 토큰을 처리하기에 가장 적합하다고 판단되는 소수의 전문가(보통 한두 개뿐이며, 이를 top-k 라우팅이라고 합니다)를 동적으로 선택합니다. 이러한 조건부 계산 덕분에 MoE 모델은 잠재적으로 수조 개의 매개변수를 가지면서도 주어진 입력에 대해서는 그중 극히 일부만 활성화하므로, 비슷한 크기의 밀집 모델에 비해 FLOPs를 감당 가능한 수준으로 유지합니다.

라우터 네트워크가 핵심(Core) 아이디어이며, 학습된 의사 결정자 역할을 합니다. 라우터는 입력 토큰의 표현을 받아 사용 가능한 각 전문가가 해당 토큰에 얼마나 적합한지를 나타내는 점수 또는 로짓을 생성합니다. 예제 코드의 _compute_routing_weights 함수에서는 이 로짓을 흔히 top-k 함수로 처리하여 점수가 가장 높은 k개의 전문가를 식별합니다. 그런 다음 선택된 전문가들의 점수는 보통 소프트맥스 함수를 사용해 정규화되어 라우팅 가중치가 됩니다. 이 가중치는 선택된 각 전문가가 해당 토큰의 최종 출력에 얼마나 기여하는지를 결정합니다. 학습 중에는 탐색을 장려하고 라우터가 초기에 소수의 전문가만 항상 선호하는 상태로 붕괴하는 것을 막기 위해 라우터 로짓에 노이즈를 추가할 수 있습니다(noise_std에서 볼 수 있습니다).

라우터가 주어진 입력 토큰에 대해 top-k 전문가를 선택하고 각각의 가중치를 계산하면, 해당 토큰은 선택된 전문가들에게만 전달됩니다. 선택된 각 전문가 네트워크(experts ModuleList에서 볼 수 있듯이 흔히 표준 FFN입니다)는 토큰을 독립적으로 처리합니다. 그런 다음 이렇게 활성화된 전문가들이 생성한 출력을 결합하여 해당 토큰의 최종 출력을 만듭니다. 이 결합은 보통 가중합으로, 선택된 각 전문가의 출력에 라우터가 계산한 해당 라우팅 가중치를 곱합니다. 예를 들어 전문가 i와 j가 가중치 w_i와 w_j로 선택되었다면, 토큰 x의 최종 출력은 w_i * expert_i(x) + w_j * expert_j(x)가 됩니다. 이를 통해 최종 표현에 가장 관련성 높은 전문가들의 특화된 지식이 반영됩니다.

MoE 모델 학습의 과제 중 하나는 모든 전문가가 효과적으로 활용되도록 하는 것입니다. 그렇지 않으면 라우터가 소수의 “인기 있는” 전문가에게 계속 과부하를 주도록 학습하고, 나머지 전문가는 충분히 발달하지 못할 수 있습니다. 이를 방지하기 위해 _compute_balance_loss 메서드에서 보여 주듯이, 보통 보조 부하 분산 손실(auxiliary load balancing loss)을 학습 목표에 포함합니다. 이 손실은 흔히 각 전문가에게 할당된 토큰 수나 각 전문가로 향하는 라우팅 가중치 합의 불균형에 페널티를 부여함으로써, 라우터가 계산 부하(즉, 입력 토큰)를 사용 가능한 모든 전문가에게 더 고르게 분배하도록 유도합니다. MoE는 라우팅을 통한 희소 활성화를 성공적으로 구현하고 부하 분산을 통합함으로써, 매우 크면서도 계산 효율적인 모델을 구축할 수 있게 합니다.

class MixtureOfExperts(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_experts=4, top_k=2, noise_std=1.0):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.num_experts = num_experts
        self.top_k = min(top_k, num_experts)
        self.noise_std = noise_std

        # Create experts
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(input_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, output_dim)
            ) for _ in range(num_experts)
        ])

        # Router network
        self.router = nn.Linear(input_dim, num_experts)

        # Save expert counts and loads for balancing loss
        self.register_buffer('expert_counts', torch.zeros(num_experts))

    def _compute_routing_weights(self, x):
        # Calculate routing weights
        routing_logits = self.router(x)  # [batch_size, seq_len, num_experts]

        # Add noise during training to encourage exploration
        if self.training and self.noise_std > 0:
            noise = torch.randn_like(routing_logits) * self.noise_std
            routing_logits = routing_logits + noise

        # Get top-k experts for each token
        routing_weights, selected_experts = torch.topk(routing_logits, self.top_k, dim=-1)

        # Normalize the routing weights with softmax
        routing_weights = F.softmax(routing_weights, dim=-1)

        return routing_weights, selected_experts

    def _compute_balance_loss(self, selected_experts, routing_weights):
        # Compute auxiliary load balancing loss
        batch_size, seq_len, _ = selected_experts.shape

        # Compute probability of each expert being selected across batch
        expert_mask = torch.zeros(batch_size, seq_len, self.num_experts, device=selected_experts.device)

        # For each position in selected_experts, increment the corresponding expert index
        for k in range(self.top_k):
            expert_mask.scatter_(-1, selected_experts[..., k:k+1], routing_weights[..., k:k+1])

        # Compute mean routing probability per expert
        expert_routing_probs = expert_mask.mean(dim=[0, 1])

        # Compute load balancing loss (all experts should receive equal probability)
        target_probs = torch.ones_like(expert_routing_probs) / self.num_experts
        balance_loss = F.mse_loss(expert_routing_probs, target_probs) * self.num_experts

        return balance_loss

    def forward(self, x):
        batch_size, seq_len, _ = x.shape

        # Compute routing weights and selected experts
        routing_weights, selected_experts = self._compute_routing_weights(x)

        # Prepare output
        output = torch.zeros(batch_size, seq_len, self.output_dim, device=x.device)

        # Dispatch to selected experts
        for k in range(self.top_k):
            # Extract the current expert indices and weights
            expert_indices = selected_experts[..., k]  # [batch_size, seq_len]
            expert_weights = routing_weights[..., k].unsqueeze(-1)  # [batch_size, seq_len, 1]

            # Dispatch to each expert
            for expert_idx in range(self.num_experts):
                # Find tokens assigned to this expert
                mask = (expert_indices == expert_idx)

                if mask.any():
                    # Gather input tokens assigned to this expert
                    expert_inputs = x[mask]

                    # Process with the expert
                    expert_outputs = self.experts[expert_idx](expert_inputs)

                    # Scatter outputs back with appropriate weights
                    output[mask] += expert_outputs * expert_weights[mask]

                    # Update expert counts (for monitoring)
                    self.expert_counts[expert_idx] += mask.sum().item()

        # Compute auxiliary load balancing loss
        balance_loss = self._compute_balance_loss(selected_experts, routing_weights)

        # Return output and auxiliary loss
        return output, balance_loss

학습률 워밍업

학습률 워밍업은 신경망 훈련 초기 단계에서 안정성을 높이고 발산을 막기 위해 널리 쓰이는 휴리스틱입니다. 훈련을 시작할 때 모델 매개변수는 보통 무작위로 초기화되므로 최적의 구성과는 거리가 먼 경우가 많습니다. 처음부터 비교적 큰 학습률을 사용하면, 초기 그래디언트 역시 크고 불규칙할 수 있어 매개변수가 급격하게 갱신될 수 있습니다. 그 결과 모델이 손실 지형의 나쁜 영역으로 밀려나거나 수치적 불안정성(예: 손실 폭발)이 생길 수도 있습니다. 학습률 워밍업은 매우 작은 학습률로 학습 과정을 시작하고, 미리 정해진 수의 초기 학습 스텝(“워밍업 스텝”) 동안 학습률을 점진적으로 높여 목표 기본값에 도달하게 함으로써 이러한 위험을 완화하며, 그 시점부터 표준 학습률 스케줄(예: 감쇠)이 시작될 수 있습니다.

이 방식은 워밍업 단계 동안 기본 학습률에 곱하는 배율을 점진적으로 키우는 것입니다. 흔히 쓰이는 전략은 선형 워밍업이며, LinearWarmupScheduler 클래스가 그 예를 보여 줍니다. 이 방식에서 스텝 $t$ 의 학습률 $η_{t}$ 는 $t < T_{warmup}$ 일 때 $η_{t} = η_{base} \times \frac{t}{T _{warmup}}$ 로 계산합니다. 여기서 $η_{base}$ 는 목표 기본 학습률이고 $T_{warmup}$ 은 전체 워밍업 스텝 수입니다. get_lr 메서드에서 볼 수 있듯이, 배율 인자 scale은 warmup_steps 동안 0에 가까운 값에서 1까지 선형으로 증가합니다. 스텝 수 last_epoch가 warmup_steps에 도달하면 학습률은 base_lrs와 같아지고 워밍업 단계가 끝납니다. 이처럼 완만하게 높여 가면 활성화 값과 그래디언트가 불안정해지기 쉬운 중요한 초기 단계에서 모델이 점진적으로 적응할 수 있습니다. 그 결과 수렴이 더 매끄러워지고, 훈련 후반에 더 높은 기본 학습률을 사용할 수 있는 경우가 많습니다.

class LinearWarmupScheduler(_LRScheduler):
    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
        self.warmup_steps = warmup_steps
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.warmup_steps:
            # During warmup: linearly increase from 0 to base LR
            scale = float(self.last_epoch + 1) / float(max(1, self.warmup_steps))
            return [base_lr * scale for base_lr in self.base_lrs]
        else:
            # After warmup: use base learning rate
            return self.base_lrs

코사인 스케줄

코사인 스케줄링(코사인 어닐링이라고도 합니다)은 학습률 스케줄 기법의 하나입니다. 이 기법의 핵심(Core) 원리는 학습 과정 전반에 걸쳐 코사인 곡선의 형태를 따라 학습률을 점진적으로 낮추는 것입니다. 특정 에폭에서 학습률을 갑자기 낮추는 스텝 감쇠 스케줄과 달리, 코사인 어닐링은 학습률을 매끄럽고 연속적으로 낮춥니다. 일반적으로 학습률은 높은 초깃값에서 시작하여 코사인 함수의 첫 반주기를 따라 감소하며, 마지막 훈련 스텝에서 미리 정한 최솟값(대개 0에 가까운 값)에 도달합니다. 이러한 매끄러운 감쇠는 최적화 과정에 도움이 된다는 것이 경험적으로 확인되었습니다. 훈련 초기에는 큰 스텝으로 손실 지형을 폭넓게 탐색하고, 후반으로 갈수록 스텝을 점점 줄여 미세 조정하면서 좋은 최솟값으로 수렴할 수 있기 때문입니다.

훈련 초기에는 학습률이 높을수록 더 빠르게 탐색할 수 있고, 나쁜 지역 최솟값에서 벗어나기도 쉽습니다. 훈련이 진행되어 모델 매개변수가 더 최적에 가까운 영역에 다가가면, 최솟값을 지나쳐 버리지 않고 더 정밀하게 수렴하도록 학습률을 낮추는 것이 매우 중요해집니다. 코사인 함수가 만드는 스케줄은 처음에는 감쇠 속도가 비교적 느리므로, 옵티마이저가 초기에 모멘텀을 유지할 수 있습니다. 이후 감쇠 속도는 스케줄 중반으로 갈수록 빨라졌다가, 최소 학습률에 가까워지면서 다시 느려집니다. 훈련 막바지의 이 마지막 감속 단계는 특히 중요합니다. 옵티마이저가 평탄할 가능성이 있는 최솟값에 조심스럽게 안착할 수 있게 해 주기 때문이며, 경험적 증거에 따르면 이러한 최솟값은 날카로운 최솟값보다 일반화 성능이 더 좋은 경우가 많습니다.

또한 코드 예제에서 볼 수 있듯이, 코사인 스케줄링은 흔히 “워밍업” 단계와 함께 사용합니다. 이 초기 단계(예: warmup_steps) 동안 학습률은 보통 아주 작은 값(또는 0)에서 주 초기 학습률까지 선형으로 증가합니다. 이 워밍업 기간은 훈련 맨 처음의 안정화에 도움이 됩니다. 특히 대규모 모델이나 데이터셋에서는 무작위로 초기화된 가중치에 큰 초기 학습률을 적용하면 불안정해지거나 발산할 수 있기 때문입니다. 워밍업이 끝나면 코사인 감쇠 단계가 시작되어, 남은 total_steps - warmup_steps 동안 학습률을 최댓값에서 목표 최솟값(base_lr * min_lr_ratio)까지 매끄럽게 낮춥니다. 이처럼 완만한 시작(워밍업)과 이론적 근거가 있는 매끄러운 감쇠(코사인 어닐링)를 결합하면 견고하고 효과적인 학습률 전략이 됩니다. 이 전략은 스텝 기반 스케줄보다 하이퍼파라미터 튜닝이 덜 필요한 경우가 많고, 모델 정확도를 높이는 경우도 많습니다.

PyTorch에는 이 기법을 구현한 CosineScheduler가 있습니다. 이를 단순화한 버전을 살펴봅시다:

class CosineAnnealingWarmupScheduler(_LRScheduler):
    def __init__(self, optimizer, warmup_steps, total_steps, min_lr_ratio=1e-4, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.min_lr_ratio = min_lr_ratio
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        if self.last_epoch < self.warmup_steps:
            # During warmup: linearly increase from 0 to base LR
            scale = float(self.last_epoch + 1) / float(max(1, self.warmup_steps))
            return [base_lr * scale for base_lr in self.base_lrs]
        else:
            # After warmup: cosine decay from base LR to min_lr
            progress = float(self.last_epoch - self.warmup_steps) / float(
                max(1, self.total_steps - self.warmup_steps)
            )
            # Cosine decay formula: min_lr + 0.5 * (base_lr - min_lr) * (1 + cos(pi * progress))
            scale = self.min_lr_ratio + 0.5 * (1.0 - self.min_lr_ratio) * (
                1.0 + math.cos(math.pi * progress)
            )
            return [base_lr * scale for base_lr in self.base_lrs]

AdamW 옵티마이저

AdamW(Adam with Decoupled Weight Decay)는 Adam과 같은 적응형 옵티마이저에서 가중치 감쇠(L2 정규화)를 구현하는 표준 방식에 숨어 있는 미묘한 문제를 해결합니다. 이 개선된 옵티마이저를 사용하면 가중치 감쇠가 모멘텀에도 분산에도 누적되지 않습니다. 기존 Adam에서는 이동 평균( $m_{t}$ 와 $v_{t}$ )을 계산하기 전에 감쇠 항( $λ \cdot weight$ )을 그래디언트에 직접 더하는 방식으로 L2 정규화를 구현하는 경우가 많습니다. 그러나 이렇게 하면 가중치 감쇠 효과가 그래디언트 모멘트에서 도출되는 적응형 학습률과 결합됩니다. 그 결과, 과거에 큰 그래디언트를 가졌던(따라서 $v_{t}$ 값이 더 큰) 매개변수는 실질적인 가중치 감쇠를 더 적게 받게 되며, 이는 균일한 정규화 압력을 적용한다는 목표에 어긋납니다. AdamW는 이 두 과정을 분리합니다. 즉, 표준 Adam 업데이트는 순수하게 그래디언트만을 기반으로 수행하고, 가중치 감쇠 단계는 별도로 가중치에 직접 적용합니다. 이를 통해 가중치가 그래디언트 이력과 무관하게 자신의 크기에 비례하여 감쇠하는 L2 정규화 본래의 동작을 사실상 되살립니다.

AdamW의 Core 업데이트 메커니즘은 대체로 표준 Adam 절차를 따르지만, 가중치 감쇠를 적용하는 방식을 수정합니다. 매개변수 $θ$ 에 대한 각 스텝 $t$ 에서, 먼저 편향된 1차 모멘트 추정값 $m_{t} = β_{1} \cdot m_{t - 1} + (1 - β_{1}) \cdot g_{t}$ 와 편향된 2차 원시 모멘트 추정값 $v_{t} = β_{2} \cdot v_{t - 1} + (1 - β_{2}) \cdot g_{t}^{2}$ 를 계산합니다. 여기서 $g_{t}$ 는 스텝 $t$ 에서의 그래디언트이고, $β_{1}$ , $β_{2}$ 는 지수 감쇠율입니다. 그런 다음 이 값들의 편향을 보정합니다: $\overset{m}{^}_{t} = m_{t} / (1 - β_{1}^{t})$ , $\overset{v}{^}_{t} = v_{t} / (1 - β_{2}^{t})$ . 결정적인 차이는 업데이트 규칙에 있습니다. AdamW는 가중치 감쇠를 $g_{t}$ 에 포함하는 대신, 먼저 매개변수에 가중치 감쇠를 직접 적용합니다: $θ_{t - 1}^{'} = θ_{t - 1} \cdot (1 - lr \cdot λ)$ . 여기서 $lr$ 은 학습률이고 $λ$ 는 weight_decay 계수입니다(p.data.mul_(1 - group['lr'] * group['weight_decay']) 줄에서 볼 수 있습니다). 그런 다음 편향 보정된 모멘트를 사용하는 표준 Adam 업데이트를 이렇게 감쇠된 가중치에 적용합니다: $θ_{t} = θ_{t - 1}^{'} - lr \cdot \overset{m}{^}_{t} / (\overset{v}{^}_{t} + ϵ)$ . 이는 이미 감쇠된 p.data에 대해 동작하는 코드 줄 p.data.addcdiv_(exp_avg, denom, value=-step_size)에 해당합니다.

딥러닝 실무에서 AdamW는 Adam과 비슷한 방식으로 사용됩니다. 모델의 매개변수와 함께 학습률( $lr$ ), 베타 값( $β_{1}$ , $β_{2}$ ), 엡실론( $ϵ$ ), 가중치 감쇠 계수( $λ$ ) 같은 주요 하이퍼파라미터를 제공하여 인스턴스를 생성합니다. 일반적으로 학습 루프 안에서 그래디언트를 계산한 후(loss.backward()) 호출되는 step() 메서드는 각 매개변수에 대해 위에서 설명한 업데이트 로직을 실행합니다. 가장 큰 장점은 일반화 성능이 향상된다는 점이며, 이는 정규화가 중요한 Transformer 같은 대형 모델을 학습할 때 특히 두드러집니다. AdamW는 가중치 감쇠의 강도가 적응형 학습률 스케일링과 무관하도록 보장하므로 하이퍼파라미터(특히 $lr$ 과 $λ$ )를 더 잘 튜닝할 수 있는 경우가 많으며, L2 정규화를 적용한 표준 Adam에 비해 처음 보는 데이터에서 더 나은 성능을 내는 모델로 이어질 수 있습니다.

PyTorch에는 고도로 최적화된 AdamW 구현이 있습니다. 이를 단순화한 버전을 살펴봅시다. 아래 예제는 최소한의 PyTorch 구현으로, 상태 변수(exp_avg, exp_avg_sq 등)를 초기화하고 매개변수 루프 안에서 분리된 가중치 감쇠와 모멘트 기반 업데이트를 수행하는 과정을 보여 줍니다.

class AdamW(Optimizer):
    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                 weight_decay=1e-2, amsgrad=False):
        defaults = dict(lr=lr, betas=betas, eps=eps,
                        weight_decay=weight_decay, amsgrad=amsgrad)
        super(AdamW, self).__init__(params, defaults)

    def step(self, closure=None):
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue

                # Get gradient
                grad = p.grad.data

                if grad.is_sparse:
                    raise RuntimeError('AdamW does not support sparse gradients')

                amsgrad = group['amsgrad']
                state = self.state[p]

                # State initialization
                if len(state) == 0:
                    state['step'] = 0
                    # Exponential moving average of gradient values
                    state['exp_avg'] = torch.zeros_like(p.data)
                    # Exponential moving average of squared gradient values
                    state['exp_avg_sq'] = torch.zeros_like(p.data)
                    if amsgrad:
                        # Maintains max of all exp. moving avg. of sq. grad. values
                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)

                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
                if amsgrad:
                    max_exp_avg_sq = state['max_exp_avg_sq']

                beta1, beta2 = group['betas']

                state['step'] += 1

                # Decay the first and second moment running average coefficient
                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

                if amsgrad:
                    # Maintains the maximum of all 2nd moment running avg. till now
                    torch.maximum(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
                    # Use the max. for normalizing running avg. of gradient
                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
                else:
                    denom = exp_avg_sq.sqrt().add_(group['eps'])

                bias_correction1 = 1 - beta1 ** state['step']
                bias_correction2 = 1 - beta2 ** state['step']
                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1

                # Apply weight decay BEFORE the optimization step
                if group['weight_decay'] != 0:
                    p.data.mul_(1 - group['lr'] * group['weight_decay'])

                # Update parameters
                p.data.addcdiv_(exp_avg, denom, value=-step_size)

        return loss

멀티 토큰 예측

멀티 토큰 예측은 자기회귀 언어 모델의 추론 속도를 높이기 위해 개발된 기법입니다. 일반적으로 자기회귀 생성은 토큰을 하나씩 예측합니다. 모델은 시퀀스를 입력받아 가장 가능성이 높은 다음 토큰 하나를 예측하고, 이를 시퀀스 뒤에 덧붙인 다음 이 과정을 반복합니다. 생성하는 토큰마다 모델의 순방향 패스를 한 번씩 온전히 거쳐야 하는 이러한 순차적 특성은 지연 시간에 민감한 애플리케이션에서 심각한 병목이 됩니다. 멀티 토큰 예측은 모델의 예측 헤드를 수정하여 현재 은닉 상태를 바탕으로 여러 미래 토큰의 확률을 동시에 출력하게 함으로써 이 문제를 극복하려 합니다. 이를 통해 주어진 길이의 시퀀스를 생성하는 데 필요한 순방향 패스 횟수를 줄입니다.

구현할 때는 보통 언어 모델의 마지막 층을 조정합니다. 멀티 토큰 예측기는 최종 은닉 상태를 다음 토큰에 대한 어휘 전체의 로짓으로 매핑하는 단일 출력 층(또는 “언어 모델 헤드”) 대신 여러 가지 전략을 사용할 수 있습니다. 한 가지 접근 방식은 별도의 헤드를 사용하는 예제 클래스에서 볼 수 있듯이, 서로 다른 미래 오프셋의 토큰을 예측하도록 각각 학습된 여러 개의 독립적인 예측 헤드를 두는 것입니다(예: 토큰 $t + 1$ 을 위한 헤드, $t + 2$ 를 위한 헤드, 이렇게 $t + N$ 까지). 이 헤드들은 보통 같은 최종 은닉 상태(예: 토큰 $t$ 에 해당하는 상태)를 입력으로 받지만, 각자 담당하는 시간 단계에 특화된 서로 다른 투영을 학습합니다. 또 다른 접근 방식은 하나의 공유 예측 헤드를 사용하는 것입니다. 이 경우 사실상 동일한 시작 표현으로부터 $N$ 개의 미래 토큰 각각에 대해 서로 다른 확률 분포를 생성해야 하므로, 은닉 상태의 학습된 변환이나 위치 정보 통합 같은 더 복잡한 메커니즘이 필요할 수 있습니다.

다중 토큰 예측기를 학습시킨다는 것은 앞선 문맥이 주어졌을 때 뒤따르는 $N$ 개 토큰의 시퀀스를 정확히 예측하도록 모델을 가르치는 것입니다. 학습 단계에서는 compute_loss 메서드에서 볼 수 있듯이, 모델이 입력 시퀀스를 받고, 다음 $N$ 개 토큰에 대한 모델의 예측을 학습 데이터에 있는 실제 $N$ 개 목표 토큰과 비교합니다. 손실 함수(보통 교차 엔트로피)는 예측된 각 위치( $t + 1$ 부터 $t + N$ 까지)마다 계산된 뒤 집계(예: 평균)되어 역전파에 사용되는 최종 손실 신호를 이룹니다.

이 방식은 속도 향상을 보여 줄 수 있지만, 다중 토큰 예측에는 몇 가지 단점이 있습니다. 더 먼 미래의 토큰일수록 예측 정확도가 떨어지는 경향이 있으며, 선택된 $N$ 개 토큰의 시퀀스가 단일 토큰 생성이었다면 택했을 최적 경로에서 벗어날 수 있습니다. 따라서 이는 흔히 생성 속도와 잠재적인 품질 저하 사이의 트레이드오프이며, 그 균형은 사용 사례에 크게 좌우됩니다.

class MultiTokenPredictor(nn.Module):
    def __init__(self, hidden_dim, vocab_size, num_predicted_tokens=2, shared_prediction_head=False):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.num_predicted_tokens = num_predicted_tokens
        self.shared_prediction_head = shared_prediction_head

        if shared_prediction_head:
            # Share the same prediction head for all positions
            self.lm_head = nn.Linear(hidden_dim, vocab_size)
        else:
            # Use separate prediction heads for each position
            self.lm_heads = nn.ModuleList([
                nn.Linear(hidden_dim, vocab_size)
                for _ in range(num_predicted_tokens)
            ])

    def forward(self, hidden_states):
        batch_size, seq_len, _ = hidden_states.shape

        # Get the hidden states for the last token
        last_hidden = hidden_states[:, -1]

        if self.shared_prediction_head:
            # Use the shared head for all predictions
            logits = self.lm_head(last_hidden)
            next_token_logits = logits.unsqueeze(1)

            # For positions beyond the first, we need to make a projection
            multi_token_logits = []
            for i in range(self.num_predicted_tokens):
                if i == 0:
                    multi_token_logits.append(logits)
                else:
                    # A simple projection for demonstration; in practice, this would be more complex
                    projected_hidden = last_hidden + i * 0.1  # Just a dummy transformation
                    multi_token_logits.append(self.lm_head(projected_hidden))

            multi_token_logits = torch.stack(multi_token_logits, dim=1)
        else:
            # Use separate heads for each position
            multi_token_logits = torch.stack([
                head(last_hidden) for head in self.lm_heads
            ], dim=1)

            # The first position is used for next-token prediction
            next_token_logits = multi_token_logits[:, 0:1]

        return next_token_logits, multi_token_logits

    def compute_loss(self, hidden_states, labels, ignore_index=-100):
        # Get predictions
        _, multi_token_logits = self.forward(hidden_states)

        # Prepare targets: shift labels to align with predictions
        # We need the next num_predicted_tokens tokens after the input
        shifted_labels = labels[:, :-self.num_predicted_tokens]
        targets = []

        for i in range(self.num_predicted_tokens):
            targets.append(labels[:, 1+i:1+i-self.num_predicted_tokens+1])

        # Stack targets
        stacked_targets = torch.stack(targets, dim=1)

        # Compute loss across all predicted positions
        loss = 0
        for i in range(self.num_predicted_tokens):
            loss += F.cross_entropy(
                multi_token_logits[:, i].view(-1, self.vocab_size),
                stacked_targets[:, i].reshape(-1),
                ignore_index=ignore_index
            )

        return loss / self.num_predicted_tokens

추측 디코딩

추측 디코딩(speculative decoding)은 대규모 자기회귀 언어 모델의 추론 과정을 가속하도록 고안된 영리한 기법으로, 텍스트 생성에 필요한 실제 경과 시간을 크게 줄여 줍니다. 표준 생성 방식은 순차적 특성 때문에 병목이 발생하는데, 계산 비용이 큰 대형 모델(“목표” 모델)이 한 번에 토큰 하나를 예측하기 위해 전체 순방향 패스를 수행해야 하기 때문입니다. 추측 디코딩은 훨씬 작고 빠른 보조 모델인 “초안” 모델을 도입합니다. 이 초안 모델은 앞으로 나올 후보 토큰들의 짧은 시퀀스(“초안”)를 빠르게 생성합니다. 그런 다음 핵심(Core) 아이디어는 큰 목표 모델을 사용하여 이 초안 시퀀스 전체를 단 한 번의 병렬 순방향 패스로 검증하는 것입니다. 이를 통해 여러 토큰을 한꺼번에 수락할 수 있으므로, 목표 모델의 계산 비용을 생성된 여러 토큰에 걸쳐 분산시킬 수 있습니다.

이 메커니즘의 관건은 초안 모델의 예측과 목표 모델의 예측을 비교하는 데 있습니다. 초안 모델이 $k$ 개 토큰의 시퀀스 $d_{1}, d_{2}, \dots, d_{k}$ 를 제안하면, 원래 입력 시퀀스에 초안을 이어 붙인 것에 대해 목표 모델을 한 번 실행합니다. 이 한 번의 패스로 초안 시퀀스 내 각 위치에서의 다음 토큰에 대한 목표 모델의 확률 분포를 얻습니다. 각 초안 토큰 $d_{i}$ 에 대해, 같은 위치에서의 목표 모델의 예측과 대조하여 그 타당성을 확인합니다. 목표 모델이 초안 토큰에 강하게 동의하면(예: 목표 모델도 $d_{i}$ 를 높은 확률로 예측했을 경우, 또는 $P_{target} (d_{i})$ 와 $P_{draft} (d_{i})$ 를 비교하는 특정 수락 규칙에 따라) 해당 토큰이 수락됩니다. 이 검증은 초안을 따라 순차적으로 진행됩니다.

이 과정은 $k$ 개의 초안 토큰이 모두 수락되거나 어떤 초안 토큰 $d_{j}$ 가 거부될 때까지 계속됩니다. 위치 $j$ 에서 거부가 발생하면, 그 앞에서 수락된 초안 토큰( $d_{1}, \dots, d_{j - 1}$ )은 모두 유지됩니다. 초안의 나머지 부분( $d_{j}, \dots, d_{k}$ )은 폐기됩니다. 중요한 점은, 위치 $j$ 에서 계산된 목표 모델의 확률 분포를 사용해 수락된 시퀀스 뒤에 덧붙일 보정된 토큰을 샘플링할 수 있다는 것이며, 이를 통해 전체 생성 결과가 통계적으로 더 강력한 목표 모델의 분포를 따르도록 보장합니다. 이 보정된 토큰은 앞서 수락된 토큰들과 함께 다음 회차 초안 생성의 입력이 됩니다. 목표 모델의 추론 단계마다 평균적으로 여러 토큰을 수락함으로써, 추측 디코딩은 생성된 텍스트의 품질에 미치는 영향을 최소화하면서 상당한 속도 향상(예: 2~3배)을 달성할 수 있습니다.

class SimpleSpeculativeDecoding:
    def __init__(self, target_model, draft_model, tokenizer, max_draft_tokens=5):
        self.target_model = target_model
        self.draft_model = draft_model
        self.tokenizer = tokenizer
        self.max_draft_tokens = max_draft_tokens

    def generate(self, prompt, max_length=100):
        # Start with the prompt's token IDs
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")

        # Keep generating until we reach max_length or an end token
        while input_ids.shape[1] < max_length:
            # Step 1: Generate multiple draft tokens
            draft_input_ids = input_ids.clone()
            draft_tokens = []

            with torch.no_grad():
                for _ in range(self.max_draft_tokens):
                    # Generate next token with draft model
                    outputs = self.draft_model(draft_input_ids)
                    next_token_logits = outputs.logits[:, -1, :]
                    next_token = torch.argmax(next_token_logits, dim=-1)

                    # Save the draft token
                    draft_tokens.append(next_token.item())

                    # Add to draft input
                    draft_input_ids = torch.cat([draft_input_ids, next_token.unsqueeze(0)], dim=1)

                    # Check for end token
                    if next_token.item() == self.tokenizer.eos_token_id:
                        break

            # Step 2: Verify with target model
            with torch.no_grad():
                # Get target model probabilities for all tokens including drafts
                verification_ids = torch.cat([input_ids, torch.tensor([draft_tokens]).to(input_ids.device)], dim=1)
                target_outputs = self.target_model(verification_ids)
                target_logits = target_outputs.logits[:, input_ids.shape[1]-1:]

                # Convert to probabilities
                target_probs = F.softmax(target_logits, dim=-1)

                # Accept tokens until mismatch or all drafts accepted
                accepted_tokens = []
                for i, token_id in enumerate(draft_tokens):
                    # Check if target model agrees with draft
                    target_prob = target_probs[0, i, token_id].item()
                    draft_prob = 1.0  # Draft chose this with highest probability

                    # Accept based on probability ratio (simplified)
                    if random.random() < min(1.0, target_prob / draft_prob):
                        accepted_tokens.append(token_id)
                    else:
                        # Rejection: get new token from target model
                        new_token = torch.argmax(target_logits[0, i]).item()
                        accepted_tokens.append(new_token)
                        break

            # Add accepted tokens to input_ids
            input_ids = torch.cat([input_ids, torch.tensor([accepted_tokens]).to(input_ids.device)], dim=1)

            # Check for end token
            if input_ids[0, -1].item() == self.tokenizer.eos_token_id:
                break

        # Decode the generated tokens
        return self.tokenizer.decode(input_ids[0])

Keyboard shortcuts

Stephen Diehl 블로그 한국어 번역