WARP: Guaranteed Inner-Layer Repair of NLP Transformers

Hsin-Ling Hsu; Min-Yu Chen; Nai-Chia Chen; Yan-Ru Chen; Yi-Ling Chang; Fang Yu

cs.LGApr 1, 2026

WARP: Guaranteed Inner-Layer Repair of NLP Transformers

Authors: Hsin-Ling Hsu, Min-Yu Chen, Nai-Chia Chen, Yan-Ru Chen, Yi-Ling Chang, Fang Yu

Organizations: Department of Management Information Systems, National Chengchi University

Abstract

Transformer-based NLP models remain vulnerable to adversarial perturbations, yet existing repair methods face a fundamental trade-off: gradient-based approaches offer flexibility but lack verifiability and often overfit; methods that do provide repair guarantees are restricted to the final layer or small networks, significantly limiting the parameter search space available for repair. We present WARP (Weight-Adjusted Repair with Provability), a constraint-based repair framework that extends repair beyond the last layer of Transformer models. WARP formulates repair as a convex quadratic program derived from a first-order linearization of the logit gap, enabling tractable optimization over a high-dimensional parameter space. Under the condition that the first-order approximation holds, this formulation induces three per-sample guarantees: (i) a positive margin constraint ensuring correct classification on repaired inputs, (ii) preservation constraints over a designated remain set, and (iii) a certified robustness radius derived from Lipschitz continuity. To ensure feasibility across varying model architectures, we introduce a sensitivity-based preprocessing step that conditions the optimization landscape accordingly. We further show that the iterative optimization procedure converges to solutions satisfying all repair constraints under mild assumptions. Empirical evaluation on encoder-only Transformers with varying layer architectures validates that these guarantees hold in practice while improving robustness to adversarial inputs. Our results demonstrate that guaranteed, generalizable Transformer repair is achievable through principled constraint-based optimization.

Figures & tables

Symbol	Shape	Description	Symbol	Shape	Description
$v(x)$	$\mathbb{R}^{d_{\text{in}}}$	Input embedding to repair layer	$U^{(t)}$	$\mathbb{R}^{d_{\text{out}}\times r}$	Top- $r$ left singular vectors of $W^{(t)}$
$W$	$\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$	Repair-layer weight (only modified param.)	$B$	$\mathbb{R}^{r\times d_{\text{in}}}$	Coefficient matrix for low-rank update
$b$	$\mathbb{R}^{d_{\text{out}}}$	Repair-layer bias (fixed)	$\beta$	$\mathbb{R}^{rd_{\text{in}}}$	Vectorisation of $B$ ; QP decision variable
$W_{c}$	$\mathbb{R}^{C\times d_{\text{out}}}$	Classification head (fixed)	$\kappa(v)$	scalar	Gap Sensitivity Norm (GSN)
$b_{c}$	$\mathbb{R}^{C}$	Classification bias (fixed)	$\gamma_{s},\gamma_{h}$	scalar	Repair / remain gap margins
$\widetilde{W}$	$\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$	Final converged repair-layer weight	$\rho,\lambda$	scalar	Regularisation / slack-penalty coeff.

Table 1: Principal notation used in § 3 .

Task	Model	$\|\mathcal{S}\|$	$\|\mathcal{P}\|$	$\|\text{Gen}\|$	$\|\text{AtkGen}\|$
SST-2	DistilBERT	108	800 †	2,000 †	1,207
SST-2	BERT	105	800 †	2,000 †	1,207
RTE	DistilBERT	47	200 ‡	500 ‡	222
RTE	BERT	50	200 ‡	500 ‡	222

Table 2: Dataset configurations.

|\mathcal{S}|

: repair set;

|\mathcal{P}|

: remain set;

|\text{Gen}|

: general accuracy evaluation set;

|\text{AtkGen}|

: attack generalization set.

\dagger

Yelp Polarity test split ( Zhang et al., 2015 ) .

\ddagger

RTE training split.

Task	Model	Repair Layer	Activation Function	Type	Init. $\kappa$
SST-2	DistilBERT	pre_classifier ( $768\!\times\!768$ )	ReLU	A	5.43
SST-2	BERT	bert.pooler.dense ( $768\!\times\!768$ )	tanh	B	0.58
RTE	DistilBERT	pre_classifier ( $768\!\times\!768$ )	ReLU	A	4.79
RTE	BERT	bert.pooler.dense ( $768\!\times\!768$ )	tanh	B	0.52

Table 3: Model configurations. GSN

\kappa

is measured at initialisation (pre-GSN-FT). Type A repair layers are fully updated during downstream fine-tuning (high initial

\kappa

); Type B receive minimal task-specific gradient updates (low initial

\kappa

SST-2 (sentiment classification)
Model	Method	Repair Acc $\uparrow$	Remain Acc $\uparrow$	General Acc $\uparrow$	Atk. Gen $\uparrow$
DistilBERT	Original	5.6%	98.1%	98.2%	22.3%
	Same-layer LoRA ( $r\!=\!2$ )	75.9%	99.0%	97.5%	56.7%
	All-layers Full FT	100.0%	99.5%	93.6%	50.2%
	WARP (ours)	100.0%	100.0%	91.4%	69.0%
BERT	Original	1.0%	98.2%	98.5%	23.2%

Table 4: Main experimental results on SST-2 and RTE. “Repair Acc” is the accuracy on the repair set

\mathcal{S}

; “Remain Acc” on the remain set

\mathcal{P}

; “General Acc” on the held-out evaluation set; “Atk. Gen” is the mean accuracy across six adversarial attack methods (higher is better for all metrics). All sets are defined with respect to the post-GSN-FT weights; see § 4.1 .

		SST-2		RTE
		DistilBERT	BERT	DistilBERT	BERT
Cert. 1	Samples with gap $\geq\gamma_{s}$	108/108	105/105	47/47	50/50
Cert. 1	gap min / mean / max	1.000/1.291/2.360	1.000/1.431/2.814	1.000/1.244/2.662	1.000/1.866/4.843
Cert. 2	Samples with remain gap $\geq\gamma_{h}$	798/800	786/800	200/200	200/200
Cert. 2	remain gap min / mean	0.300/1.059	0.300/1.915	0.300/0.946	0.300/0.951
Cert. 3	$\varepsilon^{*}$ min/median/mean/max	0.172/0.211/0.222/0.405	0.014/0.016/0.020/0.039	0.283/0.305/0.352/0.753	0.022/0.028/0.041/0.106

Table 5: Repair certificate verification (WARP, after QP repair).

	SST-2				RTE
	DistilBERT		BERT		DistilBERT		BERT
	Pre	Post	Pre	Post	Pre	Post	Pre	Post
GSN $\kappa$	5.43	7.43	0.58	1.31	4.79	8.48	0.52	1.60
Atk. Gen	22.3%	20.9%	23.2%	22.5%	37.4%	37.7%	39.3%	42.8%
Best step	—	4	—	20	—	23	—	27
$\\|W\\|_{2}$	1.506	2.238	13.408	13.423	1.140	2.949	13.413	13.558

Table 6: Effect of GSN-FT. Attack generalization is reported before and after GSN-FT to confirm that generalization gains originate from the QP step rather than preprocessing.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Description	Value
$r$	Rank of weight perturbation	2
$\gamma_{s}$	Repair set gap margin target	1.0
$\gamma_{h}$	Remain set gap margin lower bound	0.3
$\lambda$	Slack penalty coefficient	50.0
$\rho$	regularization coefficient	2.0
$T$	Maximum iterations	300

Appendix

Table 7: QP repair hyperparameters.

Parameter	Value
Max steps $T_{\mathrm{FT}}$	30
Learning rate	$1\times 10^{-3}$
Batch size	32
Auxiliary set size $N_{\mathrm{aux}}$	8
Trainable parameters	repair layer $W$ + classification head $W_{c}$ (592,130 total)
Checkpoint selection criterion	step with maximum GSN $\kappa$

Appendix

Table 8: GSN-FT hyperparameters.

	SST-2				RTE
	DistilBERT		BERT		DistilBERT		BERT
Multiplier $k$	Flipped	Cum.%	Flipped	Cum.%	Flipped	Cum.%	Flipped	Cum.%
$\leq 1\times$	0/108	0.0%	0/105	0.0%	0/47	0.0%	0/50	0.0%
$\leq 20\times$	37/108	34.3%	0/105	0.0%	13/47	27.7%	0/50	0.0%
$\leq 25\times$	78/108	72.2%	0/105	0.0%	34/47	72.3%	0/50	0.0%
$\leq 35\times$	107/108	99.1%	0/105	0.0%	47/47	100.0%	0/50	0.0%

Appendix

Table 9: Expanding-radius Monte Carlo stress test (tightness ratio: smallest multiplier

k

such that a flip is observed within

k\times\varepsilon^{*}

Distance band ( $d/\varepsilon^{*}_{\mathrm{nearest}}$ )	Samples	Accuracy
$\leq 20\times$	5	100.0%
$20\times$ – $25\times$	49	87.8%
$25\times$ – $30\times$	180	82.2%
$30\times$ – $35\times$	267	74.5%
$35\times$ – $40\times$	271	68.6%
$40\times$ – $50\times$	324	60.5%

Appendix

Table 10: Accuracy of the repaired DistilBERT model (SST-2) on adversarial samples grouped by distance to the nearest repaired embedding as a multiple of

\varepsilon^{*}

Task	Model	Method	CheckList	SememePSO	StressTest	T3	TextBugger	TextFooler	Mean
SST-2	DistilBERT	Original	8.7%	34.3%	6.4%	24.2%	29.4%	30.9%	22.3%
		Same-layer LoRA ( $r\!=\!2$ )	41.0%	50.9%	89.1%	46.4%	55.8%	56.9%	56.7%
		All-layers Full FT	59.3%	37.7%	73.7%	38.6%	46.2%	45.6%	50.2%
		WARP	70.5%	60.6%	92.9%	64.1%	59.9%	66.2%	69.0%
	BERT	Original	8.7%	34.3%	11.5%	26.8%	28.4%	29.4%	23.2%
		Same-layer LoRA ( $r\!=\!2$ )	45.7%	59.4%	80.8%	62.1%	53.8%	58.3%	60.0%

Appendix

Table 11: Per-attack-method generalization accuracy.

	DistilBERT				BERT
Rank	Atk. Gen	Gen. Acc	Conv	Time	Atk. Gen	Gen. Acc	Conv	Time
SST-2
1	68.0%	91.8%	4	34s	61.7%	95.1%	7	68s
2	69.0%	91.4%	5	33s	64.7%	92.7%	6	77s
8	68.1%	92.0%	5	160s	59.2%	94.0%	4	108s
32	68.7%	92.1%	8	989s	60.6%	93.9%	4	535s

Appendix

Table 12: Rank ablation. SST-2:

|\mathcal{S}|=108/105

|\mathcal{P}|=800

. RTE:

|\mathcal{S}|=47/50

|\mathcal{P}|=200

. “Conv” denotes QP convergence iterations. ∗ QP infeasible at rank 1.

	DistilBERT				BERT
$\|\mathcal{S}\|$	SST-2		RTE		SST-2		RTE
	Atk. Gen	Gen. Acc	Atk. Gen	Gen. Acc	Atk. Gen	Gen. Acc	Atk. Gen	Gen. Acc
10	33.2%	97.9%	36.4%	94.2%	29.2%	96.2%	38.2%	97.8%
25	51.0%	96.2%	39.2%	87.0%	42.0%	95.2%	45.9%	90.6%
50	60.5%	93.2%	—	—	60.6%	93.8%	—	—
75	65.6%	92.5%	—	—	62.5%	93.0%	—	—

Appendix

Table 13: Repair set size ablation (

|\mathcal{P}|

fixed at 800 for SST-2, 200 for RTE).

	DistilBERT (SST-2)			BERT (SST-2)		DistilBERT (RTE)		BERT (RTE)
$\|\mathcal{P}\|$	Atk. Gen	Gen. Acc	$\varepsilon^{*}_{\mathrm{med}}$	Atk. Gen	Gen. Acc	Atk. Gen	Gen. Acc	Atk. Gen
0	79.5%	0.8%	1.057	76.4%	1.1%	56.5%	15.0%	58.3%
50	76.9%	52.9%	0.305	73.1%	50.7%	41.3%	61.8%	53.1%
200	72.0%	79.8%	0.255	70.7%	80.7%	44.1%	83.8%	47.9%
400	70.3%	86.5%	0.226	69.0%	89.0%	—	—	—
800	69.0%	91.4%	0.211	64.7%	92.7%	—	—	—

Appendix

Table 14: Remain set size ablation (

|\mathcal{S}|

fixed at full size).

\varepsilon^{*}_{\mathrm{med}}

reported for DistilBERT only. Bold denotes the default setting.

Adversarial Robustness of Language Models

Adversarial robustness of language models studies how to keep their outputs and decisions reliable under manipulative prompts, input changes, and malicious fine-tuning.

See all 191 papers

A Constrained Optimization Perspective of Unrolled Transformers
Jan 24, 2026Javier Porras-Valenzuela, Samar Hadou, Alejandro RibeiroDeep Learning Optimization Transformer
Department of Electrical and Systems Engineering, University of Pennsylvania, Philadelphia, PA, United States of America.
RepTran: Search-Based Repair of Transformer Models
Jul 13, 2026Yuta Ishimoto, Paolo Arcaini, Fuyuki Ishikawa +3Transformer
The University of Osaka, Japan · National Institute of Informatics, Japan · Kyushu University, Japan +1
At the Edge of Understanding: Sparse Autoencoders Trace The Limits of Transformer Generalization
Jun 24, 2026Praneet Suresh, Jack Stanley, Sonia Joseph +2Distribution Shift Robustness Transformer Interpretability
Mila Quebec AI Institute · Meta AI · Sapient Intelligence

Jan 24, 2026J/K move · Enter open · S save

A Constrained Optimization Perspective of Unrolled Transformers

Javier Porras-Valenzuela, Samar Hadou, Alejandro Ribeiro

Department of Electrical and Systems Engineering, University of Pennsylvania, Philadelphia, PA, United States of America.

Deep Learning Optimization Transformer OOD Generalization Constrained Optimization

Open paper AskPDF arXiv

We introduce a constrained optimization framework for training transformers that behave like optimization descent algorithms. Specifically, we enforce layerwise descent constraints on the objective function and replace standard empirical risk minimization (ERM) with a primal-dual training scheme. This approach yields models whose intermediate representations decrease the loss monotonically in expectation across layers. We apply our method to both unrolled transformer architectures and conventional pretrained transformers on tasks of video denoising and text classification. Across these settings, we observe constrained transformers achieve stronger robustness to perturbations and maintain higher out-of-distribution generalization, while preserving in-distribution performance.

Abstract

Adversarial Robustness of Language Models

Explore similar work

A Constrained Optimization Perspective of Unrolled Transformers

RepTran: Search-Based Repair of Transformer Models

At the Edge of Understanding: Sparse Autoencoders Trace The Limits of Transformer Generalization