Update orthonormal_fro_reg

[alekfrohlich]

There are several subtleties involved in implementing orthonormality regularization correctly. With this pull request, I would like to initiate a broader discussion on best practices, both from a statistical perspective (e.g., unbiasedness, variance properties) and from a computational one (e.g., efficient use of GPU resources and PyTorch primitives).

In the first commit, I have implemented an unbiased estimator of the following regularization term:
[
\Omega_{1}(\theta) = \lVert C_{u_{\theta}(X)} - I_{d}\rVert_{F}^2 + 2\lVert \mathbb{E}, u_{\theta}(X)\rVert_{2}^2,
]
which directly enforces orthonormality. Here, (C_{u_{\theta}(X)}) denotes the (non-centered) covariance, and the second term enforces orthogonality with respect to the constant function (1_{\mathcal{X}}), which corresponds to the first singular function of the conditional expectation operator.

Previously, we relied on centered covariances; however, in that case the resulting regularizer is no longer an unbiased estimator of the objective above. I am very much looking forward to discussing this further and aligning on the most appropriate formulation moving forward.