Не трансформером единым: масштабируемые сферические CNN для научных приложений

Друзья, разбавлю ваш (и свой) думскроллинг обрзовательным контентом. 💔

Мир не плоский. Есть много данных, которые лучше описываются сферической топологией. Например, данные о климате, космосе, панорамные съемки, что уж скрывать, в медицине только рентген и SPECT - планарные снимки, а МРТ, КТ И PET в виде сырого сигнала - вообще-то нет.

Оказывается, обрабатывать такие данные с помощью нейронных сетей не так просто. Проекция сферы на плоскость даёт сильные искажения в районе полюсов и потерю информации. Применение обычных сверточных нейросетей (CNN) или трансформеров (ViT) наталкивается на проблему выборки: как определить равномерную сетку на сфере. Да ещё и теряется важное свойство инвариантности вращения. При повороте входных данных на любой угол выходные данные тоже должны повернутся на тот же угол. Это важно для объектов с разной 3D ориентацией, напр, молекул или человека в томографе.

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

В статье для ICML 2023 Google AI преодолели эти ограничения и опубликовали оперсорс библиотеку на JAX, где реализовали:
- сферические свертки с дополниной степенью свободы, связанной с поляризацией;
- ввели новый слой активации, уменьшающий фазовые сдвиги между слоями;
- ввели спектральную batch нормализацию;
- ввели новый остаточный блок (который для предотвращения затухания градиентов), распараллелив передачу высокочастотных и низкочастотных признаков.

Точность и эффективность проверили на задачах прогнозирования погоды до 28 дней и регрессии молекулярных свойств, получили красивые результаты, выдали нам порцию вдохновения и новую блестящую лопату копать МЛ дальше

🗞 Статья
🔣 Код