Mixture of Depths (MoD)

Mixture of Depths (MoD) est une technique de transformeur dans laquelle un routeur appris décide, bloc par bloc, quels jetons traiter et lesquels contourner le bloc via sa connexion résiduelle, de sorte que les jetons ne reçoivent pas tous la même quantité de calcul. David Raposo, Adam Santoro et leurs collègues de Google DeepMind l’ont présentée en avril 2024 dans Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models. Une capacité fixe k détermine combien de jetons chaque bloc routé traite ; le budget de calcul et les formes des tenseurs sont donc connus à l’avance, même si le choix des jetons dépend de l’entrée. Les grands modèles de langage (LLM) MoD ont égalé la qualité des modèles de référence à nombre égal de FLOPs d’entraînement et à durée d’exécution égale, tout en utilisant une fraction des FLOPs par passe avant et en accélérant l’échantillonnage de plus de 50 %. À budget d’entraînement fixe, le meilleur modèle MoD était plus grand que le meilleur modèle de référence et obtenait une perte plus faible.

Fonctionnement du routage MoD#

Avant un bloc routé, un routeur linéaire attribue un score scalaire à chaque jeton. Les k jetons ayant les scores les plus élevés passent par l’auto-attention et le MLP du bloc ; leurs sorties sont multipliées par leurs scores de routage, puis ajoutées au flux résiduel. Cette multiplication place le routeur sur le chemin du gradient, de sorte que la rétropropagation peut lui apprendre quels jetons en tirent parti. Les autres jetons passent inchangés et sont exclus de l’attention de ce bloc ; les jetons traités plus tard ne peuvent donc pas leur prêter attention dans ce bloc, et le coût d’attention diminue avec le carré de la capacité : traiter la moitié des jetons coûte un quart des produits requête-clé. Dans l’article, le meilleur réglage routait un bloc sur deux avec une capacité de 12,5 % de la séquence. Voici une version minimale d’un bloc routé, avec seulement le MLP, en PyTorch :

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Il affiche Block ran on 16 of 128 tokens; 112 skipped it unchanged : le MLP du bloc a effectué un huitième de son travail habituel, et le tenseur de sortie conserve la forme complète de la séquence pour le bloc suivant.

Routage pendant la génération#

La sélection top-k n’est pas causale : le classement d’un jeton parmi les k meilleurs de sa séquence dépend des scores des jetons qui le suivent, lesquels n’existent pas encore lorsqu’un modèle génère du texte un jeton à la fois. Les auteurs ont rendu le routage causal de deux manières : avec une perte auxiliaire d’entropie croisée binaire qui apprend au résultat du routeur à prédire l’appartenance au top-k, ce qui n’a modifié la perte de modélisation du langage que de 0,2 à 0,3 %, ou au moyen d’un petit MLP prédicteur distinct qui traite des entrées sans gradient et a rapidement atteint une précision de 97 à 99 %. Dans les deux cas, l’échantillonnage autorégressif a très peu perdu en qualité. Les routeurs entraînés présentent aussi un motif : certains jetons activent chaque bloc routé à travers les couches du modèle, tandis que d’autres contournent les blocs dès qu’ils le peuvent.

MoD et techniques apparentées#

• Mélange d’experts (MoE) : Route chaque jeton vers un ou plusieurs MLP experts d’une couche, ajoutant des paramètres tout en maintenant constant le calcul par jeton. L’article a également combiné les deux approches dans MoDE, où un jeton peut être routé vers un expert ou vers une opération sans effet ; les gains se sont alors cumulés.
• Sortie anticipée : Un jeton s’arrête après une certaine couche et saute toutes les suivantes. MoD peut contourner un bloc tout en utilisant les blocs ultérieurs.
• Fusion de jetons (ToMe) : Fusionne les jetons similaires dans un transformeur de vision afin que les couches suivantes en traitent moins, sans routeur appris.
• Élagage de modèle : Supprime définitivement des poids, de sorte que toutes les entrées utilisent le même réseau réduit.
• Routage de modèles : Dirige une requête entière vers l’un de plusieurs modèles ; il s’agit du même principe appliqué à l’échelle des modèles complets.

La plupart des réseaux ont une profondeur fixe dès leur conception. Ultralytics YOLO26 la définit pour chaque échelle de modèle : le multiplicateur de profondeur dans son fichier YAML de modèle vaut 0.50 pour les échelles n, s et m, et 1.00 pour l et x. Ainsi, une étape C3k2 indiquée avec 2 répétitions s’exécute en 1 bloc dans YOLO26n et en 2 dans YOLO26x, et chaque image les traverse tous.