# Post-Training Augmentation Invariance

Source: [Post-Training Augmentation Invariance](<https://www.sota2.com/research/paper/68fa78971e209131cea792be>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2026-06-11T05:00:05.512Z

Authors: Keenan Eikenberry, Lizuo Liu, Yoonsang Lee

Published: 2025-05-16T21:11:51.000Z

## Abstract

This work develops a framework for post-training augmentation invariance, in which our goal is to add invariance properties to a pretrained network without altering its behavior on the original, non-augmented input distribution. We define this notion precisely and additionally introduce augmented encoders, which are probabilistic encoders that formalize augmentation-based encoding processes and that serve as our fundamental object of study. We introduce two losses for augmented encoders, namely, Markov-Wasserstein minimization and Wasserstein correlation maximization, and we demonstrate empirically that both losses can be used to train lightweight, one-hidden-layer MLP adapter networks $E\_{\\theta}$ that, when appended to the latent space of a pretrained network $F$, do indeed lead to (approximate) post-training augmentation invariance. For example, on STL10 with $F=\\text{DINO}$ features, the composite network $C\\circ E\_{\\theta}\\circ F$, where $C$ is a linear classifier and where $E\_{\\theta}$ is one of our proposed adapter networks, achieves 94% classification accuracy on arbitrarily rotated images, whereas a network of the form $C\\circ F$ without the adapter $E\_{\\theta}$ drops to 71% accuracy. Similarly, we can boost noise-invariant classification results from 58% up to 86%. Significantly, we obtain these results with no fine-tuning (the weights of $F$ remain frozen throughout), and our methods introduce little corruption to the original features, since $E\_{\\theta}$ acts nearly isometrically on the non-augmented latent distribution. In contrast, we show that adapter networks trained with alternative candidate losses, specifically SimCLR and HSIC maximization, produce uncompetitive classification results and fundamentally corrupt the original latent space. Code available at https://github.com/keenan-eikenberry/augmentation\_invariance

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

- [Classification on STL10 Rotation (test)](<https://www.sota2.com/research/sota/classification-on-stl10-rotation-test.md>): Classification; STL10 Rotation (test); Accuracy (Original Data): 98.62.

- [Classification on STL10 Noise (test)](<https://www.sota2.com/research/sota/classification-on-stl10-noise-test.md>): Classification; STL10 Noise (test); Accuracy (Original): 98.58.

- [Classification on STL10 Affine (test)](<https://www.sota2.com/research/sota/classification-on-stl10-affine-test.md>): Classification; STL10 Affine (test); Accuracy (STL10 Affine Original): 98.65.

- [Classification on STL10 Crop (test)](<https://www.sota2.com/research/sota/classification-on-stl10-crop-test.md>): Classification; STL10 Crop (test); Accuracy (Original Data): 98.66.

- [Classification on STL10 Affine](<https://www.sota2.com/research/sota/classification-on-stl10-affine.md>): Classification; STL10 Affine; Accuracy (Standard): 94.23.

- [Classification on STL10 Crop](<https://www.sota2.com/research/sota/classification-on-stl10-crop.md>): Classification; STL10 Crop; Accuracy (Standard): 94.23.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2505.11702>)

[PDF](<https://arxiv.org/pdf/2505.11702.pdf>)
