LLM Architecture & Inference Mechanics

Transformer #

Most modern LLMs, such as GPT, Llama, and Mistral, follow an Autoregressive Decoder-only architecture.

This means they take text input and predict the next token one by one.

• Embedding Layer: Transforms text tokens into high-dimensional vectors d.
• Residual Connection: A structure that adds the output of each layer back to its input: x + sublayer(x). This prevents vanishing gradients and ensures information is not lost even in deep layers.
• Layer Normalization: Normalizes the output of each layer for training stability. Recently, Pre-Norm, which normalizes the input first, and RMSNorm are commonly used for model efficiency.
• Feed-Forward Network (FNN): A two-layer MLP structure that independently applies a non-linear transformation to each token vector after attention. SwiGLU activation function is typically used instead of ReLU to enhance expressiveness.

Multi-Head Self-Attention (MHSA) Mechanism #

This is a core mathematical operation that determines which information the model should focus on within the context.

Q, K, and V stand for Query, Key, and Value, respectively, and are generated by multiplying the input vector x by their respective weight matrices wq, wk, and wv.

• $QK^T$: Calculates the similarity (dot product) between tokens.
• $\sqrt{d_k}$: A scaling factor that prevents the Softmax gradient from becoming too flat as the dimension increases, which would happen due to larger dot product values.

Positional Encoding #

Since Transformers lack a recurrent structure like RNNs, token position information must be explicitly provided.

RoPE: Injects the relative distance between two tokens, instead of absolute positions, through a complex rotation matrix.

It exhibits excellent extrapolation capabilities for sentences longer than those seen during training and possesses a distance decay property, allowing it to better capture relationships between closer words than distant ones.

Inference Mechanics & Optimization #

Physical mechanisms that determine performance and resource efficiency during inference:

• KV Caching: A technique where k and v vectors computed for tokens 1 through t-1 are pre-stored in memory when generating the t-th token during the decoding phase. This avoids recomputing the entire sequence at each step, only processing the newly added token.
• Decoding Strategies:
  • Greedy Search: Selects only the token with the highest probability.
  • Nucleus Sampling (Top-p): Samples from a candidate set whose cumulative probability is p to ensure sentence diversity.
• Quantization: A technique that reduces memory footprint by lowering the precision of weight parameters (e.g., FP16 -> INT8/INT4).

Code Example #

A Transformer Block & Attention can be implemented as shown below.

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # WQ, WK, WV 가중치 행렬
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.size()
        
        # 1. 선형 변환 및 Head 분리
        # (batch, seq_len, num_heads, d_k) -> (batch, num_heads, seq_len, d_k)
        q = self.W_q(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        # 2. Scaled Dot-Product Attention 계산
        # $attn\_scores = \frac{QK^T}{\sqrt{d_k}}$
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
            
        attn_weights = torch.softmax(attn_scores, dim=-1)
        
        # 3. Value와 곱함
        output = torch.matmul(attn_weights, v)
        
        # 4. 원래 차원으로 복구
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        return self.W_o(output)