Mixture of Depths (MoD)

En las arquitecturas de aprendizaje profundo, la eficiencia computacional es fundamental, especialmente al procesar secuencias largas o entradas de alta resolución. Un enfoque novedoso asigna dinámicamente los recursos de computación permitiendo que la red decida qué partes de la entrada requieren un procesamiento completo y cuáles pueden saltarse ciertas capas con seguridad. Esta estrategia de enrutamiento dinámico reduce la complejidad computacional general sin sacrificar la potencia predictiva ni la precisión del modelo.

Link to this sectionEntender el concepto#

Mixture of Depths (MoD) es una técnica arquitectónica aplicada principalmente a arquitecturas Transformer donde el modelo aprende a omitir dinámicamente el cálculo de tokens específicos en varias capas. Los Transformers tradicionales procesan cada token a través de todas las capas, ya sea una pieza de información crucial o contenido de relleno. Por el contrario, un modelo MoD utiliza un mecanismo de enrutamiento para evaluar los tokens y asignarles una puntuación. Solo los tokens con mejor puntuación (hasta un límite de capacidad predefinido) pasan a través de los bloques de cálculo pesado, como los mecanismos de atención o las capas densas de feed-forward. Los tokens restantes omiten el bloque a través de conexiones residuales, creando efectivamente una "mezcla de profundidades" donde diferentes tokens experimentan niveles variables de profundidad de procesamiento.

Este método, popularizado por investigaciones recientes de DeepMind y documentado exhaustivamente en el repositorio arXiv, reduce drásticamente el número total de operaciones de coma flotante (FLOPs) necesarias tanto durante el entrenamiento como en la inferencia.

Link to this sectionDiferenciación de Mixture of Experts (MoE)#

Es fácil confundir este concepto con una Mixture of Experts (MoE). Aunque ambos utilizan mecanismos de enrutamiento, resuelven problemas diferentes:

• MoE enruta los tokens a diferentes subredes (expertos) dentro de una capa. La profundidad computacional sigue siendo la misma para todos los tokens, pero aumenta el número de parámetros del modelo.
• MoD enruta los tokens al bloque de cálculo o a una conexión residual. El número de parámetros se mantiene estrictamente constante, pero la profundidad computacional disminuye para los tokens menos importantes, lo que mejora directamente la latencia de inferencia.

Link to this sectionAplicaciones en el mundo real#

La capacidad de presupuestar el cómputo de forma dinámica hace que esta técnica sea muy valiosa en múltiples dominios de visión artificial y procesamiento del lenguaje natural.

1. Optimización del contexto en modelos de lenguaje: Los Large Language Models (LLMs) modernos de organizaciones como OpenAI y Anthropic procesan enormes ventanas de contexto. Al emplear el enrutamiento de profundidad dinámico, estos modelos pueden omitir palabras de relleno estructurales o repetitivas, reservando el cómputo profundo para pasos de razonamiento complejos y extracción de hechos.

2. Visión artificial de alta resolución: En sistemas de visión avanzados como el modelo Ultralytics YOLO26, procesar imágenes grandes para detección de objetos y segmentación de imágenes requiere una memoria inmensa. El enrutamiento de profundidad permite a la red omitir la extracción de características en fondos uniformes (como cielos vacíos o paredes en blanco), centrando la potencia computacional en objetos intrincados en primer plano. Esto es crucial para desplegar modelos en hardware de Edge AI con recursos limitados, optimizado por bibliotecas de optimización CUDA.

Link to this sectionEjemplo de implementación#

A continuación, un fragmento conceptual de PyTorch que demuestra cómo un mecanismo básico de enrutamiento podría omitir el cálculo para una parte de los tokens de entrada, simulando un comportamiento de enrutamiento de profundidad.

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)

Aprovechando marcos como PyTorch o TensorFlow, puedes integrar estos bloques personalizados de optimización de modelos. Además, herramientas como la Plataforma Ultralytics ayudan a los equipos a gestionar los datos de entrenamiento necesarios para entrenar con precisión estos enrutadores, además de integrarse perfectamente con ecosistemas empresariales como Google Cloud AI.