Кубики внимания для прогноза погоды ✎

Оригинал: "Earthformer: Exploring Space-Time Transformers for Earth System Forecasting"
arXiv:2207.05833 · 2022-07-12 · arXiv non-exclusive · собственный разбор · 2 мин
Новая архитектура Earthformer делит пространство и время на кубики, чтобы нейросеть училась предсказывать осадки и Эль-Ниньо точнее и быстрее.

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

Долгие годы синоптики полагались на численные модели, которые перемалывают уравнения физики атмосферы и океана. Эти машины точны, но медленны и дороги. Спутники и радары за десятилетия накопили горы данных, а нейросети уже научились предсказывать дождь по снимкам неба. Но старые архитектуры, вроде свёрточных и рекуррентных сетей, слишком жёстко привязаны к форме данных. А земная система хаотична: взмах крыла бабочки в Бразилии может отозваться торнадо в Техасе.

Трансформеры умеют ловить связи между далёкими точками, но их прямое применение к трёхмерным данным упирается в квадратичную сложность внимания. Каждый элемент хочет «посмотреть» на всех остальных, и вычисления взрываются.

Авторы Earthformer нашли способ обуздать эту жадность. Они режут входные данные на кубики — блоки размером время × высота × ширина. Внутри каждого кубика нейросеть делает self-attention, то есть взвешивает важность точек друг для друга. Кубики не видят друг друга напрямую, но их связывает небольшой набор глобальных векторов. Эти векторы собирают сводку со всех кубиков и возвращают её обратно, как курьер, который разносит новости по кварталам.

Архитектура многоуровневая: сначала данные сжимаются свёрточными слоями, потом снова разворачиваются. Декодер выдаёт прогноз не шаг за шагом, а сразу целиком, опираясь на позиционные эмбеддинги. На синтетическом наборе MovingMNIST, где цифры движутся по экрану, Earthformer достигает MSE 41.79, а на N-body MNIST, где цифры подчиняются гравитации трёх тел, — 14.82. Это лучше шести базовых моделей.

На реальных данных осадков SEVIR модель получает CSI-219 = 0.1791, тогда как у Rainformer 0.0831, а у UNet 0.0577. Для прогноза Эль-Ниньо корреляция C-Nino3.4-M достигает 0.7329, превосходя конкурентов.

Глобальные векторы дают стабильный прирост без роста вычислений: параметров становится 7.6 миллиона вместо 6.6, а GFLOPS почти не меняются. Иерархия снижает ошибку на MovingMNIST с 48.04 до 46.91. Качественно модель точнее следит за траекториями цифр в N-body MNIST, тогда как базовые методы размывают картинку или путают направление.

Раньше было неясно, как применить трансформеры к трёхмерным данным прогноза земной системы, не утонув в вычислениях. Теперь понятно: разрезание на кубики и глобальные векторы позволяют трансформеру обходить свёрточные и рекуррентные сети на задачах осадков и Эль-Ниньо. Модель можно использовать для оперативного прогноза осадков и температуры океана, а также как основу для гибридных физико-нейросетевых систем.

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

Набор N-body MNIST заставляет цифры двигаться по законам гравитации трёх тел — это хаотичная система, где малейшее изменение начальной скорости приводит к совершенно другой траектории. Модель должна угадать, куда полетит цифра «0», хотя даже точное решение неустойчиво.

Идея предсказывать хаотичную систему по ограниченным наблюдениям напоминает «Три тела» Лю Цысиня, где цивилизация пытается вычислить движение планеты в системе трёх солнц. Earthformer, как и герои романа, сталкивается с принципиальной непредсказуемостью, но учится давать лучший возможный прогноз.

Понятия
трансформеры прогноз IRI вероятностный прогноз климатическая модель весенний барьер предсказуемости климатическая сеть ERA5 прогнозирование механизм внимания сейсмометр реанализ диаграмма ховмюллера время упреждения машинное обучение многомерный индекс ЭНСО