Multi-domain: A dynamic way of training across domains and resolutions

[einrone]

Is your feature request related to a problem? Please describe.

Issue:

Currently, anemoi-core does not support dynamically switching between different global datasets, regional datasets or experiment with different resolutions during the same training without introducing additional encoders and decoders.

Each training run assumes a fixed spatial configuration (same domain and resolution), which makes it difficult to train a model with mixed datasets where samples come from different regions and resolutions during each training iteration steps.

without multi-domain:

1. Introducing additional transfer learning steps when fine-tuning to a new domain.
2. Transfer learning can cause catastrophic forgetting, where learned features of previous steps may be forgotten.
3. When training with multiple encoders, the similarities between datasets are not optimally used.
4. Harder to achieve cross-domain or cross-resolution generalization as the model is not able to continuously see different regional (or global at different resolution) samples
5. Harder to get a good prediction on a unseen domain i.e domain(s) that is not contained in the training dataset (model has not generalized)

Describe the solution you'd like

Quick summary of the solution:

multi-domain is a concept that allows the user to train a model on different datasets simultaneously, without introducing additional model components, e.g encoders and decoders. Compared to multi-dataset PR, the model architecture is a single encoder-processor-decoder structure that can be supported with different back bones. This is done by removing all graph and grid dependency in the model architecture i.e not initializing the model (during creation) with a graph. Tested with both deterministic and ensemble models (implemented in a fork on the DestinE repository). Down below there is an illustration of the model trains, see gif:

During training the dataloader holds different anemoi dataset objects which are labeled. For each training iteration the sample is randomly fetched and provides an output containing the data (shape = [time, Ensemble, Variable, Grid]) and its label, which is ties the data and its corresponding graph. An image of the output from the dataloader is provided

Additional details can be provided in a document or in a meeting.

What needs to be changed/modified/implemented:

Each of these tasks can be a separate PR or combined PRs. This can be of course be discussed

anemoi-models:

• layers/mapper.py (remove all grid dependencies)
• layers/block.py (remove all grid dependencies)
• layers/processor.py (remove all grid dependencies)
• interface/init.py (add support for multi-domain)
• models/multi-domain.py (introduce multi-domain model inheriting enc-proc-dec or ens-enc-proc-dec)

anemoi-training:

• training/train.py (accomodate for multi-domain functionalities, multiple dataset and graph objects)
• need to modify metadata (supporting arrays, indexes, etc..)
• datamodule and dataloader (separate or build in multi-domain functionalities into existing modules and dataloaders)
• train/task/base.py (accomodate for multi-domain or create separate class)
• train/task/ensforcaster.py (accomodate for multi-domain or create separate class)
• train/task/forcaster.py (accomodate for multi-domain or create separate class)

anemoi-graphs
no change required

Regarding loss functions and scalers, in the current refactor each domain has its own loss and scaler object, avoiding modifying all scalers and loss. However for an upcoming PR about loss and scalers to fit multi-domain, I am quite open to suggestions and ideas to implement this part in the best way possible to fit other use cases and multi-domain.

Describe alternatives you've considered

No response

Additional context

For those of you who are interested in running test-cases or wants to take a look at the multi-domain code base, please look at the link down below:
multi-domain-repo

Config and guidance can be provided, just let me or sophie know.

Organisation

Norwegian Meterological Institute
KNMI (Royal Netherlands Meterological Institute)

[ssmmnn11]

This will be a very nice way to train a model. I am wondering if we are basically very close with latest code changes. So one could

-> create a multi dataset config with all datasets.
-> use a single encoder and decoder (for example, link the encoder[datasetname] all to the same encoder instance, same for decoder)
-> if a dataset sample is empty, skip the dataset

-> draw a random number in each dataloader instance which dataset to keep
-> only read / return data for this dataset, otherwise return empty batch.

alternatively, the filtering could also be done inside the model, but then we have more IO than required.

Edit: even simpler: it should be enough to have a encoder / processor / decoder per dataset instead of a shared processor (but all point to the same instances / weights). Then we only need to sort out the time alignment.

[einrone]

Indeed, in the DestinE repository (our fork, where all of this has been implemented — link provided above), we support both deterministic and probabilistic forecasts. For the deterministic models, we observe significant improvements at pressure levels. For surface variables, performance is as good as or better than fine-tuned regional models at longer lead times. We have not yet evaluated CRPS, but we expect similar behavior.

The general idea is that all domains share the processor, encoder, and decoder. You mentioned weight sharing; however, having a separate instance per domain or dataset would not be ideal, since all data originate from the same input source and format, and it would unnecessarily complicate the architecture. Our goal is to generalize each component (encoder, processor, decoder) as much as possible.

I am currently in the process of implementing part of this. Since the model has been trained on large amounts of data — both from global pretraining and regional training — we aim to predict missing variables over regional domains that lack certain inputs (e.g., 2m temperature, t2m). Because the model has learned representations of such variables and encoded them in latent space, it should be able to use this knowledge to infer and output missing variables for a given domain. This can be achieved using masking, cross-attention, and related techniques (this is ongoing work).

To use the multi-dataset feature, we need to enforce a grid-independent architecture — that is, remove graph-specific information when creating the model instance. Currently, the encoder and decoder in the multi-dataset setup do not yet fully support this, which prevents us from using a single shared encoder/decoder across grids.

The current implementation is flexible: enabling dynamic_mode: True in the configuration activates multi-domain support (when set to False, the model runs in normal single-domain mode, as in ecmwf/anemoi-core). In addition, regional datasets must be specified. However, many aspects are automated, so even with 200 datasets, the configuration remains manageable and clean.

Regarding the data loader, we do not skip any data — the model consumes whatever the user provides, which is why we describe the system as flexible. For example:

Scenario 1 (only regional data in SG format; N320 is a global grid):
MEPS, ARA, CERRA, MEPS, ARA, ARA, UWC

Scenario 2 (mixture of regional and global data):
MEPS, MEPS, ARA, CERRA, AA, UWC, UWC, IFS 9km, MEPS, O96

Scenario 3 (only global data):
ARPEGE, IFS 9km, ICON, IFS 9km, IFS 9km, ICON, ARPEGE, ARPEGE

As shown above, this implementation allows the model to be trained on arbitrary combinations of analyses, forecasts, or reanalyses — whether alternating between regional and global data, or across different global grids and resolutions.

Training can be performed either in stages (global pretraining → regional fine-tuning) or in a single phase where regional and global datasets are alternated. The latter approach yields the same results as staged training and has comparable computational cost (time and GPU usage).

[ssmmnn11]

Thank you for sharing. I think we talk about the same thing. I don’t think you need many changes. Can you point to where you think the mappers are grid specific?

[einrone]

Looks like all references to sub_graph, sub_graph_edge_attributes, src_grid_size and dst_grid_size has been removed. Graphproviders is merged and all forward methods (forwardmapper, backwardmapper and processor.forward) has graph attrs and indexes as input. That is basically what has been implemented in our fork in DestinE repo. Looks great so far, when I am done with the proposal writing, I will try to set this up and see if I can reproduce my results. However I still think we need to implement the dataloader. We can just expand on nativegrid and ensgriddataset classes for example.