Big Self-Supervised Models are Strong Semi-Supervised Learners
Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, Geoffrey Hinton
Статья: https://arxiv.org/abs/2006.10029
Код: https://github.com/google-research/simclr

Продолжаем серию про большие модели. Теперь картиночная SimCLRv2 от команды Google Brain с Хинтоном. Это пока не пиллиарды и бразиллионы параметров, как в GPT-3, но тем не менее.

В каком-то смысле это адаптация идей GPT (там работы назывались похоже — “Language Models are Unsupervised Multitask Learners”, “Language Models are Few-Shot Learners”) на картиночную область, также хитрым претрейнингом добиваемся получения большой модели, которую затем легко дообучить под конкретную задачу.

Авторы обучают глубокую и широкую модель (ResNet) в режиме self-supervised и затем файнтюнят на небольшом числе размеченных примеров (в случае ImageNet это стандартные сабсеты на 1% и 10%) и ещё потом всё дистиллируют.

Итоговый пайплайн таков:
#1. Unsupervised pretraining используя SimCLRv2 (наследник также весьма свежего SimCLR, https://ai.googleblog.com/2020/04/advancing-self-supervised-and-semi.html)
#2. Supervised fine-tuning на небольшом числе размеченных примеров.
#3. Дистилляция на неразмеченных примерах из целевой задачи.

Обо всём по порядку.

Оригинальный SimCLR первой версии (“A Simple Framework for Contrastive Learning of Visual Representations”, https://arxiv.org/abs/2002.05709) использует подход contrastive learning, когда хорошие представления выучиваются через максимизацию согласия между трансформированными изображениями. Представления картинок, произошедших от общего оригинала, должны получаться близкими, а от разных — далёкими.

Каждая картинка в батче дважды аугменируется (случайный crop, изменение цвета, Gaussian blur), прогоняется через ResNet и затем MLP, а после считается contrastive loss между всеми позитивными парами (негативные пары явно не сэмплятся, но это все остальные комбинации картинок внутри батча, они влияют на нормализацию лосса).

SimCLRv2 отличается в трёх аспектах:

1) Более глубокий и широкий ResNet (вместо ResNet-50 берём ResNet-152 с 3x больше каналов и Selective Kernels (https://arxiv.org/abs/1903.06586). Даёт +29% относительного улучшения по top-1 accuracy при файнтюнинге на 1% размеченных картинок.
2) Нелинейная часть сети, projection head (MLP после ResNet), становится глубже. SimCLR отбрасывала эту часть после предобучения, здесь делаем файн-тюнинг со среднего слоя этой части. Ещё +14% относительного улучшения.
3) Добавляется механизм памяти из MoCo (https://arxiv.org/abs/1911.05722). Ещё где-то +1%.

После предобучения и файнтюнинга делают дополнительную дистилляцию без меток, когда уже обученная сеть-учитель предсказывает метки неразмеченной части (уже новой задачи), а сеть-студент пытается к этому предсказанию приблизиться. Метки, если есть, в процедуру тоже несложно добавить. Дистиллировали в такую же архитектуру (self-distillation) или в более компактную.

Обучали на 128 Cloud TPUs, во время файнтюнинга и дистилляции делали случайные кропы и горизонтальные развороты.

Результаты хорошие.

Большие модели более label-efficient и улучшаются намного сильнее на небольшом числе примеров. Более глубокие projection heads также помогают (выучивают более качественные представления), особенно когда размеченных данных мало. Дистилляция только по неразмеченным примерам работает прекрасно, и небольшое число размеченных мало что к ней добавляет.

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

Ещё из интересного, ощутимо бьют 100%-supervised baseline (здесь правда может быть корректнее сравнивать по computational budget или чему-то похожему, ибо SimCLR внутри по сути делает много всякой аугментации данных)