Смесь глубин (MoD)

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

Понимание концепции

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

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

Дифференциация от смеси экспертов (MoE)

Это понятие легко спутать с смесью экспертов (MoE). Хотя оба используют механизмы маршрутизации, они решают разные проблемы:

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

Применение в реальном мире

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

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

Пример реализации

Ниже приведен концептуальный 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 , помогают командам управлять данными обучения, необходимыми для точного обучения этих маршрутизаторов, а также обеспечивают беспроблемную интеграцию с корпоративными экосистемами, такими как Google AI.