[WIP] optimize last layer

[misko]

No description provided.

[anar-rzayev]

Hi @misko Really interesting optimization approach here. I am sorry again if interrupting any future review conversations here, but I am trying to understand the main performance gains you are targeting - it looks like you are setting mmax=0 in the final layer and truncating the Wigner coefficients. I am curious about what kind of speedup/memory savings this provides and if you have noticed any accuracy trade-offs so far? Also, quick question about the implementation - why create a new CoefficientMapping for the last layer rather than reusing the existing one? And is the Wigner truncation [:, : self.lmax + 1] directly tied to the mmax=0 change, or are those separate optimizations working together? Thanks very much as usual for all these interesting new features and open-sourcing the changes for the community

[github-actions]

This PR has been marked as stale because it has been open for 30 days with no activity.

[misko]

Hi @anar-rzayev !
Sorry i didnt notice your comment from earlier 😢

In UMA we compute the full L=2 last layer, however we only use L=0 features in the EFS head https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/models/uma/escn_md.py#L670

This means we dont actually have to compute L=1,2 in the last layer, which in a 4 layer network can save ~2-5% runtime.