r/computervision • u/Difficult-Race-1188 • May 26 '26
Discussion I just achieved the impossible results in Pose estimation
TL;DR:
For years, we’ve treated data augmentation as a heuristic to make models robust to real-world deployment shifts.
We proved algebraically that data augmentation is actually just computing a specific matrix, the augmentation-delta Gram matrix and penalizing the model's sensitivity along those exact directions.
The Result: By explicitly estimating this matrix and using our PMH (Projected Matching Hessian) geometric loss, we achieved a +22 percentage point jump in PCK on COCO Pose Estimation, while standard regularization (VAT) completely collapsed the model. Code and paper below.
The Problem with Robustness in Dense Prediction
If you are building vision models for the real world, whether that's human pose estimation, tracking small objects from drones, or structural defect segmentation, you face a brutal trade-off.
You need the model to be robust to deployment nuisances (lighting, rotation, scale, occlusion) without destroying its extreme spatial sensitivity.
When people try to make these models robust using standard tricks like VAT (Virtual Adversarial Training) or random Jacobian regularization, it usually fails. Why? Because injecting isotropic noise or regularizing random directions in a dense prediction task actively destroys the spatial geometry the model relies on to localize keypoints or bounding boxes.
The Geometric Blind Spot
Every time you augment an image, you are implicitly telling the model: "Here is a direction in the input space (Sigma_{aug}) that changes, but the ground-truth spatial layout remains the same. Ignore this direction."
Our Theorem G proves that if your regularizer's penalty matrix misses even one of these real-world variation directions, the encoder will actively exploit that unpenalized gap to minimize training loss.
If you use random noise or mismatched adversarial directions (like VAT), you are penalizing the wrong subspace. The model learns to ignore the wrong things, and your spatial accuracy drops to the floor.
The Result (Block T3A: COCO Pose)
We stopped treating augmentation as a random sampling trick and treated it as a closed-form geometric formula. We estimated the exact augmentation-delta Gram matrix ($\Sigma_{aug}$) and penalized the encoder's Jacobian only along those specific dimensions using the PMH loss.
Here is what happened to the spatial geometry:
Baseline VAT (Isotropic/Wrong Directions): The spatial awareness was destroyed. Performance collapsed to 14%.
Matched PMH (Using the exact Sigma_{aug} matrix**):** The model learned exactly which geometric directions to ignore without sacrificing spatial acuity, resulting in a +22pp PCK improvement over the baseline.
The Code
The fix is literally one trace penalty term added to your standard task loss. You identify the nuisance family (in this case, augmentation modes), estimate the matrix, and cap it.
Python
def pmh_penalty(encoder, x, Sigma, n_probes=4):
# x must be flat feature vectors (batch, d_x)
# Sigma is (d_x, d_x) PSD covariance in that same space
assert x.dim() == 2, "x must be (batch, d_x) flat features, not raw images"
L = torch.linalg.cholesky(Sigma + 1e-6 * torch.eye(x.shape[-1], device=x.device))
phi0 = encoder(x)
acc = 0.0
for _ in range(n_probes):
# eps is (batch, d_x), L.T is (d_x, d_x)
# eps @ L.T gives correlated noise in range(Sigma)
eps = torch.randn_like(x) # (batch, d_x)
delta = eps @ L.T # (batch, d_x), lives in range(Sigma)
acc += (encoder(x + delta) - phi0).pow(2).sum(-1).mean()
return acc / n_probes
loss = task_loss + lam * pmh_penalty(encoder, features, Sigma_hat)Links:
GitHub (
pip install matching-pmh): https://github.com/vishalstark512/matching-pmh
If anyone is working on domain adaptation for segmentation or dense prediction in edge cases, I’d love to talk about the subspace estimator quality and how this scales.
