Diffeomorphic gauge-fixing layer

1. Integration point: wrap a CNN, U-Net, neural operator, or latent world model whose input is a spatial tensor uₜ with shape batch × channels × height × width. Add a small WarpNet that predicts a stationary velocity field vᵤ(x) with two spatial components. Convert it to a smooth warp gᵤ = exp(vᵤ) using scaling-and-squaring. For a first MVP, use the simpler bounded map gᵤ(x) = x + α tanh(vᵤ(x)), with α less than one eighth of the smaller image dimension. 2. Pseudocode: v = WarpNet(u, reference); g = exp_velocity(v); z = grid_sample(u, inverse(g)); z_next = DynamicsNet(z); u_hat = grid_sample(z_next, g); loss = MSE(u_hat, target) + λ times warp_smoothness(g) + η times cycle_error(g, inverse(g)). The alignment loss is mean squared error between z and reference, plus the displayed Jacobian and Laplacian penalties. 3. Use bilinear grid sampling for both warps. With stationary velocities, approximate the inverse by exp(−v). Compute the mathematical components explicitly: the group action, inverse composition, smoothness energy, and equivariance consistency. Estimate the reference field as either a learnable tensor or the batch mean, and learn WarpNet and DynamicsNet jointly. For augmentation consistency, sample a random smooth warp h and add MSE(Predict(h(u)), warp(Predict(u), h)). 4. First experiment: generate two-dimensional advection or reaction-diffusion trajectories, then apply independent smooth Fourier displacement fields or thin-plate-spline warps to each trajectory. Compare a U-Net dynamics baseline, a conventional spatial-transformer baseline, and this model at matched parameter count. Evaluate one-step MSE, 20- to 50-step rollout error, accuracy on unseen warp amplitudes, and normalized equivariance error ||P(a(h,u)) − a(h,P(u))||₂ divided by ||P(u)||₂ plus 10⁻⁸. A successful result is lower long-horizon error and better accuracy with fewer training trajectories, while the extra WarpNet costs less than approximately 10 percent of baseline inference FLOPs.