라우터, Top-K, 공유 전문가

라우터가 길을 정해

라우터는 보통 토큰의 은닉 상태를 전문가별 로짓으로 바꾸는 선형 층 하나에 가까워. 크기는 작지만 어떤 전문가가 토큰을 받을지 모두 결정해. 모델의 나머지 부분과 함께 학습되며, 라우터가 잘못 배우면 전문가도 엉뚱한 학습 신호를 받아 무너져.

Top-K 선택

라우터는 토큰마다 전문가 N명에 대한 점수를 내. Top-K는 가장 높은 K개만 남기고 나머지는 0으로 만들어. 남은 점수는 softmax나 sigmoid를 거쳐 가중치가 되고, 선택된 전문가들의 출력에 곱해져 하나로 합쳐져.

K는 얼마로 잡을까

Top-1: Switch Transformer와 Llama 4 Scout가 써. 연산은 가장 적지만 한 전문가로 쏠리기 쉬워.
Top-2: Mixtral이 택한 방식이야. 두 전문가가 라우팅 실수를 보완하는 대신 연산량이 조금 늘어.
Top-6~8: DeepSeek의 세분화된 전문가 설계가 써. 작은 전문가를 많이 두고 토큰마다 여럿을 골라 더 미세한 조합을 만들어.

공유 전문가는 늘 켜져 있어

DeepSeek-V2·V3와 Gemma 4 MoE는 라우터의 선택과 관계없이 모든 토큰이 지나는 공유 전문가를 한두 개 둬. 흔한 패턴은 이 안정적인 경로가 맡고, 라우팅되는 전문가들은 더 좁은 패턴에 집중할 수 있어. 작은 구조 변화지만 학습 안정성에는 큰 도움이 돼.

Sigmoid와 softmax 라우팅

초기 MoE는 주로 전문가 로짓에 softmax를 썼어. DeepSeek-V3는 각 전문가 점수를 독립적으로 매기는 sigmoid 라우팅을 사용하고, 명시적인 보조 손실 대신 전문가별 편향을 학습해 부하를 맞춰. 주 학습 목표를 보조 손실로 흔들지 않으면서 전문가 붕괴를 피하려는 설계야.

Code

top-K와 softmax 가중치로 라우팅·python

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

class TopKRouter(nn.Module):
    def __init__(self, d_model, num_experts, k=2):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        self.k = k

    def forward(self, x):
        logits = self.gate(x)                  # (..., num_experts)
        topk_vals, topk_idx = logits.topk(self.k, dim=-1)
        weights = F.softmax(topk_vals, dim=-1)
        return topk_idx, weights

DeepSeek-V3식 편향 포함 sigmoid 라우팅·python

class SigmoidRouter(nn.Module):
    def __init__(self, d_model, num_experts, k=8):
        super().__init__()
        self.gate = nn.Linear(d_model, num_experts, bias=False)
        # Learned per-expert bias for load balancing (no aux loss needed).
        self.expert_bias = nn.Parameter(torch.zeros(num_experts))
        self.k = k

    def forward(self, x):
        scores = torch.sigmoid(self.gate(x)) + self.expert_bias  # (..., num_experts)
        topk_vals, topk_idx = scores.topk(self.k, dim=-1)
        # Sigmoid scores are independent; no softmax over selected experts.
        return topk_idx, topk_vals

라우터가 길을 정해

Top-K 선택

K는 얼마로 잡을까

공유 전문가는 늘 켜져 있어

Sigmoid와 softmax 라우팅

Code

External links

Exercise

Progress

댓글 0