Post-Training Augmentation Invariance

Keenan Eikenberry; Lizuo Liu; Yoonsang Lee

arxiv: 2505.11702 · v2 · submitted 2025-05-16 · 💻 cs.LG · stat.ML

Post-Training Augmentation Invariance

Keenan Eikenberry , Lizuo Liu , Yoonsang Lee This is my paper

Pith reviewed 2026-05-22 13:59 UTC · model grok-4.3

classification 💻 cs.LG stat.ML

keywords post-training augmentation invarianceadapter networksWasserstein distanceMarkov-Wasserstein minimizationWasserstein correlation maximizationaugmented encodersimage classification robustness

0 comments

The pith

Lightweight adapter networks trained with Wasserstein-based losses can add approximate invariance to augmentations in a frozen pretrained network while preserving its original behavior.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper develops a framework for post-training augmentation invariance that adds desired invariance properties to a pretrained network without altering its behavior on the original non-augmented input distribution. It introduces augmented encoders as probabilistic objects that formalize augmentation processes and proposes two losses, Markov-Wasserstein minimization and Wasserstein correlation maximization, to train small one-hidden-layer MLP adapters. These adapters are appended to the latent space of a frozen pretrained network and produce composite models that achieve high accuracy on rotated or noisy test images. The approach matters because it offers a way to improve robustness without the computational cost of retraining or the risk of degrading original features.

Core claim

By training one-hidden-layer MLP adapter networks E_theta using either Markov-Wasserstein minimization or Wasserstein correlation maximization, the composite network C o E_theta o F achieves 94% classification accuracy on arbitrarily rotated STL10 images (versus 71% without the adapter) and 86% on noisy images (versus 58%), with F frozen throughout and E_theta acting nearly isometrically on the non-augmented latent distribution.

What carries the argument

Augmented encoders, defined as probabilistic encoders that formalize augmentation-based encoding processes, optimized via Markov-Wasserstein minimization or Wasserstein correlation maximization to produce nearly isometric adapters on the original latent distribution.

If this is right

The adapter can be trained and appended after the pretrained network is fixed, requiring no gradient updates to the original weights.
The original latent features remain largely uncorrupted, so the same downstream classifiers or heads can continue to be used.
Both proposed losses produce better results than SimCLR or HSIC maximization, which corrupt the original latent space and yield lower classification accuracy.
The same adapter architecture works for different augmentation types including arbitrary rotations and additive noise.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The probabilistic definition of augmented encoders could be used to derive quantitative bounds on how much invariance is achieved for a given loss value.
Multiple such adapters trained for different augmentations might be composed without interfering with each other or with the original features.
The method could be tested on regression or detection tasks to check whether the same invariance gains appear outside classification.
Similar post-training adapters might be trained for other target properties such as equivariance to group actions or robustness to specific distribution shifts.

Load-bearing premise

The proposed losses enforce invariance on augmented inputs while keeping the adapter nearly isometric on the non-augmented latent distribution of the pretrained network.

What would settle it

Training an adapter with one of the proposed losses but then observing that accuracy on the original non-augmented test images drops substantially below the no-adapter baseline or that accuracy on the augmented test images fails to rise above the no-adapter baseline.

read the original 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 = DINOv2 features, the composite network C o E_theta o 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 o 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

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

The paper shows a practical adapter trick that boosts DINOv2 accuracy on rotated and noisy STL10 from 71% to 94% and 58% to 86% while keeping the base model frozen, using two new Wasserstein losses on augmented encoders.

read the letter

The core contribution is a post-training setup that trains a one-hidden-layer MLP adapter on top of a frozen pretrained encoder to enforce invariance under augmentations. They formalize this with augmented encoders and introduce Markov-Wasserstein minimization plus Wasserstein correlation maximization as the training objectives. On STL10 with DINOv2 features the adapter lifts classification accuracy substantially on rotated and noisy test images, and it does so without the original-feature corruption seen in SimCLR and HSIC baselines. Code is released, which helps reproducibility.

Referee Report

2 major / 2 minor

Summary. The paper develops a framework for post-training augmentation invariance in pretrained networks F by introducing augmented encoders and two losses (Markov-Wasserstein minimization and Wasserstein correlation maximization) to train lightweight one-hidden-layer MLP adapters E_θ. These adapters are appended to the latent space of F (e.g., DINOv2) with F frozen, yielding composite networks that achieve 94% accuracy on rotated STL10 images (vs. 71% without adapter) and 86% on noisy images (vs. 58%), while claiming to preserve behavior on the original non-augmented distribution via near-isometry of E_θ.

Significance. If the empirical gains hold under stricter controls, the approach offers a practical post-training method for adding augmentation robustness to frozen pretrained models without full retraining, which could be useful in computer vision pipelines. The code release aids reproducibility, and the direct empirical measurements on held-out augmented data avoid circularity with fitted parameters. However, the significance is tempered by the need for stronger validation of the isometry property central to the 'little corruption' claim.

major comments (2)

[Empirical results] Empirical results section (STL10 experiments with DINOv2): The assertion that E_θ 'acts nearly isometrically on the non-augmented latent distribution' and introduces 'little corruption to the original features' is supported only indirectly via downstream classification accuracy on clean test data; no direct metric such as Wasserstein distance between the original latent pushforward and the original marginal, or pairwise feature distortion statistics, is reported. This measurement is load-bearing for the central post-training invariance definition.
[Loss definitions] Loss definitions and training procedure: The losses are defined via probabilistic augmented encoders that couple original and augmented latents, yet no analysis or bound is provided showing that joint optimization cannot produce non-negligible distortion on the original marginal while still improving augmented accuracy; the isometry is presented as an observed outcome without sensitivity checks to training details or hyperparameters.

minor comments (2)

[Experimental setup] The reported accuracy figures lack error bars, number of independent runs, or statistical tests, which would strengthen assessment of the gains over baselines like SimCLR and HSIC.
Notation for augmented encoders and the precise implementation of the two proposed losses would benefit from an additional diagram or pseudocode to improve clarity for readers.

Simulated Author's Rebuttal

2 responses · 1 unresolved

We thank the referee for their constructive comments on our work. We address the major comments below and will incorporate revisions to strengthen the empirical validation of our claims.

read point-by-point responses

Referee: [Empirical results] Empirical results section (STL10 experiments with DINOv2): The assertion that E_θ 'acts nearly isometrically on the non-augmented latent distribution' and introduces 'little corruption to the original features' is supported only indirectly via downstream classification accuracy on clean test data; no direct metric such as Wasserstein distance between the original latent pushforward and the original marginal, or pairwise feature distortion statistics, is reported. This measurement is load-bearing for the central post-training invariance definition.

Authors: We agree with the referee that direct metrics would provide stronger support for the near-isometry property. In the revised manuscript, we will add computations of the Wasserstein distance between the pushforward measure of the original latents through E_θ and the original marginal distribution. Additionally, we will report average pairwise Euclidean distances or relative norm changes for clean samples to quantify any distortion introduced by the adapter. These additions will directly validate the 'little corruption' claim central to post-training augmentation invariance. revision: yes
Referee: [Loss definitions] Loss definitions and training procedure: The losses are defined via probabilistic augmented encoders that couple original and augmented latents, yet no analysis or bound is provided showing that joint optimization cannot produce non-negligible distortion on the original marginal while still improving augmented accuracy; the isometry is presented as an observed outcome without sensitivity checks to training details or hyperparameters.

Authors: While a rigorous theoretical bound on the possible distortion would be valuable, developing such a bound is non-trivial given the nature of the Wasserstein-based objectives and may be beyond the scope of this work. Empirically, we demonstrate that our method preserves clean performance in contrast to alternatives like SimCLR and HSIC, which do corrupt the latent space. To address the lack of sensitivity analysis, we will include in the revision additional experiments that vary key hyperparameters (e.g., the trade-off parameters in the losses, learning rates) and report the resulting clean and augmented accuracies, confirming that the near-isometry holds across reasonable training configurations. revision: partial

standing simulated objections not resolved

A theoretical analysis or bound guaranteeing that the joint optimization of the Markov-Wasserstein minimization and Wasserstein correlation maximization losses cannot induce significant distortion on the original marginal distribution.

Circularity Check

0 steps flagged

No significant circularity; claims rest on independent empirical evaluation.

full rationale

The paper defines post-training augmentation invariance and augmented encoders as new objects, introduces Markov-Wasserstein minimization and Wasserstein correlation maximization as losses, then trains one-hidden-layer MLP adapters and measures classification accuracy on held-out rotated and noisy STL10 images. These accuracies are external test-set measurements, not quantities that reduce to the fitted parameters or loss definitions by construction. The claim that the adapter acts nearly isometrically on the original latent distribution is presented as an observed empirical outcome rather than a definitional or self-referential consequence. No self-citations, uniqueness theorems, or ansatzes imported from prior author work appear in the load-bearing steps. The derivation chain is therefore self-contained against external benchmarks.

Axiom & Free-Parameter Ledger

1 free parameters · 1 axioms · 1 invented entities

The central claim rests on the newly defined augmented encoders as probabilistic objects and on the empirical effectiveness of the two Wasserstein losses; the adapter parameters are learned rather than free parameters chosen by hand.

free parameters (1)

Adapter parameters theta
Learned via the proposed losses on augmented data; not preset constants.

axioms (1)

domain assumption Pretrained network F produces latent features that remain useful when an adapter is appended without retraining F.
Invoked by freezing F throughout training and evaluation.

invented entities (1)

Augmented encoders no independent evidence
purpose: Probabilistic encoders that formalize augmentation-based encoding processes as the fundamental object of study.
Newly introduced to define the invariance goal precisely.

pith-pipeline@v0.9.0 · 5836 in / 1411 out tokens · 40990 ms · 2026-05-22T13:59:08.991178+00:00 · methodology

discussion (0)

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

IndisputableMonolith/Cost/FunctionalEquation.lean washburn_uniqueness_aczel unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

Markov-Wasserstein minimization loss … MW2²((Eθ ◦ F)T, δF) … anchored MSE … Wasserstein correlation maximization … WCp(π) = WDp(π) / …
IndisputableMonolith/Foundation/AbsoluteFloorClosure.lean reality_from_one_distinction unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

Definition 3.4 … (t, μX, F, V)-invariant … E ◦ F is (t, μX)-invariant and (E ◦ F)(x) = (V ◦ F)(x) … V admissible class (isometries …)

What do these tags mean?

matches: The paper's claim is directly supported by a theorem in the formal canon.
supports: The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends: The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses: The paper appears to rely on the theorem as machinery.
contradicts: The paper's claim conflicts with a theorem or certificate in the canon.
unclear: Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.