Wristband Gaussian Loss: From $O(N^2)$ to $O(N)$

Context

This is a continuation of: Wristband Gaussian Loss: Formalization and Proof. The original work used pairwise computations for the repulsion term, to measure uniformity. By rearranging equations I was able to compute the repulsion term via spectral decomposition of the Wristband Space. The original equation separated the computation in two parts, the angular term and the radial term. The second used an approximation of a Neumann series.

Work

Consider the repulsion kernel,

$$K\bigl((u,t),(u',t')\bigr) = k_\text{ang}(u, u') \cdot k_\text{rad}(t, t'),$$

where $u$, $t$ are points on the wristband space. Our angular kernel admits a Mercer expansion in the spherical harmonic basis,

$$k_\text{ang}(u, u') = \sum_{\ell=0}^{\infty} \lambda_\ell \sum_{m=1}^{N_\ell} Y_{\ell,m}(u) Y_{\ell,m}(u').$$

With the radial kernel expressed in the cosine form, $$k_\text{rad}(t, t') = \sum_{k=0}^{\infty} a_k \cos(k\pi t) \cos(k\pi t').$$

The kernel energy to be minimized becomes,

$$\boxed{\mathcal{E}_\text{sp} = \lambda_0 \sum_{k=0}^{K} \tilde{a}_k \left(\frac{1}{N}\sum_i \cos k\pi t_i\right)^{\!2} + \lambda_1 \sum_{m=1}^{d} \sum_{k=0}^{K} \tilde{a}_k \left(\frac{\sqrt{d}}{N}\sum_i u_{im} \cos k\pi t_i\right)^{\!2}},$$

where $d$ is the latent space dimension and $K$ is the truncation term for the spectral approximation. The $\tilde{a}_k$ are normalization constants.

This means going from a $O(N^2)$ pairwise energy computation to $O(Nd)$, much faster even for small $N$.

This work was also proven with Lean 4. It was proven the existence of such kernel with such properties and that minimizing its energy meant converging to a uniform distribution as intended.

I have tried to get upper-bounds of the truncation error, but they are too conservative to be informative. As a rule of thumb, if we keep $K = 6$ and $\ell = 0, 1$, you should be good.