Skip to content

필사 모드: Mixture of Experts(MoE)アーキテクチャ完全分析

日本語
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.
원문 렌더가 준비되기 전까지 텍스트 가이드로 표시합니다.

1. MoEとは何か?

**Mixture of Experts(MoE)**は、モデルの全パラメータの一部のみを活性化して演算効率を高めるアーキテクチャです。すべての入力に対してモデル全体を使用するDenseモデルとは異なり、MoEは入力に応じて**最適なエキスパート(Expert)**のみを選択して処理します。

Dense vs Sparseモデル

- **Denseモデル**: すべてのパラメータが毎回活性化(例:LLaMA、GPT-4)

- **Sparse MoE**: 全パラメータの一部のみ活性化(例:Mixtral、DeepSeek-V3)

核心的な利点は、**パラメータ数は大きいが演算量は少ない**ということです。Mixtral 8x7Bは合計46.7Bのパラメータを持ちますが、推論時に活性化されるパラメータは約12.9Bに過ぎません。

2. MoEアーキテクチャの核心構成要素

Expert Network

各Expertは独立したFFN(Feed-Forward Network)です:

class Expert(nn.Module):

def __init__(self, d_model: int, d_ff: int):

super().__init__()

self.w1 = nn.Linear(d_model, d_ff, bias=False)

self.w2 = nn.Linear(d_ff, d_model, bias=False)

self.w3 = nn.Linear(d_model, d_ff, bias=False) # SwiGLU gate

def forward(self, x: torch.Tensor) -> torch.Tensor:

SwiGLU activation

return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))

Router(Gating Network)

Routerは各トークンをどのExpertに送るかを決定します:

class TopKRouter(nn.Module):

def __init__(self, d_model: int, num_experts: int, top_k: int = 2):

super().__init__()

self.gate = nn.Linear(d_model, num_experts, bias=False)

self.top_k = top_k

def forward(self, x: torch.Tensor):

x shape: (batch, seq_len, d_model)

logits = self.gate(x) # (batch, seq_len, num_experts)

top_k_logits, top_k_indices = logits.topk(self.top_k, dim=-1)

top_k_weights = torch.softmax(top_k_logits, dim=-1)

return top_k_weights, top_k_indices

MoEレイヤー全体の実装

class MoELayer(nn.Module):

def __init__(self, d_model: int, d_ff: int,

num_experts: int = 8, top_k: int = 2):

super().__init__()

self.experts = nn.ModuleList([

Expert(d_model, d_ff) for _ in range(num_experts)

])

self.router = TopKRouter(d_model, num_experts, top_k)

self.num_experts = num_experts

def forward(self, x: torch.Tensor) -> torch.Tensor:

batch_size, seq_len, d_model = x.shape

weights, indices = self.router(x)

Reshape for expert processing

flat_x = x.view(-1, d_model)

flat_weights = weights.view(-1, weights.shape[-1])

flat_indices = indices.view(-1, indices.shape[-1])

output = torch.zeros_like(flat_x)

for i, expert in enumerate(self.experts):

Find tokens routed to this expert

mask = (flat_indices == i).any(dim=-1)

if mask.any():

expert_input = flat_x[mask]

expert_output = expert(expert_input)

Weight by router probability

idx = (flat_indices[mask] == i).float()

w = (flat_weights[mask] * idx).sum(dim=-1, keepdim=True)

output[mask] += w * expert_output

return output.view(batch_size, seq_len, d_model)

3. 主要MoEモデル分析

Mixtral 8x7B(Mistral AI)

- 8つのExpert、Top-2ルーティング

- 合計46.7Bパラメータ、活性12.9B

- Attentionレイヤーは共有、FFNのみExpertに分離

DeepSeek-V3 MoE

DeepSeek-V3はより精巧なMoE設計を採用しています:

class DeepSeekMoE(nn.Module):

"""DeepSeek-V3スタイル:共有Expert + Routed Expert"""

def __init__(self, d_model, d_ff, num_shared=1,

num_routed=256, top_k=8):

super().__init__()

すべてのトークンが通過する共有Expert

self.shared_experts = nn.ModuleList([

Expert(d_model, d_ff) for _ in range(num_shared)

])

トークンごとに選択されるRouted Expert

self.routed_experts = nn.ModuleList([

Expert(d_model, d_ff // 4) for _ in range(num_routed)

])

self.router = TopKRouter(d_model, num_routed, top_k)

def forward(self, x):

共有Expertの結果

shared_out = sum(e(x) for e in self.shared_experts)

Routed Expertの結果

weights, indices = self.router(x)

routed_out = self._route_tokens(x, weights, indices)

return shared_out + routed_out

- 1つの共有Expert + 256のRouted Expert(Top-8選択)

- 合計671Bパラメータ、活性37B

- **Auxiliary-loss-freeロードバランシング**を導入

モデル比較

| モデル | 総パラメータ | 活性パラメータ | Expert数 | Top-K |

| ------------- | ------------ | -------------- | -------- | ----- |

| Mixtral 8x7B | 46.7B | 12.9B | 8 | 2 |

| Mixtral 8x22B | 141B | 39B | 8 | 2 |

| DeepSeek-V3 | 671B | 37B | 256+1 | 8+1 |

| Qwen2.5-MoE | 14.3B | 2.7B | 60+4 | 4+4 |

4. ルーティング戦略

Token Choice vs Expert Choice

Token Choice: 各トークンがExpertを選択

def token_choice_routing(logits, top_k=2):

top_k_vals, top_k_idx = logits.topk(top_k, dim=-1)

weights = torch.softmax(top_k_vals, dim=-1)

return weights, top_k_idx

Expert Choice: 各Expertがトークンを選択

def expert_choice_routing(logits, capacity_factor=1.25):

num_tokens = logits.shape[0]

num_experts = logits.shape[1]

capacity = int(num_tokens * capacity_factor / num_experts)

expert_scores = logits.T # (num_experts, num_tokens)

top_k_vals, top_k_idx = expert_scores.topk(capacity, dim=-1)

return top_k_vals, top_k_idx

5. ロードバランシング

Expert間の負荷不均衡はMoEの核心的な課題です:

def load_balancing_loss(router_logits, top_k_indices, num_experts):

"""Auxiliary load balancing loss (Switch Transformer方式)"""

Expert別トークン比率

mask = torch.zeros_like(router_logits)

mask.scatter_(-1, top_k_indices, 1.0)

tokens_per_expert = mask.float().mean(dim=0) # (num_experts,)

Expert別ルーティング確率の平均

router_probs = torch.softmax(router_logits, dim=-1)

router_prob_per_expert = router_probs.mean(dim=0)

2つの分布の内積 = 不均衡の尺度

loss = num_experts * (tokens_per_expert * router_prob_per_expert).sum()

return loss

DeepSeek-V3の**Auxiliary-loss-free**方式は、Expert毎のbias termを動的に調整し、トークンを多く受け取るExpertのbiasを下げ、少ないExpertのbiasを上げることで、別途のlossなしに均衡を維持します。

6. 推論最適化

Expert Parallelism: Expertを複数のGPUに分散

GPU 0: Expert 0-3, GPU 1: Expert 4-7

class ExpertParallel(nn.Module):

def __init__(self, experts_per_gpu, rank, world_size):

super().__init__()

self.local_experts = nn.ModuleList([

Expert(d_model, d_ff)

for _ in range(experts_per_gpu)

])

self.rank = rank

self.world_size = world_size

def forward(self, x, indices):

All-to-All通信でトークンを再配分

dispatched = all_to_all(x, indices, self.world_size)

ローカルExpertで処理

output = self._process_local(dispatched)

結果を再結合

return all_to_all(output, indices, self.world_size)

7. クイズ

約**12.9B**パラメータです。8つのExpertのうちTop-2のみが活性化されるため、2つのExpertのFFNパラメータ + 共有されるAttentionパラメータのみが使用されます。全体46.7Bの約28%のみが活性化されることになります。

従来のMoEはロードバランシングのために別途のauxiliary lossを追加しますが、これはモデル性能を低下させる可能性があります。DeepSeek-V3は各Expertに動的な**bias term**を設け、トークンを多く受け取るExpertのbiasを下げ、少ないExpertのbiasを上げることで自然に均衡を取ります。これにより別途のlossなしで安定したロードバランシングが可能になります。

- **Token Choice**: 各トークンがTop-K Expertを選択します。実装が簡単ですが、一部のExpertにトークンが集中する不均衡が発生する可能性があります。

- **Expert Choice**: 各Expertが処理するトークンを選択します。完全なロードバランシングが保証されますが、一部のトークンがどのExpertにも選択されない可能性があります。

実際にはToken Choice + load balancing lossの組み合わせが最も多く使用されています。

クイズ

Q1: 「Mixture of Experts(MoE)アーキテクチャ完全分析」の主なトピックは何ですか?

Sparse

MoEの原理からMixtral、DeepSeek-V3のMoE実装、ルーティング戦略、ロードバランシングまで、MoEアーキテクチャを完全分析します。

Mixture of

Experts(MoE)は、モデルの全パラメータの一部のみを活性化して演算効率を高めるアーキテクチャです。すべての入力に対してモデル全体を使用するDenseモデルとは異なり、MoEは入力に応じて最適なエキスパート(Expert)のみを選択して処理します。

Expert Network 各Expertは独立したFFN(Feed-Forward Network)です: Router(Gating Network)

Routerは各トークンをどのExpertに送るかを決定します: MoEレイヤー全体の実装

Mixtral 8x7B(Mistral AI) 8つのExpert、Top-2ルーティング 合計46.7Bパラメータ、活性12.9B

Attentionレイヤーは共有、FFNのみExpertに分離 DeepSeek-V3 MoE

DeepSeek-V3はより精巧なMoE設計を採用しています: 1つの共有Expert + 256のRouted Expert(Top-8選択)

合計671Bパラメータ、活性37B Auxiliary-loss-freeロードバランシングを導入 モデル比較

Token Choice vs Expert Choice

현재 단락 (1/131)

**Mixture of Experts(MoE)**は、モデルの全パラメータの一部のみを活性化して演算効率を高めるアーキテクチャです。すべての入力に対してモデル全体を使用するDenseモデルとは異な...

작성 글자: 0원문 글자: 6,226작성 단락: 0/131