HEAL-ViT: Vision Transformers on a spherical mesh for medium-range
  weather forecasting
Vivek Ramavajjala
In recent years, a variety of ML architectures and techniques have seen
success in producing skillful medium range weather forecasts. In particular,
Vision Transformer (ViT)-based models (e.g. Pangu-Weather, FuXi) have shown
strong performance, working nearly "out-of-the-box" by treating weather data as
a multi-channel image on a rectilinear grid. While a rectilinear grid is
appropriate for 2D images, weather data is inherently spherical and thus
heavily distorted at the poles on a rectilinear grid, leading to
disproportionate compute being used to model data near the poles. Graph-based
methods (e.g. GraphCast) do not suffer from this problem, as they map the
longitude-latitude grid to a spherical mesh, but are generally more memory
intensive and tend to need more compute resources for training and inference.
While spatially homogeneous, the spherical mesh does not lend itself readily to
be modeled by ViT-based models that implicitly rely on the rectilinear grid
structure. We present HEAL-ViT, a novel architecture that uses ViT models on a
spherical mesh, thus benefiting from both the spatial homogeneity enjoyed by
graph-based models and efficient attention-based mechanisms exploited by
transformers. HEAL-ViT produces weather forecasts that outperform the ECMWF IFS
on key metrics, and demonstrate better bias accumulation and blurring than
other ML weather prediction models. Further, the lowered compute footprint of
HEAL-ViT makes it attractive for operational use as well, where other models in
addition to a 6-hourly prediction model may be needed to produce the full set
of operational forecasts required.