Mixture of Depths (MoD)

В архитектурах глубокого обучения вычислительная эффективность имеет первостепенное значение, особенно при обработке длинных последовательностей или входных данных высокого разрешения. Новый подход динамически распределяет вычислительные ресурсы, позволяя сети решать, какие части входных данных требуют полной обработки, а какие могут безопасно миновать определенные слои. Эта стратегия динамической маршрутизации снижает общую вычислительную сложность, не жертвуя при этом предсказательной мощностью или точностью модели.

Link to this sectionПонимание концепции#

Mixture of Depths (MoD) — это архитектурный метод, который в основном применяется к архитектурам Transformer, где модель учится динамически пропускать вычисления для конкретных токенов на различных слоях. Традиционные трансформеры пропускают каждый токен через каждый слой, будь то важная информация или просто «вода». В отличие от них, модель MoD использует механизм маршрутизации для оценки токенов и присвоения им баллов. Только токены с наивысшим баллом (до заданного предела емкости) проходят через тяжелые вычислительные блоки, такие как механизмы внимания или плотные полносвязные слои. Остальные токены обходят блок через остаточные связи, эффективно создавая «смесь глубин» (mixture of depths), где разные токены проходят через уровни обработки разной глубины.

Этот метод, популяризированный недавними исследованиями DeepMind и подробно описанный в репозитории arXiv, радикально сокращает общее количество операций с плавающей запятой (FLOPs), необходимых как во время обучения, так и при инференсе.

Link to this sectionОтличие от Mixture of Experts (MoE)#

Легко спутать эту концепцию с Mixture of Experts (MoE). Хотя оба подхода используют механизмы маршрутизации, они решают разные задачи:

• MoE направляет токены к различным подсетям (экспертам) внутри слоя. Вычислительная глубина остается одинаковой для всех токенов, но количество параметров модели увеличивается.
• MoD направляет токены либо в вычислительный блок, либо через skip connection. Количество параметров остается строго постоянным, но вычислительная глубина уменьшается для менее важных токенов, что напрямую улучшает задержку инференса.

Link to this sectionРеальные приложения#

Способность динамически распределять вычислительный бюджет делает этот метод весьма ценным в различных областях компьютерного зрения и обработки естественного языка.

1. Оптимизация контекста в языковых моделях: Современные большие языковые модели (LLM) от таких организаций, как OpenAI и Anthropic, обрабатывают огромные контекстные окна. Используя динамическую маршрутизацию глубины, эти модели могут пропускать структурные или повторяющиеся вспомогательные слова, резервируя глубокие вычисления для сложных этапов рассуждения и извлечения фактов.

2. Компьютерное зрение высокого разрешения: В продвинутых системах технического зрения, таких как модель Ultralytics YOLO26, обработка больших изображений для обнаружения объектов и сегментации изображений требует огромного объема памяти. Маршрутизация глубины позволяет сети пропускать извлечение признаков на однородном фоне (например, пустом небе или чистых стенах), концентрируя вычислительную мощность на сложных объектах переднего плана. Это крайне важно для развертывания моделей на оборудовании с ограниченными ресурсами Edge AI, оптимизированном с помощью библиотек CUDA optimization.

Link to this sectionПример реализации#

Ниже приведен концептуальный фрагмент PyTorch, демонстрирующий, как базовый механизм маршрутизации может пропускать вычисления для части входных токенов, имитируя поведение маршрутизации глубины.

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)

Используя такие фреймворки, как PyTorch или TensorFlow, разработчики могут интегрировать эти кастомные блоки оптимизации моделей. Кроме того, такие инструменты, как Ultralytics Platform, помогают командам управлять обучающими данными, необходимыми для точного обучения этих маршрутизаторов, а также легко интегрироваться с корпоративными экосистемами, такими как Google Cloud AI.