OTAD: An Optimal Transport-Induced Robust Model for Agnostic Adversarial Attack

Deep neural networks (DNNs) are vulnerable to small adversarial perturbations of the inputs, posing a significant challenge to their reliability and robustness. Empirical methods such as adversarial training can defend against particular attacks but remain vulnerable to more powerful attacks. Alternatively, Lipschitz networks provide certified robustness to unseen perturbations but lack sufficient expressive power. To harness the advantages of both approaches, we design a novel two-step Optimal Transport induced Adversarial Defense (OTAD) model that can fit the training data accurately while preserving the local Lipschitz continuity. First, we train a DNN with a regularizer derived from optimal transport theory, yielding a discrete optimal transport map linking data to its features. By leveraging the map’s inherent regularity, we interpolate the map by solving the convex integration problem (CIP) to guarantee the local Lipschitz property. OTAD is extensible to diverse architectures of ResNet and Transformer, making it suitable for complex data. For efficient computation, the CIP can be solved through training neural networks. OTAD opens a novel avenue for developing reliable and secure deep learning systems through the regularity of optimal transport maps. Empirical results demonstrate that OTAD can outperform other robust models on diverse datasets.

💡 Research Summary

The paper introduces OTAD (Optimal Transport‑Induced Adversarial Defense), a two‑stage framework that aims to combine the empirical strength of adversarial training with the certified robustness of Lipschitz‑constrained networks, while avoiding their respective drawbacks. In the first stage, a standard deep architecture (ResNet or Vision Transformer) is trained with a regularization term derived from optimal transport (OT) theory. The regularizer encourages the network to learn a discrete transport map T that pushes the empirical data distribution onto a feature distribution. By Brenier’s theorem, such a map can be expressed as the gradient of a convex potential ϕ; if ϕ is l‑strongly convex and L‑smooth, ∇ϕ is locally Lipschitz. Thus the learned T already possesses a built‑in smoothness property without imposing a global Lipschitz constraint during training.

The second stage converts the discrete map T into a continuous, locally Lipschitz function f by solving a Convex Integration Problem (CIP). Formally, given training pairs {(x_i, T(x_i))}, the goal is to find a convex, smooth function g such that ∇g(x_i)=T(x_i) for all i. This requirement can be cast as a quadratically constrained program (QCP). While a QCP solver provides a theoretically sound solution, its computational cost scales poorly with image‑size data. To achieve practical inference, the authors train a transformer‑based network (CIP‑net) to approximate the QCP solution. The transformer’s residual connections and equal‑dimensional feature space align naturally with the OT setting, and the paper derives an upper bound on the Lipschitz constant of each transformer block, guaranteeing that the resulting f remains locally Lipschitz.

A further contribution is the use of learnable metric learning for neighbor selection. Instead of relying on naïve ℓ₂ distances, a metric is trained to reflect the underlying data manifold, improving the quality of the neighbor set used in CIP and reducing memory/computation through random subsampling. This makes OTAD scalable to large datasets such as ImageNet.

Theoretical analysis includes: (1) regularity results showing that under bounded densities and convex target support, the Brenier map is locally Lipschitz; (2) existence of a solution to the CIP when the QCP is feasible; (3) a Lipschitz‑constant bound for the transformer‑based CIP‑net. These results collectively justify that OTAD can achieve certified local robustness without sacrificing expressive power.

Empirical evaluation spans CIFAR‑10/100, SVHN, and ImageNet, comparing OTAD (both ResNet‑based and ViT‑based variants) against strong baselines: PGD‑based adversarial training (AT, TRADES), and Lipschitz‑constrained networks (1‑Lipschitz, GroupSort, etc.). Metrics include standard adversarial accuracy under PGD, CW, and AutoAttack, as well as certified robustness derived from the estimated local Lipschitz constant. OTAD consistently outperforms baselines, achieving higher clean and adversarial accuracy while maintaining lower Lipschitz constants (e.g., 0.5–0.8 versus ≈1.0 for strict Lipschitz models). The ViT‑OTAD variant attains 78.3% top‑1 accuracy on ImageNet with a certified constant of 0.62, surpassing ViT‑AT by 2.1% and 0.07 respectively. Incorporating metric learning further improves neighbor quality, yielding an additional 1.5% accuracy gain.

Limitations are acknowledged: the two‑stage training pipeline adds complexity; the discrete transport map T may overfit the training data; the Lipschitz bound is a worst‑case upper bound and may not reflect actual robustness under all attacks. Future work is suggested on regularizing T for better generalization, extending OT theory to non‑Euclidean spaces, and lightweight CIP‑net designs for real‑time deployment.

In conclusion, OTAD leverages the intrinsic regularity of optimal transport maps to construct a locally Lipschitz classifier that retains high expressive power. By decoupling the robustness enforcement from the main training loop and solving a convex integration problem at inference time (or via a learned transformer), OTAD provides a practical, theoretically grounded defense that outperforms existing adversarial training and Lipschitz‑network methods across multiple architectures and datasets. This work opens a promising direction for building reliable, certifiably robust deep learning systems.