임베딩 (Embeddings)

토큰을 연속적인 벡터 공간으로 매핑하는 과정. 모델이 텍스트의 의미를 이해하는 기반이 됨.

Token Embedding

개념

어휘 크기: V (예: 50,000)
임베딩 차원: d (예: 768)
임베딩 행렬: E ∈ R^(V × d)

토큰 ID i → E[i] → d차원 벡터

구현

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        # x: (batch, seq_len) - 토큰 ID
        # 출력: (batch, seq_len, d_model)
        return self.embedding(x) * (self.d_model ** 0.5)  # 스케일링

# 사용
vocab_size = 50000
d_model = 768
embedding = TokenEmbedding(vocab_size, d_model)

token_ids = torch.tensor([[1, 5, 100, 2]])  # (1, 4)
vectors = embedding(token_ids)  # (1, 4, 768)

초기화

# 표준 정규분포
nn.init.normal_(embedding.weight, mean=0, std=0.02)

# Xavier/Glorot
nn.init.xavier_uniform_(embedding.weight)

# 사전 학습된 임베딩 로드
pretrained = torch.load('embeddings.pt')
embedding.weight.data = pretrained

Position Embedding

Transformer는 순서 정보가 없으므로 위치 정보 추가 필요.

Sinusoidal (원본 Transformer)

import math

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()

        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * 
            -(math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        seq_len = x.size(1)
        return x + self.pe[:, :seq_len]

Learned (GPT)

class LearnedPositionalEmbedding(nn.Module):
    def __init__(self, max_seq_len, d_model):
        super().__init__()
        self.position_embedding = nn.Embedding(max_seq_len, d_model)

    def forward(self, x):
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device)
        return x + self.position_embedding(positions)

Rotary Position Embedding (RoPE)

LLaMA, Qwen 등에서 사용. 상대적 위치 정보를 회전으로 인코딩.

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=2048, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

        t = torch.arange(max_seq_len).float()
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer('cos_cached', emb.cos())
        self.register_buffer('sin_cached', emb.sin())

    def forward(self, q, k, seq_len):
        cos = self.cos_cached[:seq_len]
        sin = self.sin_cached[:seq_len]

        q_embed = (q * cos) + (self._rotate_half(q) * sin)
        k_embed = (k * cos) + (self._rotate_half(k) * sin)

        return q_embed, k_embed

    def _rotate_half(self, x):
        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
        return torch.cat((-x2, x1), dim=-1)

ALiBi (Attention with Linear Biases)

BLOOM, MPT 등에서 사용. 위치 임베딩 대신 attention bias.

def get_alibi_slopes(num_heads):
    """ALiBi slopes 계산"""
    def get_slopes_power_of_2(n):
        start = 2 ** (-(2 ** -(math.log2(n) - 3)))
        return [start * (start ** i) for i in range(n)]

    if math.log2(num_heads).is_integer():
        return get_slopes_power_of_2(num_heads)
    else:
        closest_power = 2 ** math.floor(math.log2(num_heads))
        return (
            get_slopes_power_of_2(closest_power) +
            get_alibi_slopes(2 * closest_power)[0::2][:num_heads - closest_power]
        )

class ALiBiAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        slopes = torch.tensor(get_alibi_slopes(num_heads)).view(num_heads, 1, 1)
        self.register_buffer('slopes', slopes)

    def forward(self, q, k, v):
        seq_len = q.size(2)

        # 위치 차이 행렬
        positions = torch.arange(seq_len, device=q.device)
        distance = positions.unsqueeze(0) - positions.unsqueeze(1)

        # ALiBi bias
        alibi_bias = self.slopes * distance.abs().unsqueeze(0)

        # Attention with bias
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))
        scores = scores - alibi_bias

        return torch.matmul(F.softmax(scores, dim=-1), v)

임베딩 공간

의미적 관계

유사한 의미의 단어는 가까운 벡터.

def cosine_similarity(v1, v2):
    return torch.dot(v1, v2) / (v1.norm() * v2.norm())

# 예: "king" - "man" + "woman" ≈ "queen"
king = embedding(tokenizer.encode("king"))
man = embedding(tokenizer.encode("man"))
woman = embedding(tokenizer.encode("woman"))

result = king - man + woman
# result와 가장 가까운 토큰 찾기

시각화

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

# 임베딩 추출
words = ["king", "queen", "man", "woman", "cat", "dog"]
embeddings = [embedding(tokenizer.encode(w)).mean(dim=1).squeeze() for w in words]
embeddings = torch.stack(embeddings).detach().numpy()

# t-SNE 적용
tsne = TSNE(n_components=2, perplexity=min(5, len(words)-1))
reduced = tsne.fit_transform(embeddings)

# 플롯
plt.scatter(reduced[:, 0], reduced[:, 1])
for i, word in enumerate(words):
    plt.annotate(word, (reduced[i, 0], reduced[i, 1]))
plt.show()

Sentence/Document Embedding

문장이나 문서 전체를 하나의 벡터로 표현.

평균 풀링

def mean_pooling(token_embeddings, attention_mask):
    """마스크를 고려한 평균"""
    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size())
    sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
    sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
    return sum_embeddings / sum_mask

[CLS] 토큰 (BERT)

# BERT의 첫 번째 토큰이 문장 표현
outputs = bert_model(input_ids, attention_mask)
sentence_embedding = outputs.last_hidden_state[:, 0, :]  # [CLS] 토큰

Sentence Transformers

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2')

sentences = ["This is a sentence", "This is another sentence"]
embeddings = model.encode(sentences)

# 유사도 계산
from sklearn.metrics.pairwise import cosine_similarity
similarity = cosine_similarity([embeddings[0]], [embeddings[1]])

Weight Tying

임베딩 행렬과 LM head 가중치 공유.

class GPTModel(nn.Module):
    def __init__(self, vocab_size, d_model, ...):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

        # Weight tying
        self.lm_head.weight = self.token_embedding.weight

장점: - 파라미터 수 감소 - 정규화 효과 - 일반화 성능 향상

임베딩 모델 비교

모델	차원	특징
all-MiniLM-L6-v2	384	빠르고 효율적
all-mpnet-base-v2	768	높은 품질
text-embedding-ada-002	1536	OpenAI API
text-embedding-3-large	3072	최고 성능
bge-large-en-v1.5	1024	오픈소스 SOTA

실무 트레이드오프

Position Encoding 선택

방식	장점	단점	사용 모델
Sinusoidal	외삽 가능, 학습 불필요	성능 낮음	원본 Transformer
Learned	높은 성능	외삽 불가 (max_len 고정)	GPT-2, BERT
RoPE	외삽 가능, 상대 위치	구현 복잡	LLaMA, Qwen, Mistral
ALiBi	외삽 우수, 학습 불필요	성능 약간 낮음	BLOOM, MPT

임베딩 모델 선택 가이드

# 사용 케이스별 추천

# 1. 빠른 프로토타입 (속도 우선)
model = SentenceTransformer('all-MiniLM-L6-v2')  # 384차원, 빠름

# 2. 높은 품질 (정확도 우선)
model = SentenceTransformer('all-mpnet-base-v2')  # 768차원

# 3. 다국어 지원
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

# 4. 검색 특화 (Retrieval)
model = SentenceTransformer('multi-qa-mpnet-base-dot-v1')

# 5. 비대칭 검색 (Query vs Document)
model = SentenceTransformer('msmarco-distilbert-base-v4')

임베딩 차원 선택

차원별 특성:

384 (MiniLM):
  + 저장 공간 절약
  + 유사도 계산 빠름
  - 정보 손실 가능

768 (BERT-base):
  + 균형 잡힌 선택
  + 대부분의 태스크에 충분

1536+ (OpenAI ada-002):
  + 높은 표현력
  - 저장 비용 증가
  - 차원의 저주 (고차원에서 유사도 차이 감소)

컨텍스트 길이 확장 (RoPE)

# LLaMA 등에서 RoPE 기반 컨텍스트 확장
# YaRN, Code LLaMA 방식

def apply_yarn_scaling(rope_freq, scale_factor, original_max_len):
    """YaRN: Yet another RoPE extensioN"""
    # NTK-aware interpolation
    base = 10000
    dim = len(rope_freq) * 2

    # Dynamic NTK scaling
    new_base = base * (
        (scale_factor * original_max_len / original_max_len) 
        ** (dim / (dim - 2))
    )

    return rope_freq * (base / new_base)

# 4K → 16K 확장 가능 (품질 약간 저하)

흔한 실수

실수	증상	해결책
정규화 누락	코사인 유사도 이상	F.normalize() 적용
배치 처리 미사용	속도 느림	encode() 배치 사용
풀링 방식 불일치	성능 저하	모델 권장 풀링 확인
길이 초과	정보 손실	긴 텍스트 청킹

임베딩 저장 및 검색

import numpy as np
import faiss

class EmbeddingStore:
    """효율적인 임베딩 저장/검색"""

    def __init__(self, dimension, use_gpu=False):
        self.dimension = dimension

        # FAISS 인덱스 (IVF + PQ로 메모리 절약)
        quantizer = faiss.IndexFlatL2(dimension)
        self.index = faiss.IndexIVFPQ(
            quantizer, dimension, 
            100,      # nlist: 클러스터 수
            8,        # M: PQ 서브벡터 수
            8         # nbits: 각 서브벡터 비트
        )

        if use_gpu:
            self.index = faiss.index_cpu_to_gpu(
                faiss.StandardGpuResources(), 0, self.index
            )

    def train_and_add(self, embeddings):
        """학습 후 추가 (IVF 필수)"""
        embeddings = np.array(embeddings).astype('float32')
        self.index.train(embeddings)
        self.index.add(embeddings)

    def search(self, query_embedding, k=10):
        """유사 벡터 검색"""
        query = np.array([query_embedding]).astype('float32')
        distances, indices = self.index.search(query, k)
        return indices[0], distances[0]

# 100만 개 768차원 임베딩:
# - 원본: ~3GB
# - PQ 압축: ~100MB
# - 검색 속도: <10ms

참고 자료

• Sentence-BERT
• RoPE Paper
• ALiBi Paper
• Sentence Transformers
• YaRN: Efficient Context Window Extension
• MTEB Benchmark (임베딩 모델 비교)