Mixture of Depths (MoD)

Dans les architectures de deep learning, l’efficacité computationnelle est primordiale, en particulier lors du traitement de longues séquences ou d’entrées haute résolution. Une nouvelle approche alloue dynamiquement les ressources de calcul en permettant au réseau de déterminer quelles parties de l’entrée nécessitent un traitement complet et lesquelles peuvent ignorer certaines couches sans risque. Cette stratégie de routage dynamique réduit la complexité computationnelle globale sans sacrifier la capacité prédictive ni la précision du modèle.

Comprendre le concept#

Le mélange de profondeurs (MoD) est une technique architecturale principalement appliquée aux architectures Transformer, dans lesquelles le modèle apprend à ignorer dynamiquement le calcul pour certains tokens à différentes couches. Les Transformer traditionnels traitent chaque token à travers chaque couche, qu’il s’agisse d’une information essentielle ou d’un contenu de remplissage. En revanche, un modèle MoD utilise un mécanisme de routage pour évaluer les tokens et leur attribuer un score. Seuls les tokens obtenant les scores les plus élevés — jusqu’à une limite de capacité prédéfinie — sont transmis aux blocs de calcul intensif, tels que les mécanismes d’attention ou les couches denses à propagation avant. Les tokens restants contournent le bloc via des connexions résiduelles, créant ainsi un « mélange de profondeurs » où différents tokens sont traités à des niveaux de profondeur variables.

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

Différencier le mélange de profondeurs du mélange d’experts (MoE)#

Il est facile de confondre ce concept avec un mélange d’experts (MoE). Bien que les deux utilisent des mécanismes de routage, ils répondent à des problèmes différents :

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

Applications concrètes#

La capacité à allouer dynamiquement le budget de calcul rend cette technique particulièrement utile dans de nombreux domaines de la vision par ordinateur et du traitement automatique du langage naturel.

1. Optimisation du contexte dans les modèles de langage : Les grands modèles de langage (LLMs) modernes d’organisations comme OpenAI et Anthropic traitent des fenêtres de contexte immenses. En utilisant un routage dynamique de la profondeur, ces modèles peuvent ignorer les mots de remplissage structurels ou répétitifs, réservant les calculs approfondis aux étapes de raisonnement complexe et à l’extraction de faits.

2. IA vision 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 quantité considérable de mémoire. Le routage de la profondeur permet au réseau d’ignorer l’extraction de caractéristiques sur les arrière-plans uniformes (comme les ciels dégagés ou les murs vierges), afin de concentrer la puissance de calcul sur les objets complexes au premier plan. Cette capacité est essentielle pour déployer des modèles sur du matériel d’IA en périphérie aux ressources limitées, optimisé par des bibliothèques d’optimisation CUDA.

Exemple d’implémentation#

Voici un extrait conceptuel en PyTorch illustrant comment un mécanisme de routage élémentaire pourrait ignorer le calcul pour une partie des tokens d’entrée, simulant un comportement de routage de la 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 s’appuyant sur des frameworks comme le framework PyTorch ou TensorFlow, les développeurs peuvent intégrer ces blocs personnalisés d’optimisation de modèles. En outre, des outils comme l’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 fluide aux écosystèmes d’entreprise comme Google Cloud AI.