Top-K Routing (MoE)

Définition

Le Top-K Routing est le mecanisme de selection d'experts dans les architectures Mixture of Experts (MoE) des LLMs. Un reseau de routage (router) calcule un score d'affinite entre chaque token entrant et chaque expert disponible, puis selectionne les K experts ayant les scores les plus eleves pour traiter ce token. Typiquement K=1 (Mixtral de base) ou K=2 (Mixtral 8x7B, DeepSeek-MoE). L'architecture du router est simple : une couche lineaire W_r de dimension (d_model, n_experts) produit des logits pour chaque expert, normalises par softmax pour obtenir des probabilites de selection. Les Top-K experts (scores les plus eleves) sont selectionnes, et la sortie finale est la moyenne ponderee de leurs outputs par leurs scores de routage. Le Load Balancing est le defi principal du Top-K Routing : sans contrainte, le router tend a toujours selectionner les memes experts (quelques experts 'populaires'), laissant les autres inutilises. Une loss auxiliaire de load balancing (Auxiliary Loss) penalise les distributions de routage inegales, forcant une distribution plus uniforme des tokens entre experts. DeepSeek-MoE a introduit des innovations importantes au Top-K Routing : (1) experts ultra-fins (64 experts de petite taille plutot que 8 experts de grande taille), (2) experts partages (quelques experts activés pour tous les tokens + experts specialises routes par Top-K), (3) balance loss plus fine par niveau de couche. Ces innovations ont permis de reduire le cout d'activation tout en ameliorant les performances. Les MoE avec Top-K Routing obtiennent souvent des performances superieures a un Dense Model de meme cout de calcul, car differents tokens (code, maths, factuel) peuvent etre routes vers des experts specialises. La taille totale des poids est grande (Mixtral 8x7B = 47B poids), mais seulement 2 experts sont actives par token, donnant un cout d'inference equivalent a un modele ~13B.

Implementation Router MoE

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

class TopKRouter(nn.Module):
    def __init__(self, d_model, n_experts, top_k=2):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k
        self.router = nn.Linear(d_model, n_experts, bias=False)

    def forward(self, x):
        # x: [batch, seq_len, d_model]
        logits = self.router(x)  # [batch, seq_len, n_experts]
        scores = F.softmax(logits, dim=-1)

        # Selectionner Top-K experts
        top_k_scores, top_k_indices = torch.topk(scores, self.top_k, dim=-1)

        # Normaliser les poids des Top-K experts (somme a 1)
        top_k_scores = top_k_scores / top_k_scores.sum(dim=-1, keepdim=True)
        return top_k_scores, top_k_indices

Parametres typiques des MoE populaires

Modele	Nb experts	Top-K	Params actifs
Mixtral 8x7B	8	2	~13B
Mixtral 8x22B	8	2	~39B
DeepSeek-MoE 16B	64	6	~3B
DeepSeek V3 (671B)	256	8	~37B
GPT-4 (estime)	~16-32	2	~100-200B

Articles liés

Expert en cybersécurité offensive et intelligence artificielle. Pentest, audit et développement IA sur-mesure.

Services

Ressources

Projets & Outils

Contact

Nous contacter

Mentions légales Confidentialité Politique éditoriale Plan du site ⭐ Source préférée

Un projet cybersécurité ?

Expert dispo · Réponse 24h

Devis