Wristband Gaussian Loss: Formalization and Proof

Story

Mikhail Parakhin (CTO of Shopify) published the Wristband Gaussian Loss. A training regularizer that forces a neural encoder's outputs to be exactly distributed as a standard Gaussian. Available on Github.

The core idea is a coordinate change: every encoder output gets mapped to a direction (where it points on the unit sphere) and a radial percentile (where its magnitude sits in the chi-squared distribution). The insight is that the encoder is $\mathcal{N}(0, I)$ if and only if this pair is uniformly spread across the cylinder.

What I did specifically was to formalize the author's Python code into Lean and prove (with the help of LLMs):

• the equivalence between the standard normal space and this Wristband space.
• the unique minimization, if the repulsion loss the author picked is in fact uniquely minimized at the uniform distribution.

Notes

All theoretical ideas and the original implementation are Mikhail Parakhin's.
My contribution is the Lean formalization.