Describe the task. It can be a feature, documentation, etc.
for xai work:
We are using WeatherGenerator for Layer-wise Relevance Propagation (LRP) interpretability work. This requires computing gradients with respect to the ERA5 input tokens, rather than the model weights, so the autograd computation graph needs to remain intact through the forward pass.
While looking at this workflow, we noticed a small autograd-related point in EmbeddingEngine.forward in engines.py that is relevant for interpretability use cases. The current code is:
tokens_all.scatter_(0, scatter_idxs, torch.cat(x_embeds) + pe_embed[pe_idxs])
Our understanding is that, because tokens_all is a requires_grad=False tensor, the in-place scatter_ operation does not preserve the autograd graph back to the ERA5 input tokens. The forward-pass values appear to be unchanged; the difference is in whether gradients can be tracked for interpretability workflows.
We have tested a few possible ways of handling this locally, and would welcome your guidance on which approach would best fit the WeatherGenerator codebase.
Agreed upon config flag change
Hedgedoc URL, if you are keeping notes, plots, logs in hedgedoc.
No response
Area
Describe the task. It can be a feature, documentation, etc.
for xai work:
We are using WeatherGenerator for Layer-wise Relevance Propagation (LRP) interpretability work. This requires computing gradients with respect to the ERA5 input tokens, rather than the model weights, so the autograd computation graph needs to remain intact through the forward pass.
While looking at this workflow, we noticed a small autograd-related point in EmbeddingEngine.forward in engines.py that is relevant for interpretability use cases. The current code is:
tokens_all.scatter_(0, scatter_idxs, torch.cat(x_embeds) + pe_embed[pe_idxs])
Our understanding is that, because tokens_all is a requires_grad=False tensor, the in-place scatter_ operation does not preserve the autograd graph back to the ERA5 input tokens. The forward-pass values appear to be unchanged; the difference is in whether gradients can be tracked for interpretability workflows.
We have tested a few possible ways of handling this locally, and would welcome your guidance on which approach would best fit the WeatherGenerator codebase.
Agreed upon config flag change
Hedgedoc URL, if you are keeping notes, plots, logs in hedgedoc.
No response
Area