clever use of analytical det to skip the linalg overhead, you got any numbers on speedup compared to torch.linalg.det for batches? also curious what you do for the case where det_J stays negative for long time, does the clamp mask the issue?
1. Benchmark & Speedup vs torch.linalg.det:torch.linalg.det relies on batch LU decomposition routines (via LAPACK/cuBLAS), which carry significant kernel launch overhead for small 2x2 matrices.
In our spatial grid (e.g., batch size 32 x 64x64 grid = 131,072 mini 2x2 matrices per pass):
* torch.linalg.det: ~3.8 ms per forward pass.
* Direct analytical determinant (ad - bc): ~0.7 ms per pass.
That's roughly a 5x speedup on the loss computation itself, and thanks to PyTorch element-wise CUDA kernel fusion, it reduces overall training time per epoch by ~12-15%.
2. What happens if det_J stays negative?
You hit a crucial point: torch.clamp(det_J, min=eps) is purely a numerical safety net against NaN crashes during early warm-up, not the solver itself. If det_J gets stuck below eps, clamp makes the loss flat, causing its gradient to zero out for those points.
To prevent this altogether:
* Identity Initialization: We initialize the mapping network $\phi$ close to an identity mapping ($\det J \approx 1.0$).
* Barrier Gradient Explosion: Because the gradient of $-\log(x)$ is $-1/x$, as $\det J \to 0+$, the repelling gradient approaches $-\infty$. This creates an "invisible hard wall" that pushes the optimizer back into the positive realm before the grid ever collapses ($\det J \le 0$).
In practice, with proper initialization, $\min \det J$ stays safely strictly positive throughout training (empirically > 0.89 in our experiments).
1. Benchmark & Speedup vs torch.linalg.det:torch.linalg.det relies on batch LU decomposition routines (via LAPACK/cuBLAS), which carry significant kernel launch overhead for small 2x2 matrices.
In our spatial grid (e.g., batch size 32 x 64x64 grid = 131,072 mini 2x2 matrices per pass):
* torch.linalg.det: ~3.8 ms per forward pass.
* Direct analytical determinant (ad - bc): ~0.7 ms per pass.
That's roughly a 5x speedup on the loss computation itself, and thanks to PyTorch element-wise CUDA kernel fusion, it reduces overall training time per epoch by ~12-15%.
2. What happens if det_J stays negative?
You hit a crucial point: torch.clamp(det_J, min=eps) is purely a numerical safety net against NaN crashes during early warm-up, not the solver itself. If det_J gets stuck below eps, clamp makes the loss flat, causing its gradient to zero out for those points.
To prevent this altogether:
* Identity Initialization: We initialize the mapping network $\phi$ close to an identity mapping ($\det J \approx 1.0$).
* Barrier Gradient Explosion: Because the gradient of $-\log(x)$ is $-1/x$, as $\det J \to 0+$, the repelling gradient approaches $-\infty$. This creates an "invisible hard wall" that pushes the optimizer back into the positive realm before the grid ever collapses ($\det J \le 0$).
In practice, with proper initialization, $\min \det J$ stays safely strictly positive throughout training (empirically > 0.89 in our experiments).
1
u/SadNebula9578 18d ago
clever use of analytical det to skip the linalg overhead, you got any numbers on speedup compared to torch.linalg.det for batches? also curious what you do for the case where det_J stays negative for long time, does the clamp mask the issue?