📊 statistics

Post-Training Augmentation Invariance

This paper introduces a post-training framework that appends lightweight, trainable adapter networks to frozen pretrained models to achieve robust augmentation invariance (e.g., to rotation and noise) without degrading original performance, utilizing novel Markov-Wasserstein and Wasserstein correlation losses that outperform existing methods like SimCLR and HSIC.

Original authors: Keenan Eikenberry, Lizuo Liu, Yoonsang Lee

Published 2026-04-24

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Keenan Eikenberry, Lizuo Liu, Yoonsang Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, world-class art critic named F. This critic has spent years studying millions of paintings and has developed an incredible ability to recognize objects, styles, and details. However, there's a catch: F is a bit rigid. If you show them a painting of a cat, they recognize it perfectly. But if you show them the same painting rotated 90 degrees, or covered in a little bit of static noise, F gets confused and might think it's a different object entirely.

In the world of AI, F is a pre-trained model (like DINOv2 or CLIP). It's already trained and "frozen" (its brain can't be changed without a massive, expensive retraining process).

The problem is: How do we make this rigid critic flexible enough to handle rotated or noisy images without breaking their original, perfect ability to recognize the "normal" images?

This paper proposes a clever solution: Don't retrain the critic. Just give them a pair of smart glasses.

The Core Idea: The "Adapter" Glasses

Instead of trying to retrain the giant, expensive model F (which is like trying to re-teach a PhD professor basic math), the authors attach a tiny, lightweight add-on network called an Adapter (let's call it E).

Think of E as a pair of smart glasses that F wears.

The Input: An image comes in.
The Glasses (E): The image passes through the glasses first. If the image is rotated, the glasses subtly "pre-process" it so it looks "normal" to the critic. If it's noisy, the glasses filter the noise.
The Critic (F): The critic sees the processed image and gives their usual, perfect judgment.
The Result: The system works on both normal and weird images, and the critic's original knowledge remains intact.

The Challenge: Don't Break the Glasses!

Here is the tricky part. If you just train the glasses to fix rotated images, you might accidentally distort the image so much that the critic can no longer recognize the original un-rotated images. You don't want the glasses to turn a normal cat into a dog just to make it look "upright."

The authors needed a way to train these glasses to be invariant (ignoring the rotation/noise) while remaining isometric (preserving the original shape and relationships of the data).

The Two Secret Recipes (Loss Functions)

To train these glasses, the authors invented two new mathematical "recipes" (loss functions) to guide the learning process. They compared these against two popular, existing recipes (SimCLR and HSIC) and found the existing ones failed miserably.

1. The "Anchor" Recipe (Markov-Wasserstein Minimization)

The Analogy: Imagine you are trying to teach a dancer to perform a move regardless of the music tempo.
How it works: You tell the dancer: "When the music is normal, stand exactly where you are (don't move). When the music is fast (augmented), move to the same spot as if the music were normal."
The Magic: This recipe uses a concept called Optimal Transport (think of it as the most efficient way to move furniture from one room to another). It forces the "rotated" version of the image to be mapped to the exact same location in the critic's mind as the "normal" version, but it strictly forbids the glasses from scrambling the positions of the "normal" images.
Result: The critic sees the rotated cat as the same cat, and the original cat is still recognized perfectly.

2. The "Correlation" Recipe (Wasserstein Correlation Maximization)

The Analogy: Imagine you have a map of a city. You want to make sure that no matter how you rotate the map, the distance between two landmarks (like the library and the park) stays the same.
How it works: This recipe tries to maximize the "statistical connection" between the input and the output while ensuring the structure of the map (the distances between points) is preserved. It's like saying, "Make sure the rotated image and the original image are best friends in the critic's mind, but don't let them forget who their other friends are."
Bonus: This recipe is so good at preserving structure that it can also shrink the data (dimensionality reduction). It can take a huge, complex image and compress it into a smaller, simpler version without losing the important details.

What Happened When They Tested It?

The authors tested these "glasses" on several famous AI models (DINOv2, CLIP, SwAV) using the STL10 dataset (a collection of animal and object photos).

The "Before" Scenario: Without the glasses, if you rotated an image, the AI's accuracy dropped from 98% to 71%. It was confused.
The "After" Scenario: With the new "Anchor" glasses (MaWa), the AI handled rotated images with 94% accuracy, while still keeping 98% accuracy on normal images.
The "Noise" Scenario: When they added static noise to the images, the AI without glasses dropped to 58% accuracy. With the glasses, it jumped to 86%.

The "Bad" Recipes (Why SimCLR and HSIC Failed)

The authors tried using two other popular methods (SimCLR and HSIC) to train the glasses.

The Analogy: Imagine trying to fix the glasses by telling the dancer to "group all cats together and all dogs together."
The Result: While this made the dancer good at grouping, it completely scrambled the map. The glasses distorted the original image so badly that the critic could no longer recognize the original, un-rotated images. The "structure" of the data was destroyed.
The Lesson: You can't just use standard "grouping" techniques for this job; you need methods that respect the geometry of the original data.

The One Exception: The "Supervised" Critic

There was one case where the "Anchor" recipe failed: when the critic was a Supervised ResNet50 (a model trained by humans with labels).

Why? This model had already learned to be very sensitive to specific details. The "rotated" version of an image was so different from the "original" in this model's mind that the glasses couldn't find a way to map them together without breaking the original.
The Takeaway: The method works best on models that learned features on their own (Self-Supervised) rather than models that were heavily guided by human labels.

Summary

This paper is like a universal adapter for AI models.

Problem: Pre-trained AI models are great but brittle; they fail when images are rotated or noisy.
Solution: Don't retrain the giant model. Just add a tiny, lightweight "adapter" layer.
Innovation: They invented two new mathematical rules (based on moving data efficiently) to train this adapter so it fixes the image without distorting the model's original knowledge.
Result: AI models become robust against rotations and noise, staying accurate on weird inputs while remaining perfect on normal ones, all without the cost of retraining the massive model.

It's a way of making AI "flexible" without breaking its "muscle memory."

1. Problem Statement

The paper addresses a critical gap in modern machine learning: Post-Training Augmentation Invariance.

Context: Large-scale pretrained models (e.g., DINOv2, CLIP, SwAV) provide powerful general-purpose features but are not inherently invariant to specific augmentations (e.g., arbitrary rotations, noise, affine transformations) that may be crucial for downstream tasks.
The Challenge: Retraining foundation models from scratch is computationally prohibitive. Standard fine-tuning often distorts the pretrained latent space, degrading performance on the original, non-augmented distribution.
Goal: Develop a framework to add invariance properties to a frozen pretrained network $F$ without corrupting its existing capabilities or the geometric structure of its latent space. The solution must be lightweight and applicable post-training.

2. Methodology

The authors propose a framework based on Augmented Encoders and two novel loss functions derived from Optimal Transport theory.

A. Core Concepts

Augmented Encoders ( $E_T$ ): Instead of a deterministic encoder, the authors formalize the process as a probabilistic encoder that maps an input $x$ to a distribution over latent codes. This distribution is formed by encoding multiple augmented views of $x$ .
Invariance Definition: They define $(t, \mu_X, F, \mathcal{V})$ -invariance. An adapter network $E_\theta$ $E_{θ}$ appended to a frozen encoder $F$ $F$ is invariant if:
1. The composite network $E_\theta \circ F$ produces the same latent representation for an input and its augmented versions (invariance).
2. The adapter preserves the metric structure of the original latent distribution $F_\sharp \mu_X$ up to a map $V$ in an admissible class (e.g., isometries). This ensures the original features remain usable for other tasks.

B. Proposed Loss Functions

The paper introduces two training objectives for the lightweight adapter network $E_\theta$ (a one-hidden-layer MLP):

Markov-Wasserstein (MaWa) Minimization:
- Mechanism: This is an anchored Mean Squared Error (MSE) loss. It treats the encoder as a Markov-Wasserstein kernel.
- Objective: It minimizes the distance between the distribution of encoded augmented views and the distribution of clean (non-augmented) encoded views, while anchoring the clean view to the identity.
- Formula:
  $\mathcal{L}_{MaWa} = \frac{1}{N} \sum_{i=1}^N \left( \frac{1}{s+1} \| (E_\theta \circ F)(x_i) - F(x_i) \|^2 + \frac{1}{s+1} \sum_{k=1}^s \| (E_\theta \circ F)(t(a_i^k, x_i)) - F(x_i) \|^2 \right)$
- Key Feature: The "anchoring" term prevents the collapse of the latent space (a common failure mode in naive invariance losses) by forcing the network to preserve the original feature $F(x_i)$ for non-augmented inputs.
Wasserstein Correlation (WaCo) Maximization:
- Mechanism: This loss maximizes the Wasserstein Correlation between the input distribution and the encoded augmented distribution.
- Objective: It encourages the joint distribution of (input, encoded-augmented) to be as far as possible from the product distribution (independence), subject to normalization. This forces the encoder to map augmented views of the same input to the same point in the latent space.
- Key Feature: Unlike MaWa, WaCo does not strictly require a frozen network $F$ to anchor against and allows for dimensionality reduction while maintaining invariance. It acts as an approximate local isometry.

C. Baselines for Comparison

The authors compare their methods against:

SimCLR: A standard contrastive learning loss.
HSIC Maximization: A kernel-based independence criterion often used in representation learning.
End-to-End (EC): Training a classifier directly on augmented data without an adapter.

3. Key Contributions

Formal Framework: The first precise mathematical definition of post-training augmentation invariance that explicitly requires the preservation of the pretrained latent space's metric structure.
Augmented Encoders: Introduction of probabilistic encoders as the fundamental object of study for augmentation-based learning.
Novel Losses: Proposal of MaWa minimization and WaCo maximization, which leverage Optimal Transport to achieve invariance without structural corruption.
Empirical Validation: Demonstration that simple, one-hidden-layer MLP adapters can achieve near-perfect invariance on frozen networks, outperforming complex fine-tuning or contrastive approaches.

4. Experimental Results

Experiments were conducted on STL10, TinyImageNet, and various pretrained backbones (DINOv2, SwAV, CLIP, ResNet50).

Performance on Invariance:
- Rotation: On STL10 with DINOv2, the baseline ( $C \circ F$ ) dropped to 71% accuracy on rotated images. The MaWa adapter ( $C \circ E_\theta \circ F$ ) achieved 94%.
- Noise: Baseline accuracy dropped to 58%; MaWa boosted it to 86%.
- Generalization: The method worked consistently across different backbones (Transformers and ResNets) and augmentation types (Rotation, Affine, Noise, Crop).
Structure Preservation (Crucial Finding):
- MaWa & WaCo: These methods preserved the original latent space geometry. For MaWa on DINOv2, the relationship between input and output distances was nearly isometric ( $R^2 = 0.96$ , NRMSD = 0.02).
- SimCLR & HSIC: These methods corrupted the latent space. They achieved some invariance but destroyed the metric structure ( $R^2 \approx 0.04$ ), making the features unusable for downstream tasks that rely on the original geometry.
Failure Cases:
- On the supervised ResNet50, MaWa failed due to high "collision rates" (where augmented features of different classes overlap in the latent space). However, WaCo remained robust in this regime, suggesting it is more flexible for high-collision scenarios.
Dimensionality Reduction: WaCo successfully reduced the dimension of DINOv2 features (384 $\to$ 96) while maintaining invariance and competitive accuracy.

5. Significance and Conclusion

Efficiency: The method requires no fine-tuning of the massive pretrained model. Only a tiny adapter (one hidden layer) is trained, making it computationally cheap and scalable.
Safety: Unlike fine-tuning or contrastive learning, which can "break" the pretrained features, this approach guarantees that the original capabilities of the model are preserved (structure preservation).
Versatility: It enables the use of powerful foundation models in scenarios requiring specific invariances (e.g., medical imaging requiring rotation invariance, satellite imagery requiring affine invariance) without retraining the base model.
Theoretical Insight: The paper highlights that Optimal Transport-based objectives (specifically those preserving metric structure) are uniquely suited for post-training adaptation, whereas mutual information or contrastive methods are better suited for learning representations from scratch but fail at preserving existing structure.

In summary, the paper provides a mathematically rigorous and empirically validated solution to "fix" pretrained models for specific augmentation needs without sacrificing their general utility, using lightweight adapters trained on novel Optimal Transport losses.