Mixture of Depths (MoD)

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

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

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

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

Отличие от смеси экспертов (MoE)#

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

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

Практические применения#

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

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

2. Компьютерное зрение с высоким разрешением: в передовых системах компьютерного зрения, таких как модель Ultralytics YOLO26, обработка больших изображений для обнаружения объектов и сегментации изображений требует огромного объёма памяти. Маршрутизация по глубине позволяет сети пропускать извлечение признаков на однородном фоне, например на участках с чистым небом или однотонными стенами, и сосредоточить вычислительные ресурсы на сложных объектах переднего плана. Это крайне важно для развертывания моделей на аппаратном обеспечении граничного ИИ с ограниченными ресурсами, оптимизированном с помощью библиотек оптимизации 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 Platform, помогают командам управлять данными для обучения, необходимыми для точного обучения этих маршрутизаторов, и обеспечивают бесшовную интеграцию с корпоративными экосистемами, такими как Google Cloud AI.