MoE와 다중 토큰 예측

Jayce_97·2025년 3월 19일
0

DeepSeek

목록 보기
3/4

이전 포스팅 글에서 DeepSeek에 대해 작성하며, MoE와 다중 토큰 예측을 다루어 좀 더 자세히 작성하기 위한 글입니다.

MoE

MoE (Mixture of Experts)란

MoE(Mixture of Experts)는 여러 개의 전문가(Experts) 모델을 두고, 입력에 따라 적절한 전문가를 선택해 학습과 예측을 수행하는 모델 구조입니다.
-> 큰 모델을 효율적으로 사용하면서 연산량을 줄일 수 있음

MoE의 핵심 개념

여러 개의 전문가(Experts) 신경망이 존재
입력 데이터를 보고 Gating Network가 가장 적절한 전문가를 선택
선택된 전문가만 활성화되어 예측 수행 (모든 전문가가 활성화되는 것이 아님)
특정 데이터 유형에 특화된 전문가를 선택하므로 효율적

MoE 구조

Experts (전문가 네트워크)

각각 특정한 데이터 패턴을 학습하는 전문가 네트워크들
예를 들어, 자연어 처리 모델에서

  • Expert A: 문법적 분석
  • Expert B: 감정 분석
  • Expert C: 요약
    -> 입력 데이터에 따라 적절한 전문가를 선택해서 사용

Gating Network (게이팅 네트워크)

입력 데이터를 보고 어떤 전문가를 선택할지 결정 하는 네트워크
Softmax를 이용해 각 전문가에 대한 가중치 계산
일부 전문가만 활성화되도록 하여 연산량 감소

Sparse Activation (희소 활성화)

모든 전문가를 사용하는 것이 아니라 소수의 전문가만 활성화
예를 들어, 총 16개의 전문가 중 2~4개만 사용
모델이 크더라도 연산량이 줄어듦

MoE 장단점

MoE의 장점

  • 효율적인 연산: 일부 전문가만 활성화되므로 계산량이 줄어듦
  • 특정 태스크에 특화된 학습 가능: 각 전문가가 특정 데이터 패턴을 더 잘 학습
  • 대규모 모델 확장 가능: 전문가 수를 늘려도 일부만 사용하므로 메모리 사용 절약

MoE의 단점

  • Gating Network가 잘못 동작하면 비효율 발생
  • 일부 전문가만 자주 선택되면 학습 불균형 문제 발생
  • 일반적인 Dense 모델보다 구현이 복잡

구현 방법

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

# Expert 정의
class Expert(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Expert, self).__init__()
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return F.relu(self.fc(x))

# Gating Network 정의
class GatingNetwork(nn.Module):
    def __init__(self, input_dim, num_experts):
        super(GatingNetwork, self).__init__()
        self.fc = nn.Linear(input_dim, num_experts)

    def forward(self, x):
        return F.softmax(self.fc(x), dim=-1)

# MoE 모델 정의
class MoE(nn.Module):
    def __init__(self, input_dim, output_dim, num_experts):
        super(MoE, self).__init__()
        self.experts = nn.ModuleList([Expert(input_dim, output_dim) for _ in range(num_experts)])
        self.gate = GatingNetwork(input_dim, num_experts)

    def forward(self, x):
        gate_weights = self.gate(x)  # 어떤 전문가를 선택할지 결정
        expert_outputs = torch.stack([expert(x) for expert in self.experts], dim=1)
        output = torch.sum(gate_weights.unsqueeze(-1) * expert_outputs, dim=1)  # 선택된 전문가 결과 조합
        return output

# 모델 생성
model = MoE(input_dim=128, output_dim=10, num_experts=4)
x = torch.randn(32, 128)  # Batch size = 32
output = model(x)

결론

  • MoE는 여러 전문가 네트워크를 활용하여 효율적이고 확장 가능한 모델 구조
  • Sparse Activation을 통해 연산량 절감
  • Gating Network가 중요 → 잘못 학습되면 성능 저하 가능
  • DeepSeek-V3 같은 최신 모델에서도 MoE를 활용하여 빠르고 정확한 예측 가능

Multi-Token Prediction (MTP)

Multi-Token Prediction (MTP)란

기존의 Autoregressive 방식(한 번에 한 토큰씩 예측)과 달리, MTP는 한 번에 여러 개의 토큰을 동시에 예측하는 방법입니다.
-> 한 번의 Forward Pass에서 여러 토큰을 예측하여 속도를 향상

MTP의 구조 및 개념

기존 Autoregressive 방식 vs. MTP

방식설명장점단점
Autoregressive한 번에 한 개의 토큰을 생성, 이전 토큰을 기반으로 다음 토큰 예측높은 정확도속도가 느림
MTP한 번에 여러 개의 토큰을 병렬로 예측속도가 빠름품질 저하 가능

MTP 방식의 동작 원리

  • 예측할 토큰 수(k)를 설정
  • 모델이 한 번의 Forward Pass에서 k개의 토큰을 동시에 예측
  • 각 토큰의 위치에 대해 병렬적으로 Loss를 계산하고 업데이트
    -> 즉, 예측 속도를 높이기 위한 방법으로 사용됨

MTP의 장점

  • 속도 향상: 한 번에 여러 개의 토큰을 예측하므로 텍스트 생성 속도가 증가
  • 더 나은 병렬 처리: GPU 활용도가 증가하여 연산 효율성이 높아짐
  • 디코딩 과정의 연산 감소: Autoregressive 방식보다 예측 과정이 최적화됨

MTP의 단점

정확도 저하 가능성
*여러 토큰을 동시에 예측하면 문맥을 정확하게 반영하기 어려움

  • 특히, 긴 문장에서 문맥 오류 발생 가능 -> Beam Search 등과 결합하기 어려움
  • 기존 Greedy Decoding, Beam Search 같은 기법과 자연스럽게 통합하기 어려움 -> 모델 학습이 더 복잡함
  • 기존 방식보다 학습 과정이 더 정교하게 조정되어야 함

MTP 구현 방법

기존 Autoregressive 방식 (GPT-Style)

import torch
import torch.nn as nn

class AutoRegressiveModel(nn.Module):
    def __init__(self, vocab_size, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.rnn = nn.GRU(hidden_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        output, _ = self.rnn(x)
        output = self.fc(output[:, -1, :])  # 마지막 토큰만 예측
        return output  # (batch_size, vocab_size)

# 모델 생성
model = AutoRegressiveModel(vocab_size=5000, hidden_dim=256)

MTP 방식 (Multi-Token Prediction)

class MultiTokenPredictionModel(nn.Module):
    def __init__(self, vocab_size, hidden_dim, k=5):  # k: 동시에 예측할 토큰 수
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.rnn = nn.GRU(hidden_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size * k)  # k개의 토큰을 예측

    def forward(self, x):
        x = self.embedding(x)
        output, _ = self.rnn(x)
        output = self.fc(output[:, -1, :])  # (batch_size, vocab_size * k)
        output = output.view(output.size(0), -1, vocab_size)  # (batch_size, k, vocab_size)
        return output  # k개의 토큰을 예측하는 형태

# 모델 생성
model = MultiTokenPredictionModel(vocab_size=5000, hidden_dim=256, k=5)

기존 Autoregressive 모델과 차이점으로는

  • 기존 모델은 한 번에 1개의 토큰을 예측
  • MTP 모델은 한 번에 k개의 토큰을 예측

결론

  • MTP는 기존 Autoregressive 방식보다 텍스트 생성 속도를 증가시킴
  • 여러 토큰을 병렬 예측하여 디코딩 속도를 높일 수 있음
  • 하지만 정확도 저하 가능성이 있고 Beam Search 같은 방식과 결합이 어려움
  • 최신 대규모 모델에서도 효율적인 디코딩 방식으로 연구 중

Greedy Decoding

한 번에 가장 확률이 높은 단어만 선택하는 방식

  • 매 단계에서 가장 확률이 높은 단어를 선택
  • 선택한 단어를 기반으로 다음 단어를 예측

장점

  • 연산이 간단하고 빠름
  • 구현이 쉬움

단점

  • 항상 최적의 문장을 생성하는 것이 아님 (전역 최적 해를 찾지 못할 수도 있음)
  • 초반에 잘못된 단어를 선택하면 수정 불가능

매 단계에서 확률이 높은 여러 개의 후보를 유지하면서 탐색하는 방식

  • 한 번에 하나의 최적 경로만 따르는 Greedy Decoding과 달리, Beam Search는 여러 개의 후보(Beam)를 유지하면서 가장 좋은 문장을 탐색
  • Beam Width 𝑘를 설정하여, 각 단계에서 확률이 높은 𝑘개의 문장을 유지

장점

  • Greedy Decoding보다 더 좋은 문장을 생성할 확률이 높음
  • 전역 최적 해를 찾을 가능성이 증가

단점

  • 연산량이 많음 (Beam Width가 클수록 계산 비용 증가)
  • 완전한 최적 해를 보장하지는 않음
profile
Jayce_97
AI (ML/DL) 학습
이전 포스트

DeepSeek-V3: MoE 기반 AI 모델

다음 포스트

DeepSeek의 MLA (Multi-Head Latent Attention)

0개의 댓글

관련 채용 정보