HEALPix curriculum

[TillHae]

Describe the task. It can be a feature, documentation, etc.

Add HEALPix curriculum to the WeatherGenerator for more efficient training.

In my bachelor thesis, I discovered that starting at a lower HEALPix Level (HL) and sequentially increasing it during training results in efficiency gains. This coarse-to-fine curriculum reduces compute time by almost 95% in the early stages of training up to HL 6. Furthermore, it continues to save compute to a smaller extent as the curriculum progresses, all without losing final model performance. When training to even finer resolutions like HL 8 it should even result in a greater efficiency gain. We should integrate a configuration mechanism to automate these HL transitions and dataset swaps natively during training.

Hedgedoc URL, if you are keeping notes, plots, logs in hedgedoc.

No response

Area

• datasets, data readers, data preparation and transfer
• model
• science
• infrastructure and engineering
• evaluation, export and visualization
• documentation
• performance

[TillHae]

Assign Me

[github-actions]

Issue assigned to @TillHae.

[clessig]

@TillHae : do you have any quantitative results for how updating the healpix level increases training efficiency? Is this only at the input and output or also in latent space?

[TillHae]

The latent space isn't affected when updating the HEALPix level. On Monday, I presented at the High-Res meeting, and they asked me to run training on more samples to see how it impacts generalization. Regarding efficiency, it improves a lot, as you can see in the following plot. The black line is the baseline, which was run on HEALPix level 6 for the entire time.

[clessig]

The latent space isn't affected when updating the HEALPix level. On Monday, I presented at the High-Res meeting, and they asked me to run training on more samples to see how it impacts generalization. Regarding efficiency, it improves a lot, as you can see in the following plot. The black line is the baseline, which was run on HEALPix level 6 for the entire time.

Could you share the loss curves for the standard headline variables.

[clessig]

In latent space, the healpix level is always 6 in these experiments and it only changes for embedding and decoding?

[TillHae]

In latent space, the healpix level is always 6 in these experiments and it only changes for embedding and decoding?

Sorry, I mixed things up. The latent space does change by a factor of 4 but the model size doesn't. The only thing where we encounter a problem is with the attention mechanism. After a transition the old attention map doesn't make any sense and this is the loss spike that is then visible in the plot.

[TillHae]

Could you share the loss curves for the standard headline variables.

I didn't even know that it was possible to plot the training loss by variable. Which script do I need to run to get those results?

[clessig]

Could you share the loss curves for the standard headline variables.

I didn't even know that it was possible to plot the training loss by variable. Which script do I need to run to get those results?

uv run plot_train --channels z_500 u_850 v_850 t_850 q_850 2t 10u 10v

[clessig]

In latent space, the healpix level is always 6 in these experiments and it only changes for embedding and decoding?

Sorry, I mixed things up. The latent space does change by a factor of 4 but the model size doesn't. The only thing where we encounter a problem is with the attention mechanism. After a transition the old attention map doesn't make any sense and this is the loss spike that is then visible in the plot.

The attention map (you mean attention values, i.e. normalized softmax of the score matrix?) is determined by the input tokens and not learned. Do you mean something else?

[TillHae]

In latent space, the healpix level is always 6 in these experiments and it only changes for embedding and decoding?

Sorry, I mixed things up. The latent space does change by a factor of 4 but the model size doesn't. The only thing where we encounter a problem is with the attention mechanism. After a transition the old attention map doesn't make any sense and this is the loss spike that is then visible in the plot.

The attention map (you mean attention values, i.e. normalized softmax of the score matrix?) is determined by the input tokens and not learned. Do you mean something else?

What I meant is that the query and key vectors generated when the input tokens are applied to the learned weights of the global AE make no sense and therefore the "attention map" should also make no sense.

Here is the calculation: query and key calculation in ll. 595-600 of the attention

I think what you mean is ll. 953-975 of the PyTorch attention, which we use. But as I just pointed out the input for this calculation should be impacted by the learned weights of the global AE, which results in the output of the calculation being basically garbage at the beginning of a HL transition. But tbh I don't have a full overview of what is actually going on throughout the entire training process.

I think in the process of looking at everything I also found dead code:
self.att in ll. 583-588 of the attention is actually never used in the forward call of the MultiSelfAttentionHead.

[TillHae]

Could you share the loss curves for the standard headline variables.

I didn't even know that it was possible to plot the training loss by variable. Which script do I need to run to get those results?

uv run plot_train --channels z_500 u_850 v_850 t_850 q_850 2t 10u 10v

Hedgedoc with plots