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

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

Контекст

Прогноз погоды и климата почти полвека опирается на численное моделирование физических процессов. Такие системы дороги в вычислениях и плохо впитывают новые потоки данных со спутников и радаров. За последнее десятилетие накопились огромные объёмы наблюдений, и глубокое обучение уже показало успехи в отдельных задачах, например в предсказании осадков. Однако архитектуры на основе свёрточных и рекуррентных сетей накладывают жёсткие предположения о структуре данных. Для хаотичной земной системы, где малые изменения начальных условий могут резко менять поведение, эти предположения могут не выполняться. Трансформеры умеют ловить дальние зависимости, но их прямое применение к трёхмерным данным упирается в квадратичную сложность внимания. Работа Earthformer стоит особняком среди соседей по архиву: рядом лежат модели для медицинской сегментации и слабо контролируемого обучения, а прямых конкурентов по прогнозу земной системы почти нет.

Методы

Авторы строят Earthformer как иерархический трансформер с кодировщиком и декодером. Ключевой блок — Cuboid Attention: входной тензор размером время × высота × ширина разрезается на непересекающиеся кубики, внутри каждого кубика параллельно выполняется self-attention. Разные способы разрезания (локальные или разреженные, со сдвигом) порождают разные паттерны внимания. Чтобы кубики обменивались информацией, вводится небольшой набор глобальных векторов: они собирают статистику со всех кубиков и возвращают её обратно. Архитектура имеет несколько уровней иерархии с понижением и повышением разрешения через свёрточные слои и интерполяцию. Декодер генерирует прогноз не авторегрессивно, а сразу из позиционных эмбеддингов. Модель обучается на синтетических наборах MovingMNIST и N-body MNIST, где цифры движутся по законам гравитации трёх тел, а затем проверяется на реальных данных об осадках и температуре океана.

Результаты

На синтетических данных Earthformer с паттерном Axial и глобальными векторами достигает MSE 41.79 на MovingMNIST и 14.82 на N-body MNIST, обходя все шесть базовых моделей. На реальном наборе SEVIR для прогноза осадков модель получает CSI-219 = 0.1791, что заметно выше, чем у Rainformer (0.0831) и UNet (0.0577). На данных ICAR-ENSO для прогноза Эль-Ниньо Earthformer показывает корреляцию C-Nino3.4-M = 0.7329, превосходя Rainformer (0.7106) и E3D-LSTM (0.7040). Глобальные векторы дают стабильный прирост без увеличения вычислительной сложности: число параметров растёт с 6.6 до 7.6 миллионов, а GFLOPS почти не меняются. Иерархическая структура снижает MSE на MovingMNIST с 48.04 до 46.91 при той же глубине. Качественно модель точнее предсказывает траектории цифр в N-body MNIST, тогда как базовые методы размывают изображение или ошибаются в движении.

Значение

Earthformer показывает, что трансформеры можно адаптировать к трёхмерным пространственно-временным данным без взрывного роста вычислений. В отличие от соседних работ, которые используют свёрточные или рекуррентные сети, здесь внимание ограничено кубиками, а глобальные векторы связывают локальные блоки. Это позволяет ловить дальние зависимости в хаотичной системе, сохраняя линейную сложность по размеру кубика. Работа задаёт новый ориентир для архитектур прогноза земной системы: паттерн Axial с глобальными векторами оказывается наиболее эффективным. Для науки о данных это пример того, как универсальный строительный блок можно настраивать под конкретную задачу, не теряя общности.

Развитие

Модель пока детерминирована и не оценивает неопределённость, из-за чего прогнозы могут быть размытыми. Вероятностное расширение Earthformer стало бы естественным шагом. Другое направление — включение физических знаний: сейчас модель чисто data-driven и не использует законы сохранения или уравнения динамики. Гибридные подходы, объединяющие нейросеть с численным моделированием, уже показывают перспективы. Возможен и автоматический поиск паттернов кубического внимания через Neural Architecture Search. Также интересно масштабирование на глобальные данные с более высоким разрешением.

Влияние

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

Следующие шаги

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

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

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

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

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

Ключевые числа

  • MSE на N-body MNIST: 14.82
  • CSI-219 на SEVIR: 0.1791
  • Корреляция C-Nino3.4-M: 0.7329
  • Число параметров Earthformer: 7.6 млн
  • Размер кубика по умолчанию: 3×2×2
Понятия
трансформеры прогноз IRI вероятностный прогноз климатическая модель весенний барьер предсказуемости климатическая сеть ERA5 прогнозирование механизм внимания сейсмометр реанализ диаграмма ховмюллера время упреждения машинное обучение многомерный индекс ЭНСО