Mixture of Depths (MoD)

Dans les architectures d'apprentissage profond, l'efficacité computationnelle est primordiale, surtout lors du traitement de longues séquences ou d'entrées haute résolution. Une approche innovante alloue dynamiquement les ressources de calcul en permettant au réseau de décider quelles parties de l'entrée nécessitent un traitement complet et lesquelles peuvent ignorer certaines couches en toute sécurité. Cette stratégie de routage dynamique réduit la complexité computationnelle globale sans sacrifier la puissance prédictive ou la précision du modèle.

Link to this sectionComprendre le concept#

Le Mixture of Depths (MoD) est une technique architecturale principalement appliquée aux architectures Transformer où le modèle apprend à ignorer dynamiquement le calcul pour des jetons spécifiques à différentes couches. Les transformers traditionnels traitent chaque jeton à travers chaque couche, qu'il s'agisse d'une information cruciale ou d'un contenu de remplissage. En revanche, un modèle MoD utilise un mécanisme de routage pour évaluer les jetons et leur attribue un score. Seuls les jetons ayant les scores les plus élevés—jusqu'à une limite de capacité prédéfinie—sont transmis aux blocs de calcul lourds, tels que les mécanismes d'attention ou les couches denses feed-forward. Les jetons restants contournent le bloc via des connexions résiduelles, créant ainsi un "mélange de profondeurs" où différents jetons subissent des niveaux de profondeur de traitement variés.

Cette méthode, popularisée par des recherches DeepMind récentes et documentée en détail dans le dépôt arXiv, réduit considérablement le nombre total d'opérations en virgule flottante (FLOPs) nécessaires lors de l'entraînement et de l'inférence.

Link to this sectionDifférencier du Mixture of Experts (MoE)#

Il est facile de confondre ce concept avec un Mixture of Experts (MoE). Bien que les deux utilisent des mécanismes de routage, ils résolvent des problèmes différents :

• MoE achemine les jetons vers différents sous-réseaux (experts) au sein d'une couche. La profondeur computationnelle reste la même pour tous les jetons, mais le nombre de paramètres du modèle augmente.
• MoD achemine les jetons soit vers le bloc de calcul, soit vers une connexion de saut. Le nombre de paramètres reste strictement constant, mais la profondeur computationnelle diminue pour les jetons moins importants, améliorant directement la latence d'inférence.

Link to this sectionApplications concrètes#

La capacité à budgétiser dynamiquement le calcul rend cette technique très précieuse dans de multiples domaines de la computer vision et du traitement du langage naturel.

1. Optimisation du contexte dans les modèles de langage : Les Large Language Models (LLMs) modernes issus d'organisations comme OpenAI et Anthropic traitent des fenêtres de contexte massives. En employant un routage de profondeur dynamique, ces modèles peuvent ignorer les mots de remplissage structurels ou répétitifs, réservant le calcul intensif aux étapes de raisonnement complexes et à l'extraction factuelle.

2. Vision IA haute résolution : Dans les systèmes de vision avancés comme le modèle Ultralytics YOLO26, le traitement de grandes images pour la détection d'objets et la segmentation d'images nécessite une mémoire immense. Le routage de profondeur permet au réseau de contourner l'extraction de caractéristiques sur les arrière-plans uniformes (comme les ciels vides ou les murs nus), en concentrant la puissance de calcul sur les objets de premier plan complexes. Ceci est crucial pour déployer des modèles sur du matériel edge AI à ressources limitées optimisé par des bibliothèques de CUDA optimization.

Link to this sectionExemple d'implémentation#

Voici un extrait PyTorch conceptuel démontrant comment un mécanisme de routage de base pourrait ignorer le calcul pour une partie des jetons d'entrée, simulant un comportement de routage de profondeur.

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)

En exploitant des frameworks comme le PyTorch framework ou TensorFlow, tu peux intégrer ces blocs personnalisés d'optimisation de modèle. De plus, des outils comme la Ultralytics Platform aident les équipes à gérer les données d'entraînement nécessaires pour entraîner précisément ces routeurs, tout en s'intégrant de manière transparente aux écosystèmes d'entreprise comme Google Cloud AI.