Mixture of Depths (MoD)

In Deep-Learning-Architekturen ist Recheneffizienz von größter Bedeutung, insbesondere bei der Verarbeitung langer Sequenzen oder hochauflösender Eingaben. Ein neuartiger Ansatz weist Rechenressourcen dynamisch zu, indem das Netzwerk entscheidet, welche Teile der Eingabe vollständig verarbeitet werden müssen und welche bestimmte Schichten sicher umgehen können. Diese dynamische Routing-Strategie reduziert die gesamte Rechenkomplexität, ohne die Vorhersageleistung oder Genauigkeit des Modells zu beeinträchtigen.

Das Konzept verstehen#

Mixture of Depths (MoD) ist eine Architekturtechnik, die hauptsächlich auf Transformer-Architekturen angewendet wird. Dabei lernt das Modell, die Berechnung für bestimmte Token in verschiedenen Schichten dynamisch zu überspringen. Herkömmliche Transformer verarbeiten jedes Token in jeder Schicht, unabhängig davon, ob es sich um eine wichtige Information oder um Füllinhalt handelt. Im Gegensatz dazu verwendet ein MoD-Modell einen Router-Mechanismus, um Token zu bewerten und ihnen eine Punktzahl zuzuweisen. Nur die Token mit den höchsten Punktzahlen – bis zu einer vordefinierten Kapazitätsgrenze – werden durch die aufwendigen Berechnungsblöcke geleitet, etwa Aufmerksamkeitsmechanismen oder dichte Feedforward-Schichten. Die übrigen Token umgehen den Block über Residualverbindungen. Dadurch entsteht effektiv eine „Mischung von Tiefen“, bei der verschiedene Token unterschiedlich tief verarbeitet werden.

Diese Methode, die durch aktuelle Forschungen von DeepMind bekannt wurde und ausführlich im arXiv-Archiv dokumentiert ist, reduziert die Gesamtzahl der während des Trainings und der Inferenz erforderlichen Gleitkommaoperationen (FLOPs) drastisch.

Abgrenzung zu Mixture of Experts (MoE)#

Dieses Konzept lässt sich leicht mit einem Mixture of Experts (MoE) verwechseln. Obwohl beide Routing-Mechanismen verwenden, lösen sie unterschiedliche Probleme:

• MoE leitet Token innerhalb einer Schicht an verschiedene Teilnetzwerke (Experten) weiter. Die Berechnungstiefe bleibt für alle Token gleich, aber die Anzahl der Modellparameter steigt.
• MoD leitet Token entweder an den Berechnungsblock oder an eine Skip-Verbindung weiter. Die Anzahl der Parameter bleibt strikt konstant, aber die Berechnungstiefe nimmt für weniger wichtige Token ab, wodurch sich die Inferenzlatenz direkt verbessert.

Anwendungen in der Praxis#

Die Möglichkeit, das Rechenbudget dynamisch festzulegen, macht diese Technik in zahlreichen Bereichen der Computer Vision und der Verarbeitung natürlicher Sprache besonders wertvoll.

1. Kontextoptimierung in Sprachmodellen: Moderne große Sprachmodelle (LLMs) von Organisationen wie OpenAI und Anthropic verarbeiten riesige Kontextfenster. Durch den Einsatz dynamischen Routings über die Verarbeitungstiefe können diese Modelle strukturelle oder sich wiederholende Füllwörter überspringen und die aufwendige Berechnung für komplexe Schlussfolgerungen und die Extraktion von Fakten reservieren.

2. KI für hochauflösende Bildverarbeitung: In fortschrittlichen Bildverarbeitungssystemen wie dem Modell Ultralytics YOLO26 erfordert die Verarbeitung großer Bilder für die Objekterkennung und die Bildsegmentierung enorme Speicherkapazitäten. Das Routing über die Verarbeitungstiefe ermöglicht es dem Netzwerk, die Merkmalsextraktion in gleichförmigen Hintergründen zu umgehen, etwa bei leerem Himmel oder einfarbigen Wänden, und die Rechenleistung auf komplexe Objekte im Vordergrund zu konzentrieren. Dies ist entscheidend für die Bereitstellung von Modellen auf ressourcenbeschränkter Edge-KI-Hardware, die durch Bibliotheken zur CUDA-Optimierung optimiert wird.

Implementierungsbeispiel#

Das folgende konzeptionelle PyTorch-Codebeispiel zeigt, wie ein einfacher Routing-Mechanismus die Berechnung für einen Teil der Eingabetoken überspringen und so ein Verhalten zur Steuerung der Verarbeitungstiefe simulieren kann.

import torch
import torch.nn as nn


class MixtureOfDepthsBlock(nn.Module):
    def __init__(self, d_model, capacity_factor=0.5):
        super().__init__()
        self.capacity_factor = capacity_factor
        self.router = nn.Linear(d_model, 1)
        self.heavy_compute = nn.Sequential(nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model))

    def forward(self, x):
        # x shape: (batch_size, seq_len, d_model)
        seq_len = x.size(1)
        capacity = int(seq_len * self.capacity_factor)

        # 1. Compute routing scores
        scores = self.router(x).squeeze(-1)  # Shape: (batch_size, seq_len)

        # 2. Identify top-k tokens to process
        topk_indices = torch.topk(scores, capacity, dim=1).indices

        # 3. Create an output tensor mirroring the input (residual baseline)
        output = x.clone()

        # 4. Apply heavy computation only to the selected tokens
        for b in range(x.size(0)):
            selected_tokens = x[b, topk_indices[b]]
            processed_tokens = self.heavy_compute(selected_tokens)
            output[b, topk_indices[b]] += processed_tokens

        return output


# Example usage
dummy_input = torch.randn(2, 64, 128)  # Batch=2, Seq=64, Dim=128
mod_block = MixtureOfDepthsBlock(d_model=128, capacity_factor=0.5)
output = mod_block(dummy_input)
print(f"Output shape: {output.shape}")  # Expect (2, 64, 128)

Durch die Nutzung von Frameworks wie dem PyTorch-Framework oder TensorFlow können Entwickler diese benutzerdefinierten Blöcke zur Modelloptimierung integrieren. Darüber hinaus unterstützen Tools wie die Ultralytics Platform Teams bei der Verwaltung der Trainingsdaten, die zum präzisen Trainieren dieser Router erforderlich sind, und lassen sich nahtlos in Unternehmensökosysteme wie Google Cloud AI integrieren.