Mixture of Depths (MoD)

Nelle architetture di deep learning, l'efficienza computazionale è fondamentale, specialmente 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 saltare in sicurezza determinati strati. Questa strategia di routing dinamico riduce la complessità computazionale complessiva senza sacrificare la potenza predittiva o l'accuratezza del modello.

Link to this sectionComprendere il concetto#

La Mixture of Depths (MoD) è una tecnica architettonica applicata principalmente alle architetture Transformer in cui il modello impara a saltare dinamicamente il calcolo per token specifici a vari strati. I transformer tradizionali elaborano ogni token attraverso ogni strato, che si tratti di un'informazione cruciale o di contenuto di riempimento. 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—vengono fatti passare attraverso i blocchi di calcolo pesante, come i meccanismi di attenzione o i layer feed-forward densi. I token rimanenti bypassano il blocco tramite connessioni residuali, creando effettivamente una "miscela di profondità" in cui diversi token sperimentano livelli variabili di profondità di elaborazione.

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

Link to this sectionDifferenziare dalla Mixture of Experts (MoE)#

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

• MoE instrada i token verso diverse sotto-reti (esperti) all'interno di uno strato. 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 o verso una skip connection. Il numero di parametri rimane rigorosamente costante, ma la profondità computazionale diminuisce per i token meno importanti, migliorando direttamente la latenza di inferenza.

Link to this sectionApplicazioni nel mondo reale#

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

1. Ottimizzazione del contesto nei modelli linguistici: I moderni Large Language Models (LLMs) di organizzazioni come OpenAI e Anthropic elaborano enormi finestre di contesto. Impiegando il routing dinamico della profondità, questi modelli possono saltare parole di riempimento strutturali o ripetitive, riservando il calcolo profondo a passaggi di ragionamento complessi ed estrazione di fatti.

2. AI di visione ad alta risoluzione: In sistemi di visione avanzati come il modello Ultralytics YOLO26, l'elaborazione di immagini di grandi dimensioni per object detection e image segmentation richiede un'enorme 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 computazionale su oggetti in primo piano intricati. Questo è cruciale per implementare modelli su hardware edge AI con risorse limitate, ottimizzato dalle librerie di CUDA optimization.

Link to this sectionEsempio di Implementazione#

Di seguito è riportato uno snippet concettuale di PyTorch che mostra come un meccanismo di routing di base potrebbe saltare il calcolo per una porzione di token di input, simulando un comportamento di depth-routing.

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 PyTorch framework o TensorFlow, gli sviluppatori possono integrare questi blocchi personalizzati di model optimization. Inoltre, strumenti come la Ultralytics Platform aiutano i team a gestire i training data necessari per addestrare accuratamente questi router, integrandoli perfettamente con ecosistemi aziendali come Google Cloud AI.