Scaling Deep Learning: Reducing GPU Overhead with Riemannian Manifold Projection

[Nossari]

Hi! I faced this exact issue in my architecture recently. Standard Euclidean layers were failing to maintain stability, so I developed a solution using a RotationalLinearLayer that constrains weight transformations to the Lie Group SO(n).
​This approach completely neutralized the gradient failure by keeping the Frobenius norm invariant throughout the training trajectory. You might find the mathematical implementation and the logic behind the dynamic kinetic damping useful for your current setup.
​You can check the implementation details here:
https://github.com/Nossari/Non-Euclidean-Deep-Learning-Framework/tree/main
​I'd be happy to discuss the math behind the manifold projection if you're interested!

[sowndappan5]

Interesting approach. Have you benchmarked the SO(n)-constrained layer against standard linear layers in large-scale distributed training? I'm curious whether the manifold projection introduces measurable communication or optimization overhead when used with DeepSpeed ZeRO.