Mixture of Depths (MoD)

En las arquitecturas de aprendizaje profundo, la eficiencia computacional es primordial, especialmente al procesar secuencias largas o entradas de alta resolución. Un enfoque novedoso asigna dinámicamente los recursos de cómputo al permitir que la red decida qué partes de la entrada requieren un procesamiento completo y cuáles pueden omitir determinadas capas sin riesgo. Esta estrategia de enrutamiento dinámico reduce la complejidad computacional general sin sacrificar la capacidad predictiva ni la precisión del modelo.

Comprender el concepto#

La mezcla de profundidades (MoD) es una técnica arquitectónica aplicada principalmente a las arquitecturas Transformer, en las que el modelo aprende a omitir dinámicamente el cómputo de tokens específicos en distintas capas. Los Transformer tradicionales procesan cada token a través de todas las capas, independientemente de que sea una información crucial o contenido de relleno. En cambio, un modelo MoD utiliza un mecanismo de enrutamiento para evaluar los tokens y asignarles una puntuación. Solo los tokens con mayor puntuación —hasta un límite de capacidad predefinido— pasan por los bloques de cómputo intensivo, como los mecanismos de atención o las capas densas de propagación hacia delante. Los tokens restantes omiten el bloque mediante conexiones residuales, creando de forma efectiva una «mezcla de profundidades» en la que distintos tokens experimentan diferentes niveles de profundidad de procesamiento.

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

Diferencias con la mezcla de expertos (MoE)#

Es fácil confundir este concepto con una mezcla de expertos (MoE). Aunque ambas utilizan mecanismos de enrutamiento, resuelven problemas distintos:

• MoE dirige los tokens a distintas subredes (expertos) dentro de una capa. La profundidad computacional es la misma para todos los tokens, pero aumenta el número de parámetros del modelo.
• MoD dirige los tokens al bloque de cómputo o a una conexión de salto. El número de parámetros permanece estrictamente constante, pero la profundidad computacional disminuye para los tokens menos importantes, lo que mejora directamente la latencia de inferencia.

Aplicaciones en el mundo real#

La capacidad de asignar dinámicamente un presupuesto de cómputo hace que esta técnica sea muy valiosa en múltiples ámbitos de la visión artificial y el procesamiento del lenguaje natural.

1. Optimización del contexto en modelos de lenguaje: Los modelos de lenguaje grandes (LLMs) modernos de organizaciones como OpenAI y Anthropic procesan ventanas de contexto enormes. Mediante el enrutamiento dinámico por profundidad, estos modelos pueden omitir palabras de relleno estructurales o repetitivas y reservar el cómputo profundo para los pasos de razonamiento complejo y la extracción de información factual.

2. IA de visión de alta resolución: En sistemas avanzados de visión, como el modelo Ultralytics YOLO26, procesar imágenes grandes para la detección de objetos y la segmentación de imágenes requiere una cantidad ingente de memoria. El enrutamiento por profundidad permite que la red omita la extracción de características en fondos uniformes, como cielos despejados o paredes lisas, y concentre la potencia de cómputo en objetos complejos del primer plano. Esto es fundamental para implementar modelos en hardware de IA perimetral con recursos limitados, optimizado mediante bibliotecas de optimización de CUDA.

Ejemplo de implementación#

A continuación se muestra un fragmento conceptual de PyTorch que demuestra cómo un mecanismo de enrutamiento básico podría omitir el cómputo de una parte de los tokens de entrada, simulando un comportamiento de enrutamiento por 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)

Al aprovechar frameworks como el framework PyTorch o TensorFlow, los desarrolladores pueden 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 estos enrutadores con precisión, además de integrarse sin problemas con ecosistemas empresariales como Google Cloud AI.