AIFS graph recipe and model config

[pimmeerdink]

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

I'm trying to replicate results from the AIFS paper, but would need to understand the recipe of the graph in order to start training the model myself and working with it. Are there plans to release the graph recipe and model config for the AIFS model?

Describe the solution you'd like

Ideally, there would be a subdirectory with all the necessary input files to construct AIFS. Even better: a notebook or example code that defines the AIFS model from scratch.

Describe alternatives you've considered

I have looked into loading the model from the weights present in the huggingface checkpoint, however here we can load the resulting graph (._graph_data), but it is not clear what reproducible recipe led to this graph. The same goes for the model config.

Additional context

No response

Organisation

No response

[semvijverberg]

I would be keen on the model config as well!

[HCookie]

Hi,
Thanks for your interest in the anemoi project and in the AIFS.

As you correctly identified, at the moment the checkpoint we are providing through huggingface does not provide the necessary metadata for loading within anemoi. Primarily the release of the AIFS is intended for inference and associated experiments.

Regarding the configs necessary to reconstruct the model, as anemoi is intended for more applications then just the AIFS, a subdirectory wouldn't make too much sense.

If you wish to begin experimenting with graph transformers and the architecture that becomes the AIFS, the default configs provide a good starting point.
The documentation 'should' also provide information on how to go about building a dataset which becomes the foundation that a graph is built on.

Let us know if you have any further questions.

[pimmeerdink]

Hi @HCookie and team,

Thank you for your response and for providing clarity on the current purpose of the AIFS release. I appreciate the work that has gone into this project and the resources made available so far.

I understand the broader scope of the anemoi project and the reasoning behind not including a dedicated subdirectory for AIFS configurations. However, having access to the specific graph recipe and model configuration used to generate the current checkpoint would be incredibly helpful for those of us looking to adapt and experiment further with the model architecture, especially in novel use cases.

While I see the potential of starting with the default configs, knowing the exact recipe that led to the AIFS graph would allow me to understand its nuances better and make modifications while maintaining its integrity. It would also ensure reproducibility of the AIFS model, which is a key aspect of research and experimentation.

Would it be possible to share a detailed example (e.g., a notebook or code snippet) illustrating how the graph for AIFS was constructed and the corresponding model configuration? A high-level walkthrough or annotated example in the documentation would also be greatly appreciated.

Thank you again for your work on this project and for engaging with the community. Looking forward to your thoughts!