Mixture of Depths (MoD)

Nelle architetture di deep learning, l'efficienza computazionale è fondamentale, soprattutto quando si elaborano sequenze lunghe o input ad alta risoluzione. Un approccio innovativo alloca dinamicamente le risorse di calcolo consentendo alla rete di decidere quali parti dell'input richiedono un'elaborazione completa e quali possono tranquillamente bypassare determinati livelli. Questa strategia di routing dinamico riduce la complessità computazionale complessiva senza sacrificare la capacità predittiva o l'accuratezza del modello.

Comprendere il concetto#

La Miscela di profondità (MoD) è una tecnica architetturale applicata principalmente alle architetture Transformer, in cui il modello impara a saltare dinamicamente il calcolo per token specifici a diversi livelli. I Transformer tradizionali elaborano ogni token attraverso ogni livello, indipendentemente dal fatto che contenga un'informazione fondamentale o del contenuto irrilevante. Al contrario, un modello MoD utilizza un meccanismo di routing per valutare i token e assegnare loro un punteggio. Solo i token con il punteggio più alto, fino a un limite di capacità predefinito, passano attraverso i blocchi di calcolo più onerosi, come i meccanismi di attenzione o i livelli feed-forward densi. I token rimanenti bypassano il blocco tramite connessioni residuali, creando di fatto una "miscela di profondità" in cui token diversi vengono elaborati a livelli di profondità differenti.

Questo metodo, reso popolare da recenti ricerche di DeepMind e ampiamente documentato nel repository arXiv, riduce drasticamente il numero totale di operazioni in virgola mobile (FLOPs) necessarie sia durante l'addestramento sia durante l'inferenza.

Differenze rispetto alla Miscela di esperti (MoE)#

È facile confondere questo concetto con una Miscela di esperti (MoE). Sebbene entrambi utilizzino meccanismi di routing, risolvono problemi diversi:

• MoE instrada i token verso diverse sotto-reti (esperti) all'interno di un livello. La profondità computazionale rimane la stessa per tutti i token, ma il numero di parametri del modello aumenta.
• MoD instrada i token verso il blocco di calcolo oppure verso una connessione skip. Il numero di parametri rimane rigorosamente costante, ma la profondità computazionale diminuisce per i token meno importanti, migliorando direttamente la latenza dell'inferenza.

Applicazioni nel mondo reale#

La capacità di assegnare dinamicamente un budget di calcolo rende questa tecnica estremamente preziosa in diversi ambiti della computer vision e dell'elaborazione del linguaggio naturale.

1. Ottimizzazione del contesto nei modelli linguistici: i moderni modelli linguistici di grandi dimensioni (LLMs) di organizzazioni come OpenAI e Anthropic elaborano finestre di contesto enormi. Utilizzando il routing dinamico della profondità, questi modelli possono saltare parole strutturali o ripetitive e irrilevanti, riservando il calcolo approfondito ai passaggi di ragionamento complessi e all'estrazione di informazioni fattuali.

2. IA per la visione ad alta risoluzione: nei sistemi avanzati di visione come il modello Ultralytics YOLO26, l'elaborazione di immagini di grandi dimensioni per il rilevamento degli oggetti e la segmentazione delle immagini richiede una quantità enorme di memoria. Il routing della profondità consente alla rete di bypassare l'estrazione delle caratteristiche su sfondi uniformi, come cieli vuoti o pareti bianche, concentrando la potenza di calcolo sugli intricati oggetti in primo piano. Questo è fondamentale per distribuire i modelli su hardware edge AI con risorse limitate, ottimizzato mediante librerie di ottimizzazione CUDA.

Esempio di implementazione#

Di seguito è riportato uno snippet concettuale in PyTorch che mostra come un meccanismo di routing di base possa saltare il calcolo per una parte dei token di input, simulando un comportamento di routing della profondità.

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)

Sfruttando framework come il framework PyTorch o TensorFlow, gli sviluppatori possono integrare questi blocchi personalizzati di ottimizzazione del modello. Inoltre, strumenti come la piattaforma Ultralytics aiutano i team a gestire i dati di addestramento necessari per addestrare accuratamente questi router, integrandosi al contempo perfettamente con ecosistemi enterprise come Google Cloud AI.