Как увеличивать робастность модели и нужно ли вообще это делать?

Владимир «ternaus» Игловиков недавно опубликовал на Хабре пост про свою библиотеку Albumentations. Этот пост — результат серии бесед с людьми из индустрии, которые используют эту либу. Я, кстати, тоже поучаствовал, правда, уже после публикации поста — получилось весьма прикольно, постарался насыпать медицинской специфики и заодно задумал в будущем написать пост про аугментации в medical imaging.

Но сегодня не об этом — среди доменов и способов использования библиотеки в посте упомянут интересный пункт — «тестирование моделей перед деплоем на устойчивость к дрифту данных». Мы как раз недавно столкнулись с этой историей. Команда Flux (рентген и флюшка лёгких) решила проверить модельку на устойчивость к небольшим трансформациям. Вот какую картину дал сдвиг изображения на пиксель

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

Определений слова «робастность» можно встретить очень много, но в самых общих чертах — это способность модели сохранять своё качество работы при различных изменениях входных данных. Изменения можно рассматривать относительно конкретной картинки, а можно относительно обучающей выборки. Типы изменений при этом могут быть самые разные:

• натуральные — доменные сдвиги (новый тип оборудования или болезнь), редкие данные и выбросы, natural corruptions (засветы, неверная укладка пациента);

• баги — на обучении и на инференсе используется разный тип интерполяции или другая версия библиотеки по работе с изображениями;

• adversarial — специально созданная трансформация данных, которая визуально или семантически не отличается от исходных данных.

Что делать?

Как увеличивать робастность модели и нужно ли вообще это делать?

Ответ на второй вопрос не так очевиден, как может показаться. Увеличение робастности — не бесплатная процедура, чаще всего оно сопровождается снижением метрик на стандартных тестовых сетах. Выбор баланса между «классическим качеством» и робастностью зависит от домена, цены разных видов ошибок, стабильностью и дисперсией распределения данных и, конечно, от самой задачи. К примеру, в классическом Kaggle-соревновании со случайным train/test-сплитом вряд ли стоит тратить на робастность много усилий =) Подробнее об этом trade-off и вообще об обучении робастных моделей можно узнать, например, из этого видео.

Способы увеличения робастности модели можно разделить на три большие группы:

• Дата-методы. Сюда можно включить аугментации (включая Test-Time Augmentation), доразметку данных, генерацию синтетических данных.

• Архитектурные методы. Использование специальных слоёв или модулей, постпроцессинг с использованием доменного знания.

• Оптимизационные методы. Adversarial training, доменная адаптация, flat minima optimizers.

Adversarial training

Подробнее хочется остановиться на области adversarial training. Для начала определим, что такое adversarial attack. Будем считать, что это какое-то изменение входного изображения (или его фичей внутри сети), которое определяется двумя свойствами:

• Множество пертурбаций — насколько сильно мы разрешаем изменять входное изображение. Понятно, что если закрасить всю картинку чёрным цветом, то модель выдаст другое предсказание, но такие случаи рассматривать неинтересно. Множество пертурбаций может быть задано Lp-нормой (L-inf- самое большое изменение среди всех пикселей, L2 — евклидово расстояние между исходным и изменённым изображениями, L0 — количество изменённых пикселей), пределами геометрических трансформаций (сдвиг по пикселям, угол поворота) или максимально допустимым шумом в пространстве фичей.

• Механизм оптимизации — алгоритм поиска худшей пертурбации внутри заданного множества.

Как сделать сеть робастной к атакам внутри заданного множества? Очень просто — давайте во время обучения искать такие пертурбации картинки, которые ведут к высокому лоссу, и делать по ним gradient descent. Процесс adversarial training таким образом как раз и описывается двумя этапами:

• Внутренняя петля — поиск худшего примера внутри множества пертурбаций.

• Внешняя петля — backprop на этом примере.

Последнее недостающее звено — имплементация поиска худшей пертурбации. Тут тоже есть варианты:

• Случайный выбор — по сути аугментация внутри заданного множества.

• Worst-of-K — генерим K аугментаций, считаем лосс, делаем backprop только по худшей.

• Grid Search — дискретизируем пространство аугментаций и ищем «худшую».

• First-order adversarial learning — методы, которые позволяют найти пертурбацию, максиимизирующую лосс, с помощью процедуры оптимизации. К этой группе относятся Fast Gradient Sign Method и Projected Gradient Descent. О них можно подробнее узнать из того же видео или почитать здесь.

Датасеты

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

• ImageNet-A — картинки из интернета, на которых ошибается стандартный ResNet-50.

• ImageNet-C, ImageNet-P — тестовый сет ImageNet с синтетическими аугментациями типа блюра или изменения яркости.

• ImageNet-O — Out-of-Domain сет, картинки с классами, которых нет в датасете ImageNet-1k.

• ImageNet-R — изображения искусства, мультфильмов, граффити, татуировок, игр для оценки робастности к доменным сдвигам.

• ImageNet-S — картинки, полученные с разными интерполяциями при ресайзе и с помощью разных библиотек (Pillow, OpenCV).

Нюансы

Почему не все модельки по умолчанию пытаются сделать робастными к всевозможным трансформациям?

1) Часто это ведёт к ухудшению метрик на «стандартных» тестовых сетах и не всегда выгодно для конкретной задачи

2) Многие методы робастного обучения значительно замедляют обучение.

3) Робастность к одному типу шума совсем не означает робастность к другому. Более того, например, робастность к L-inf-пертурбациям совсем необязательно ведёт к робастности к L2-трансформациям

4) Сложно априори выбрать правильный сет пертурбаций, который улучшит качество и надёжность модели в продакшн-среде

Тестирование моделей на устойчивость хотя бы к простым геометрическим или пиксельным преобразованиям, однако, ещё никому не вредило. По результатам тестирования уже можно будет принимать решение — нужно ли что-то делать.

Если вы хотите узнать ещё больше об организации процессов ML-разработки, подписывайтесь на наш Телеграм-канал Варим ML