OMNI: один компактный энкодер собрал знания 10 моделей для патологии

Команда исследователей опубликовала на arXiv работу Scaling Multi-Teacher Distillation for Digital Pathology. В ней показано, как знания десяти передовых foundation-моделей для цифровой патологии сжать в один компактный энкодер. Результат называется OMNI. По заявлению авторов, такие студенты до 10 раз эффективнее крупных моделей-учителей. При этом они сопоставимы с ними по качеству, а на значительной части задач их обходят. Проверка охватила более 10 foundation-моделей и 39 задач на уровне отдельных тайлов и целых слайдов.

Где ломается мультиучительская дистилляция

В цифровой патологии за последнее время появилось много сильных foundation-моделей. Каждая обучена на своих данных и по своему рецепту, и у каждой свои сильные стороны. Идея мультиучительской дистилляции (MTD) напрашивается сама. Можно собрать нескольких учителей и передать их знания одной студенческой сети. Тогда не придётся гонять десяток тяжёлых энкодеров на гигапиксельных снимках.

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

Четыре решения вместо одного

Чтобы студент сохранял особенности каждого учителя, авторы меняют и архитектуру, и функцию потерь. Для каждого учителя вводятся обучаемые токены. Это даёт сети отдельную «точку входа» под каждого эксперта. Схема внимания специально настроена под такую конструкцию. В энкодер добавлены дополнительные слои mixture-of-experts, которые разводят вычисления по специализированным веткам. Контрастивный лосс дополнительно разделяет представления разных учителей.

В результате метод лучше масштабируется по числу учителей. Именно это не получалось у базового MTD.

Что проверили и что выпустили

Ключевые параметры работы:

• 10 передовых foundation-моделей использованы как учителя;
• более 10 foundation-моделей оценено в сравнении;
• 39 задач, включая уровни тайлов и целых слайдов;
• студенты до 10x эффективнее крупных учителей;
• качество на уровне учителей, а местами выше.

По итогам авторы выпускают коллекцию OMNI — набор эффективных foundation-моделей, обученных на знаниях десяти учителей. Отдельно они проанализировали выученные представления. Оказалось, что они комплементарны, то есть несут разную информацию. Этим авторы объясняют, почему их легко агрегировать на downstream-этапе.

Что менять в практике

Для команд, которые строят пайплайны в патологии, главный практический вывод касается стоимости. Анализ слайдов упирается в вычисления: один слайд разбивается на огромное число тайлов, и каждый пропускается через энкодер. Модель, которая в десять раз эффективнее при сопоставимом качестве, напрямую снижает стоимость инференса и позволяет запускать обработку на более скромном железе.

Второй вывод — методический. Если вы дистиллируете из нескольких моделей, не стоит сливать их в единый целевой вектор. Отдельные токены, MoE-слои и контрастивное разделение, вероятно, пригодятся и вне патологии, в любых доменах с множеством сильных специализированных моделей. Это пока предположение, в статье оно не проверялось.

Ограничения стоит учитывать. Сведения основаны на аннотации препринта, детальные метрики по задачам нужно смотреть в самой работе. Условия лицензирования весов OMNI и результаты на ваших данных надо проверять отдельно. Для клинического применения одной точности на бенчмарках недостаточно: нужны валидация на данных конкретной лаборатории и регуляторные процедуры.