Mixture of Depths (MoD)

Mixture of Depths (MoD) ist ein Transformer-Verfahren, bei dem ein erlernter Router für jeden Block entscheidet, welche Tokens verarbeitet werden und welche den Block über seine Restverbindung überspringen. So erhalten verschiedene Tokens unterschiedlich viel Rechenaufwand. David Raposo, Adam Santoro und ihre Kollegen bei Google DeepMind stellten das Verfahren im April 2024 in Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models vor. Eine feste Kapazität k legt fest, wie viele Tokens jeder geroutete Block verarbeitet. So stehen Rechenbudget und Tensorformen im Voraus fest, auch wenn die Auswahl der Tokens von der Eingabe abhängt. MoD-große Sprachmodelle (LLMs) erreichten bei gleichem Trainingsaufwand in FLOPs und gleicher Laufzeit eine vergleichbare Qualität wie die Basislinien. Gleichzeitig verwendeten sie pro Vorwärtsdurchlauf nur einen Bruchteil der FLOPs und generierten beim Sampling über 50 % schneller. Bei festem Trainingsbudget war das beste MoD-Modell größer als die beste Basislinie und erreichte einen niedrigeren Loss.

So funktioniert das Routing bei MoD#

Vor einem gerouteten Block weist ein linearer Router jedem Token einen Skalarwert zu. Die k Tokens mit den höchsten Werten durchlaufen die Selbstaufmerksamkeit und das MLP des Blocks. Ihre Ausgaben werden mit den Router-Werten multipliziert und zum Reststrom addiert. Durch die Multiplikation mit dem Wert liegt der Router im Gradientenpfad, sodass die Rückpropagation ihm beibringen kann, von welchen Tokens die Verarbeitung profitiert. Die übrigen Tokens werden unverändert weitergereicht und bei der Aufmerksamkeit dieses Blocks ausgelassen. Spätere Tokens können sie dort also nicht berücksichtigen, und die Aufmerksamkeitskosten sinken quadratisch mit der Kapazität: Die Verarbeitung der Hälfte der Tokens kostet ein Viertel der Query-Key-Produkte. Die beste Einstellung in der Arbeit leitete jeden zweiten Block weiter und setzte die Kapazität auf 12,5 % der Sequenz. Eine minimale Version eines einzelnen gerouteten Blocks, hier nur mit dem MLP, in 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")

Die Ausgabe lautet Block ran on 16 of 128 tokens; 112 skipped it unchanged: Das MLP des Blocks hat ein Achtel seines üblichen Aufwands benötigt, und der Ausgabetensor behält für den nächsten Block die volle Sequenzform.

Routing bei der Generierung#

Die Top-k-Auswahl ist nicht kausal: Ob ein Token zu den k besten seiner Sequenz gehört, hängt von den Werten der Tokens ab, die nach ihm kommen. Diese existieren noch nicht, wenn ein Modell Text Token für Token generiert. Die Autoren machten das Routing auf zwei Arten kausal: Entweder verwendeten sie einen zusätzlichen binären Kreuzentropie-Loss, der den Router darauf trainiert, anhand seiner Ausgabe die Zugehörigkeit zu den Top-k vorherzusagen. Dadurch verschob sich der Sprachmodellierungs-Loss nur um 0,2–0,3 %. Oder sie verwendeten ein kleines separates MLP, das Vorhersagen anhand von Eingaben mit gestoppter Gradientenweitergabe traf und schnell eine Genauigkeit von 97–99 % erreichte. In beiden Fällen ging beim autoregressiven Sampling nur wenig Qualität verloren. Trainierte Router zeigen außerdem ein Muster: Einige Tokens durchlaufen jeden gerouteten Block über die gesamte Modelltiefe hinweg, andere überspringen Blöcke, wann immer möglich.

MoD und verwandte Verfahren#

• Mixture of Experts (MoE): Leitet jedes Token innerhalb einer Schicht an ein oder wenige Experten-MLPs weiter. Dadurch kommen Parameter hinzu, während der Rechenaufwand pro Token gleich bleibt. Die Arbeit kombinierte beide Verfahren auch als MoDE: Dabei kann ein Token an einen Experten oder an eine wirkungslose Operation weitergeleitet werden, wodurch sich die Vorteile addierten.
• Vorzeitiger Abbruch: Ein Token wird nach einer bestimmten Schicht gestoppt und überspringt alle folgenden Schichten. MoD kann einen Block überspringen und spätere Blöcke trotzdem nutzen.
• Token-Merging (ToMe): Führt ähnliche Tokens in einem Vision Transformer zusammen, damit spätere Schichten weniger davon verarbeiten müssen. Ein erlernter Router kommt dabei nicht zum Einsatz.
• Modellpruning: Entfernt Gewichte dauerhaft, sodass jede Eingabe dasselbe kleinere Netz erhält.
• Modellrouting: Leitet eine gesamte Anfrage an eines von mehreren Modellen weiter – dasselbe Prinzip auf Ebene vollständiger Modelle.

Die meisten Netze haben nach ihrer Erstellung eine feste Tiefe. Ultralytics YOLO26 legt sie für jede Modellgröße fest: Der Tiefenmultiplikator in der Model-YAML-Datei beträgt für die Größen n, s und m 0,50 und für l und x 1,00. Daher wird eine C3k2-Stufe mit zwei Wiederholungen in YOLO26n einmal und in YOLO26x zweimal ausgeführt. Jedes Bild durchläuft alle Blöcke.