MoCo

Материал из MachineLearning.

Версия от 17:45, 4 июля 2026; Kirill Bazhutov (Обсуждение | вклад)
(разн.) ← Предыдущая | Текущая версия (разн.) | Следующая → (разн.)
Перейти к: навигация, поиск
Статья написана с использованием LLM Gemini и проверена участником Kirill Bazhutov 21:45, 4 июля 2026 (MSD)


Momentum Contrast (MoCo,Momentum Contrast) — метод самоконтролируемого обучения представлений (-supervised learning, SSL) для задач компьютерного зрения, предложенный исследователями Facebook AI Research (FAIR) в 2019 году. Метод базируется на парадигме контрастивного обучения, в рамках которой задача обучения информативных признаков (representation learning) формулируется как задача поиска по словарю (dictionary look-up).

Содержание

Формальная математическая постановка

В контрастивном обучении модель учится сближать в векторном пространстве представления разных аугментированных видов одного и того же изображения и отдалять представления других изображений. Для каждого изображения генерируются два случайных аугментированных вида. Одно из них кодируется в вектор запроса (query) q. Второе кодируется в вектор положительного ключа (positive key) k^+. Также имеется набор из K векторов отрицательных ключей (negative keys) k^-, представляющих другие изображения.

Цель модели — научиться находить правильный ключ k^+ среди множества всех доступных ключей \{k^+, k^-_1, \dots, k^-_K\} для заданного запроса q. Эта задача решается путём минимизации InfoNCE — контрастивной функции потерь, введённой в методе Contrastive Predictive Coding. Она математически эквивалентна функции кросс-энтропии для задачи многоклассовой классификации на K+1 классов, где логиты выражены через скалярные произведения:

L_q = - \log \frac{\exp(q \cdot k^+ / \tau)}{\exp(q \cdot k^+ / \tau) + \sum_{i=1}^K \exp(q \cdot k_i^- / \tau)}

где:

  • q \cdot k — скалярное произведение, измеряющее косинусное сходство (векторы q и k предварительно L_2-нормализованы).
  • \tau (temperature) — гиперпараметр температуры, контролирующий масштаб распределения Softmax и жёсткость штрафа за трудноразличимые отрицательные примеры (hard negatives).

Проблематика размера словаря

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

  1. Сквозное обратное распространение (End-to-end): Векторы q и k вычисляются энкодерами с градиентами в рамках одного прохода. В этом случае размер словаря K жёстко ограничен размером мини-пакета (mini-batch), который упирается в объём доступной видеопамяти (GPU).
  2. Банк памяти (Memory bank): Словарь хранится в памяти и состоит из представлений всех объектов датасета, вычисленных на предыдущих эпохах. Размер словаря огромен, но ключи могут становиться устаревшими и менее согласованными, поскольку были вычислены предыдущими версиями энкодера на разных этапах оптимизации сети.

Архитектурные решения MoCo

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

Словарь как очередь (Dictionary as a queue)

Словарь отрицательных примеров реализуется как структура данных очередь (FIFO) фиксированного размера K. На каждой итерации текущий мини-пакет новых ключей добавляется в очередь, а самый старый мини-пакет удаляется. Это полностью отвязывает размер словаря от размера мини-пакета (например, размер очереди может быть K=65536 при размере батча 256), не требуя вычисления градиентов для всех ключей в словаре.

Импульсное обновление (Momentum update)

Использование очереди приводит к тому, что ключи генерируются энкодером, веса которого меняются с каждой итерацией оптимизатора. Чтобы обеспечить плавность (консистентность) представлений ключей в очереди, MoCo использует два раздельных энкодера:

  • f_q с весами \theta_q (query encoder) — обновляется стандартным алгоритмом обратного распространения ошибки.
  • f_k с весами \theta_k (key encoder) — не участвует в обратном распространении и обновляется как скользящее экспоненциальное среднее (Exponential Moving Average, EMA) от весов \theta_q:
\theta_k \leftarrow m \theta_k + (1 - m) \theta_q

где m \in [0, 1) — коэффициент импульса (momentum). На практике m берется близким к единице (например, m=0.999), что делает изменения key encoder крайне медленными и обеспечивает согласованность элементов в словаре.

Shuffle Batch Normalization

В оригинальной реализации MoCo также использовалась процедура перемешивания при батч-нормализации (shuffle batch normalization). Поскольку query и key энкодеры при стандартном использовании BatchNorm могут «подглядывать» статистики одного мини-пакета, на нескольких GPU элементы мини-пакета для key энкодера перемешивались перед вычислением BatchNorm и затем возвращались в исходный порядок. Это предотвращало утечку информации между положительными парами через внутрипакетные статистики нормализации.

Псевдокод алгоритма

Концептуальная реализация логики MoCo в стиле PyTorch:

# f_q, f_k: энкодеры для query и key (например, ResNet-50)
# queue: словарь (очередь) размером K x D
# m: коэффициент импульса
# tau: температура

f_k.params = f_q.params # инициализация весов

for x in loader:
    # генерация двух случайных аугментаций одного изображения
    x_q, x_k = aug(x), aug(x)

    q = f_q.forward(x_q) # запросы: N x D, с градиентами
    q = normalize(q, dim=1)

    with torch.no_grad(): # ключи вычисляются без градиентов
        # импульсное обновление весов f_k
        for param_q, param_k in zip(f_q.params, f_k.params):
            param_k.data = m * param_k.data + (1 - m) * param_q.data
        
        k = f_k.forward(x_k) # положительные ключи: N x D
        k = normalize(k, dim=1)
        k = k.detach() # отвязка от вычислительного графа

    # вычисление логитов InfoNCE
    # положительные логиты: Nx1
    l_pos = bmm(q.view(N,1,D), k.view(N,D,1))
    # отрицательные логиты: NxK
    l_neg = mm(q.view(N,D), queue.view(D,K))
    
    # полные логиты: Nx(1+K)
    logits = cat([l_pos, l_neg], dim=1)
    
    # контрастивный лосс (эквивалентен кросс-энтропии, цель - индекс 0)
    labels = zeros(N)
    loss = CrossEntropyLoss(logits / tau, labels)
    
    # шаг оптимизации
    loss.backward()
    update(f_q.params)

    # обновление очереди (одновременное удаление старых и добавление новых)
    dequeue_and_enqueue(queue, k)

Эволюция: MoCo v2 и MoCo v3

Оригинальный метод (MoCo v1) использовал стандартный ResNet в качестве базового энкодера. После публикации конкурирующего подхода SimCLR архитектура была доработана:

  • MoCo v2: Авторы интегрировали успешные решения из SimCLR, в частности нелинейный проекционный блок (MLP-head) после базового энкодера, более агрессивные цветовые аугментации данных (color jitter) и косинусное расписание скорости обучения (cosine learning rate schedule). Это позволило MoCo v2 превзойти SimCLR при значительно меньших требованиях к размеру батча и объёму вычислений.
  • MoCo v3: Эмпирическая адаптация принципов MoCo к архитектурам типа Vision Transformer (ViT). В этой версии авторы отказались от использования явной очереди, компенсировав её отсутствие использованием больших размеров батча (до 4096). Однако фундаментальный механизм импульсного обновления энкодера (momentum encoder) был сохранён, так как он оказался критически важным для стабилизации обучения трансформеров в SSL-задачах.

См. также

Литература

  • Oord A. v. d., Li Y., Vinyals O. Representation learning with contrastive predictive coding // arXiv preprint arXiv:1807.03748. — 2018.
  • Wu Z., Xiong Y., Yu S. X., Lin D. Unsupervised feature learning via non-parametric instance discrimination // CVPR. — 2018. — С. 3733–3742.
  • He K., Fan H., Wu Y., Xie S., Girshick R. Momentum contrast for unsupervised visual representation learning // CVPR. — 2020. — С. 9729–9738.
  • Chen T., Kornblith S., Norouzi M., Hinton G. A Simple Framework for Contrastive Learning of Visual Representations // ICML. — 2020. — С. 1597–1607.
  • Chen X., Fan H., Girshick R., He K. Improved baselines with momentum contrastive learning // arXiv preprint arXiv:2003.04297. — 2020.
  • Chen X., Xie S., He K. An empirical study of training self-supervised vision transformers // ICCV. — 2021. — С. 9640–9649.
Личные инструменты