Riemannian inertial updates for normalized embeddings

(1) Integration point: apply this optimizer only to a tensor of unit-normalized vectors, initially class prototypes in a contrastive-learning head or token/feature prototypes in a retrieval layer. Keep the backbone optimizer unchanged for the first experiment. Let x have shape [N,D], with every row normalized, and let v have the same shape as a persistent optimizer state.

(2) Pseudocode:

initialize x = normalize(x)
initialize v = zeros_like(x)
for each minibatch:
    loss = network_loss(backbone(batch), x)
    g = grad(loss, x)
    gx = g - row_sum(g*x, axis=D, keepdim=True)*x
    v = beta*v - eta*gx
    x_trial = x + v
    x = x_trial / (row_norm(x_trial) + eps)
    v = v - row_sum(v*x, axis=D, keepdim=True)*x

Use beta=exp(-eta/m) as a damping-inspired default and compare beta values 0.9, 0.95, and 0.99. A semi-implicit variant can use v=(beta*v-eta*gx)/(1+eta/m), with m and damping exposed as hyperparameters.

(3) Computed from the paper's mathematics: the tangent projector, unit-sphere constraint, inertial velocity, and retraction. Estimated empirically: eta, beta or m, gradient-noise sensitivity, and whether coupling across prototypes helps. Optionally add the paper's mean-field coupling as c_i=lambda*(mu-(mu dot x_i)*x_i), where mu=(1/N)sum_j x_j, but begin with lambda=0 to isolate the optimizer effect.

(4) Cheap experiment: train a ResNet-18 or ViT-Tiny on CIFAR-100 with a supervised contrastive head containing 100 unit class prototypes. Compare SGD with momentum, AdamW plus post-update normalization, and this spherical inertial optimizer at matched parameter-update FLOPs. Record training loss versus optimizer steps, validation accuracy, prototype norm error, gradient-to-update ratio, and oscillation after learning-rate reductions. Success means faster loss descent or higher accuracy at equal steps without norm drift or prototype collapse; target at least 10% fewer steps to the same validation accuracy.