[Feature request]: Support using intermediate embeddings

[laserkelvin]

Feature/behavior summary

Refactor Embeddings data structure, and their use by OutputBlocks to allow the use of intermediate embeddings for modeling. In some models such as MACE, the output of the model is given as some reduction over projections of intermediate layers, i.e. $E_f = E_0 + E_1 + \ldots + E_l$ for $l$ layers.

The current implementation hinders this mode of usage a little, as we need to be able to store intermediate embeddings and then use them correctly after every layer is computed, which would be fine if the intermediate embeddings were the same shape and could be concatenated along a single (new) dimension, and the output blocks just broadcasts. In the case of MACE and other equivariance preserving models, the intermediate layers may have different shapes and can't be concatenated.

Request attributes

• [X] Would this be a refactor of existing code?
• [ ] Does this proposal require new package dependencies?
• [ ] Would this change break backwards compatibility?
• [ ] Does this proposal include a new model?
• [ ] Does this proposal include a new dataset?
• [ ] Does this proposal include a new task/workflow?

Related issues

72 is the main issue, but #83 is a related WIP PR

Solution description

One way would be to refactor Embeddings to allow intermediate embeddings; i.e. for node/system levels, we expect either a Tensor or a list[Tensor]. This would then need the logic of OutputBlock to be modified so that we use different output heads per intermediate embedding, then reduce them for the final output.

For now I don't think this will be backwards breaking.

Additional notes

No response

[laserkelvin]

Closing this for now, as we've agreed that it's not entirely necessary for MACE for now. We can re-open this issue if there are new use cases that demand this model, but having thought about it for a bit, it's not an entirely trivial task.