Breaking the Adversarial Robustness-Performance Trade-off in Text Classification via Manifold Purification

Conference: AAAI 2026
arXiv: 2511.07888
Code: To be confirmed
Area: AI Safety/Adversarial Robustness
Keywords: Adversarial Defense, Text Classification, Manifold Correction, Normalizing Flow, Riemannian Geometry, OOD Detection, Geodesic Purification

TL;DR

The paper proposes the Manifold-Correcting Causal Flow (MC²F) framework. It learns the manifold density of clean data embeddings using a Stratified Riemannian Continuous Normalizing Flow (SR-CNF) for adversarial sample detection. It then utilizes a Geodesic Purification Solver to project embeddings detected as adversarial back to the clean manifold along the shortest path. This approach comprehensively outperforms state-of-the-art (SOTA) adversarial robustness on three datasets (SST-2, AGNews, and YELP) without any loss in (and even slightly improving) clean data accuracy.

Background & Motivation

• Background: Pre-trained language models (PLMs, such as BERT) have achieved remarkable success in text classification tasks, but they remain highly vulnerable to adversarial attacks (e.g., TextFooler, BERT-Attack). Tiny, semantically imperceptible text perturbations can lead to completely incorrect model predictions.
• Limitations of Prior Work: Existing defense methods (such as adversarial training, embedding denoising, etc.) face a pervasive robustness-accuracy trade-off, where enhancing adversarial robustness inevitably degrades performance on clean data. This trade-off is unacceptable in safety-critical applications.
• Key Challenge: Adversarial training (AT) forcibly improves robustness via data augmentation but incurs high computational costs and often suffers from "robustness illusions" due to gradient masking. Although purification methods avoid modifying model training, they struggle with limited purification effects due to a lack of precise geometric modeling of the embedding space.
• Key Insight: Empirical analysis reveals that clean and adversarial text embeddings occupy geometrically separable, distinct manifold regions within the BERT embedding space. Consequently, adversarial defense can be reframed from a brute-force training problem into a geometric correction problem.

Core Problem

Can we detect adversarial samples as outliers and project them back to the clean manifold by precisely modeling the manifold structure of clean data embeddings, thereby simultaneously achieving high robustness and zero accuracy loss?

Method

Empirical Foundation: Manifold Separability Hypothesis

The paper first conducts a systematic geometric analysis on the SST-2 dataset:

1. Visual Evidence: Three dimensionality reduction methods (PCA, t-SNE, and UMAP) all demonstrate visible, separate clustering between clean and adversarial embeddings.
2. Statistical Distance Evidence: The MMD, JSD, and Wasserstein distances between clean and adversarial distributions are significantly larger than the distributional distances within the clean data itself.
3. Local Intrinsic Dimension (LID) Evidence: The average LID value of adversarial embeddings (28.20) is significantly higher than that of clean embeddings (23.74), with a p-value close to \(10^{-43}\)—implying that adversarial perturbations systematically push embeddings into regions of higher geometric complexity.

Based on these findings, two hypotheses are established: (1) Manifold Separability—clean and adversarial embeddings are statistically and geometrically separable; (2) Stratified Manifold Structure—the embedding space consists of sub-manifolds with different intrinsic dimensions.

Overall Architecture

MC²F consists of two core modules: (1) SR-CNF for adversarial sample detection; (2) Geodesic Purification Solver for embedding correction. During inference, the system computes the log-likelihood \(\log p(z_{in})\) for the input embedding \(z_{in}\). If it falls below a threshold \(\tau\), purification is triggered; otherwise, it passes through directly.

Key Designs

1. Stratified Riemannian Continuous Normalizing Flow (SR-CNF)

   • Function: Learns the probability density of clean data embeddings \(p_{clean}(z)\), which is used to detect OOD adversarial samples.
   • Mechanism: Instead of assuming a fixed geometry, a Mixture-of-Experts (MoE) network is used to learn a position-dependent Riemannian metric tensor \(G(z) = \sum_{k=1}^{K} \alpha_k(z) E_{\psi_k}(z)\).
   • The gating network \(g_\phi(z)\) outputs weights, and each of the \(K\) expert networks specializes in the local geometry of a specific stratum.
   • Ensuring positive definiteness: Each expert outputs \(L_k(z)L_k(z)^T + \epsilon I\).
   • A CNF is defined on the learned Riemannian manifold, and the log-likelihood is calculated via Riemannian divergence (Equations 3-4).
   • Detection Mechanism: An input embedding is classified as an adversarial sample if \(\log p(z_{in}) < \tau\).
   • Design Motivation: The embedding space is not a single homogeneous manifold but a stratified structure composed of different intrinsic dimensions. The MoE adaptively learns this stratified geometry.

2. Geodesic Purification Solver

   • Function: Projects embeddings detected as adversarial back to the clean manifold along the geodesic (the shortest path on the manifold).
   • Formulation: Minimizes the path energy functional \(\mathcal{L}[\gamma] = \int_0^1 \langle \gamma'(t), \gamma'(t) \rangle_{G(\gamma(t))} dt\).
   • Boundary Conditions: \(\gamma(0) = z_{adv}\) and \(\gamma(1) = z_{corr} \in \mathcal{M}_{clean}\).
   • Solver: Discretizes the path and minimizes the energy functional of the path points using gradient descent. The constraint \(\log p(z_{corr}) \geq \tau\) is implemented via soft penalty.
   • Design Motivation: Instead of arbitrary denoising, the solver finds the geometrically closest clean representation to preserve maximum semantic information.

3. Multi-Objective Training Paradigm

   • Density Estimation Loss \(\mathcal{L}_{NLL}\): The negative log-likelihood of standard normalizing flows, driving the learning of clean data distributions.
   • Topological Regularization \(\mathcal{L}_{topo}\): Based on differentiable persistent homology, this calculates the Wasserstein distance between the persistence diagrams of a batch of clean embeddings and their counterparts in the latent space, ensuring the flow transformation preserves the global topological structure.
   • Causal Semantic Regularization \(\mathcal{L}_{causal}\): Models the purification process as a causal intervention (removing the confounding effect of adversarial perturbations). It uses the Fisher-Rao distance to constrain the classifier's output distribution of the purified embedding to match that of the original clean embedding.
   • Total Loss: \(\mathcal{L}_{total} = \mathcal{L}_{NLL} + \lambda_{topo}\mathcal{L}_{topo} + \lambda_{causal}\mathcal{L}_{causal}\)

Key Experimental Results

Main Results (3 Datasets × 3 Attack Methods)

Dataset	Method	Clean%	BERT-Attack Aua%	TextFooler Aua%	TextBugger Aua%
SST-2	Fine-tune	92.71	3.83	6.10	28.70
SST-2	SD (SOTA)	91.36	36.46	46.30	54.50
SST-2	MC²F	92.71	40.05	52.60	61.50
AGNews	Fine-tune	94.68	4.09	14.70	40.00
AGNews	SD (SOTA)	93.81	38.60	49.30	60.10
AGNews	MC²F	95.13	45.30	53.80	64.30
YELP	Fine-tune	95.19	5.40	5.20	29.60
YELP	SD (SOTA)	93.45	39.61	47.80	55.10
YELP	MC²F	95.26	48.50	54.00	63.20

Ablation Study (AGNews, TextFooler Attack)

Configuration	Clean%	Aua%	#Query
MC²F (Full)	95.13	53.8	561.4
w/o \(\mathcal{L}_{NLL}\)	93.22	32.6	366.7
w/o \(\mathcal{L}_{topo}\)	93.41	32.9	375.4
w/o \(\mathcal{L}_{causal}\)	94.76	48.6	479.1

Key Findings

• Zero accuracy loss with even slight improvement: MC²F achieves a Clean% of 95.13% on AGNews (compared to 94.68% for fine-tuning) and 95.26% on YELP (compared to 95.19% for fine-tuning), completely breaking the robustness-accuracy trade-off.
• Significant increase in attack query budget: When facing BERT-Attack on YELP, MC²F requires 586.4 queries (compared to only 320.7 for SD), indicating that its decision boundary is much harder to explore.
• Topological regularization makes the largest contribution: After removing \(\mathcal{L}_{topo}\), the Aua% plummets from 53.8% to 32.9%, validating that preserving the global topological structure of the manifold is crucial to preventing fragile representations.
• All three losses are indispensable: Removing any of the loss terms leads to a significant drop in robustness and/or accuracy.

Highlights & Insights

• "Detection-Purification" paradigm replacing adversarial training: Instead of modifying the model training process, this approach acts as an input filter in the embedding space during inference. It is decoupled from any downstream model, offering high generalizability.
• A complete workflow from empirical evidence to method design: The framework validates the "manifold separability" hypothesis from multiple perspectives (PCA, t-SNE, UMAP, MMD, JSD, LID) before designing the method accordingly, rather than designing a method first and searching for experimental support later.
• MoE for stratified Riemannian geometry: Utilizing a Mixture-of-Experts allows the framework to adaptively capture the non-homogeneous geometric structures of the embedding space, which is far more flexible than fixed metrics.
• Importance of topological regularization: By using persistent homology to constrain the flow transformation to keep topological invariants, its critical role in adversarial robustness is explicitly verified for the first time.

Limitations & Future Work

• The inference phase requires additional density estimation and potential geodesic optimization steps, the computational overhead of which is not reported in detail; this could become a bottleneck for real-time applications.
• The model was only validated on BERT-base; its generalizability to larger models (e.g., RoBERTa-large, LLMs) has not been tested.
• The detection threshold \(\tau\) is determined using the validation set, but in real deployments, the clean and adversarial distributions may suffer from continuous domain shift.
• Experiments only cover word-level attacks (TextFooler, BERT-Attack, TextBugger); sentence-level (paraphrase) or character-level adversarial attacks have not been tested.
• The number of iterations and convergence analysis for the geodesic solver are insufficient.
• The purification process might introduce subtle semantic drift, which in extreme cases could affect the performance of non-adversarial samples.

Related Work & Insights

• vs. Adversarial Training (FreeLB/WLRE): Adversarial training modifies model weights, which typically degrades Clean% by 0.5-1.5%. As a post-processing module, MC²F does not alter the model, resulting in no degradation—and even slight improvements—in Clean%.
• vs. Subspace Defense (SD): SD removes adversarial components via subspace projection, but the projection is linear. MC²F performs geodesic projection via learned non-linear Riemannian metrics, which more accurately conforms to the curved geometry of the embedding space.
• vs. DAD (Zhang et al. 2025): DAD uses MMD for detection and a denoiser for purification. MC²F's detection utilizing Riemannian CNF is more precise (relying on density estimation rather than a two-sample test), and its purification provides geometric optimality guarantees (geodesics vs. heuristic denoising).
• Insights: The logic of "learning data manifolds \(\rightarrow\) OOD detection \(\rightarrow\) geometric projection correction" within this framework can be generalized to adversarial defense in image and multimodal domains.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Unifies stratified Riemannian CNF, geodesic purification, and topological regularization into an adversarial defense framework for the first time, backed by complete theoretical derivations.
• Experimental Thoroughness: ⭐⭐⭐⭐ Tested on 3 datasets × 3 attacks × 4 baselines, complete with ablation studies and preliminary empirical analyses, though it lacks evaluations on larger models and more attack types.
• Writing Quality: ⭐⭐⭐⭐⭐ The logical chain from empirical hypotheses to method design and experimental validation is exceptionally clear, with rigorous mathematical formulations.
• Value: ⭐⭐⭐⭐ Resolves the long-standing pain point of the robustness-accuracy trade-off in text classification, offering high practical utility.

Related Papers

• [AAAI 2026] TopoReformer: Mitigating Adversarial Attacks Using Topological Purification in OCR Models
• [ICML 2026] Position: Embodied AI Requires a Privacy-Utility Trade-off
• [ICLR 2026] Rethinking Pareto Frontier: On the Optimal Trade-offs in Fair Classification
• [ICCV 2025] Failure Cases Are Better Learned But Boundary Says Sorry: Facilitating Smooth Perception Change for Accuracy-Robustness Trade-Off in Adversarial Training
• [AAAI 2026] Breaking the Dyadic Barrier: Rethinking Fairness in Link Prediction Beyond Demographic Parity