RepTran: Search-Based Repair of Transformer Models
RepTran combines neuron-aware localization with Differential Evolution, repairing 74.7% of targeted faults across 18 ViT benchmarks.
Key Findings
Methodology
RepTran repairs Transformer feed-forward networks (FFNs). It computes VDiff and MisAct for intermediate neurons, forming NeuronScore, then modulates Arachne’s Forward Impact and Gradient Loss to obtain Weight Suspiciousness. The top-Nw weights are optimized with Differential Evolution (DE). Its fitness function rewards correct predictions on already-correct inputs while improving the targeted misclassification set, enabling selective repair rather than full retraining.
Key Results
- Across 18 fault benchmarks built from CIFAR-100 and Tiny-ImageNet, RepTran achieved a 74.7% mean repair rate and 95.2% maximum. Arachne applied to ViT FFNs reached only 17.1% on average and 44.4% at best, showing the value of Transformer-aware localization.
- RepTran averaged 476.19 seconds, compared with 620.93 seconds for Arachne. It remained competitive across SRC-TGT, TGT-FP, and TGT-FN faults and their three frequency ranks, rather than helping only one error category.
- Effect-size analysis against ArachneW found higher repair rates for RepTran regardless of the number of selected weights. The underlying ViT reached 91.18% test accuracy on CIFAR-100 and 86.34% on Tiny-ImageNet after two fine-tuning epochs.
Significance
The work addresses a practical reliability problem: deployed Transformers may make systematic, localized mistakes, while full retraining is costly and can damage established behavior. It transfers software-engineering ideas of fault localization and automated patch search to ViTs, emphasizing targeted behavioral correction. The results suggest that structure-aware repair is useful when labeled repair data and compute are limited, especially in safety-sensitive AI-enabled systems.
Technical Contribution
RepTran explicitly injects FFN-neuron behavior into weight localization. VDiff compares covariance-based behavioral diversity between correct and incorrect samples, while MisAct measures mean activation under the fault. Their normalized product modulates Arachne’s FI and GL. DE then searches only selected FFN parameters with a fitness function balancing repair and regression avoidance, yielding an integrated localization-and-search pipeline rather than generic fine-tuning.
Novelty
Most DNN repair methods target CNNs or generic feed-forward networks. Arachne searches weights but does not exploit the knowledge-like role of Transformer FFN neurons. RepTran’s fundamental novelty is therefore structural: it combines neuron-level behavioral statistics with weight-level bidirectional influence and applies the result to semantically concentrated FFN regions. It is a specialized search-based repair strategy, not a new Transformer architecture.
Limitations
- The evaluation covers only ViT image classification, CIFAR-100, Tiny-ImageNet, and the final encoder-block FFN. It does not establish effectiveness for LLMs, language tasks, multi-layer faults, or distribution shifts.
- The method assumes a predefined misclassification type and sufficient representative samples. Rare, heterogeneous, or changing faults may produce unreliable NeuronScores and difficult DE searches.
- DE has computational and hyperparameter costs, while the fitness objective provides empirical rather than formal safety guarantees.
Future Work
Future work should extend repair to multiple layers, LLMs, other Transformer tasks, and automatically discovered fault categories. Combining RepTran with provable methods such as PRoViT could provide both empirical generalization and guarantees. Important directions include online repair, robustness and fairness checks, continual regression testing, and evaluation on industrial failure logs.
AI Executive Summary
Transformers now support foundation models, autonomous-driving perception and medical imaging, yet their failures are often systematic rather than global: a model may repeatedly confuse one class while performing well elsewhere. Full retraining is expensive and may erase useful behavior. Existing DNN repair methods, including Arachne, largely overlook the distinctive role of Transformer feed-forward networks (FFNs).
RepTran treats repair as a two-stage search problem. It measures how FFN intermediate neurons differ between correct and faulty inputs using VDiff and MisAct, combines this NeuronScore with Arachne’s Forward Impact and Gradient Loss, and selects the most suspicious weights. Differential Evolution then searches their values. A candidate is rewarded when it fixes the target errors without breaking already-correct predictions.
On 18 fault benchmarks derived from CIFAR-100 and Tiny-ImageNet, RepTran repaired 74.7% of targeted faults on average and reached 95.2% in the best case. Arachne averaged 17.1% and peaked at 44.4%; RepTran also reduced mean runtime from 620.93 to 476.19 seconds. Its advantage persisted over ArachneW across weight budgets. The evidence is promising but bounded: experiments used ViT image classifiers and only the final FFN block. RepTran is therefore a strong specialized repair framework, not yet a universal or formally guaranteed solution.
Deep Analysis
Background
Transformers combine multi-head self-attention with stacked FFNs and now dominate language and vision systems. DNN repair methods such as Arachne, CARE and PRDNN established selective correction; ROME and MEND edit factual associations. However, these approaches rarely exploit FFN intermediate neurons as semantic or knowledge-bearing units. This paper evaluates ViT and focuses on its final block, where prior work suggests relevant knowledge is concentrated.
Core Problem
The objective is not general accuracy improvement, but reducing a specified misclassification pattern while preserving unrelated correct behavior. The problem is difficult because models contain many parameters, error responsibility is unclear, Transformer components differ from CNNs, and repair data are limited. Random or broad fine-tuning can search inefficiently or introduce regressions.
Innovation
- ��Introduces NeuronScore from covariance-based VDiff and misbehavior activation MisAct.
- ��Uses NeuronScore to modulate Arachne’s Forward Impact and Gradient Loss, producing Weight Suspiciousness.
- ��Applies Differential Evolution only to selected FFN weights, with a repair-versus-preservation fitness.
- ��Demonstrates that Transformer-aware localization outperforms generic weight selection on 18 benchmarks.
Methodology
- ��Input: a fine-tuned ViT, correct set Icor, and target-misclassification set Imis.
- ��FFN: FFN(z)=Waftσ(Wbefz+bbef)+baft; experiments use D=768 and Dhidden=3072.
- ��Neuron localization: VDiff compares absolute covariance-row sums between correct and faulty activations; MisAct is mean activation on Imis; NeuronScore=VDiffnorm×MisActnorm.
- ��Weight localization: for each connected weight w, ModFI=FI×NeuronScore and ModGL=GL×NeuronScore; WeightSusp=pModFI+(1-p)ModGL.
- ��Search: DE initializes a population, mutates and selects candidates for Top-Nw weights. Fitness averages 1 for correct predictions and 1/(1+L) for incorrect ones, while α balances preservation and repair.
Experiments
The study uses CIFAR-100, with 50,000 training and 10,000 test images, and Tiny-ImageNet, with 100,000 training, 10,000 validation and 10,000 test images. An 80:20 train/repair split and a disjoint test set support generalization analysis. Eighteen benchmarks cover SRC-TGT, TGT-FP and TGT-FN faults at ranks 1–3. Baselines are random selection, Arachne, ArachneW, and, for RQ5, PRoViT variants. Nw is 11, 236, 472 or 944; p defaults to 0.5.
Results
RepTran achieved 74.7% mean repair and 95.2% maximum, versus 17.1% mean and 44.4% maximum for Arachne, while runtime averaged 476.19 versus 620.93 seconds. Effect sizes showed RepTran superior to ArachneW across weight budgets. Layer analysis found consistently higher final-layer neuron scores across C100 fault types and ranks, supporting the chosen repair location. Base ViT accuracies were 91.18% on C100 and 86.34% on TinyImg.
Applications
The framework could patch deployed visual models after systematic errors emerge—for example, reducing false positives in autonomous-driving perception or false negatives in medical imaging. Users need model-weight access, labeled correct and faulty examples, and an independent regression suite. It is best viewed as a controlled maintenance tool, not a replacement for safety validation or retraining.
Limitations & Outlook
The evidence is limited to ViT image classification, two datasets and one final FFN block; it does not test LLMs, multi-layer faults, out-of-distribution data or continual repair. DE introduces search time and parameter sensitivity, and the objective supplies no formal safety guarantee. Future work should combine broader architectures, provable repair, robustness testing and realistic industrial fault streams.
Plain Language Accessible to non-experts
Imagine a huge factory that sorts packages. Most packages are handled correctly, but a particular type is repeatedly sent to the wrong destination. Rebuilding the entire factory would be expensive and might ruin stations that already work well.
RepTran first watches correct and incorrect packages moving through the factory. It looks for small stations that behave very differently during wrong deliveries or become unusually active. It then marks only the few connecting screws most likely to matter, rather than touching every machine.
Next, many repair teams try slightly different settings for those screws. Each proposal earns points for correcting the bad deliveries and keeping all previously correct deliveries correct. The teams repeatedly keep better proposals, like controlled trial production.
On 18 fault groups from two image datasets, this process fixed 74.7% of target errors on average and up to 95.2%, while Arachne averaged 17.1%. The factory analogy has limits: the method still needs representative examples and has only been tested on one kind of vision model.
ELI14 Explained like you're 14
Imagine a photo-guessing game with a robot that is usually smart but keeps calling one kind of cat a fox. Retraining the whole robot could take ages and might make it forget things it already knows. RepTran tries a smaller repair: find the few internal number-settings most connected to that mistake.
First, it compares pictures the robot got right with pictures it got wrong. If an internal part becomes unusually active or changes a lot during wrong guesses, it gets a higher suspicion score. The method also checks which connections have the biggest effect on the final answer. So it is not randomly poking the robot—it is making a shortlist.
Then many repair attempts compete. Each attempt changes only the shortlisted numbers and gets points for two goals: fewer target mistakes and no new mistakes on pictures that were already right. A search method called Differential Evolution keeps improving the best attempts.
The experiments repaired 74.7% of target errors on average and 95.2% at best, compared with 17.1% for Arachne. Pretty impressive! But it was tested on ViT image models, not chatbots, and a repair that looks good in the lab still needs serious safety testing before real-world use.
Glossary
Feed-Forward Network (FFN)
A Transformer module that applies two linear transformations with a nonlinear activation between them. In this paper, Wbef and Waft in the FFN are the repair targets.
RepTran localizes and edits FFN weights.
Vision Transformer (ViT)
A vision model that represents image patches as tokens and processes them with Transformer encoder blocks. The evaluated model has 12 blocks, D=768 and Dhidden=3072.
The subject architecture in the experiments.
Differential Evolution (DE)
A population-based numerical optimization algorithm using mutation, competition and selection. It searches candidate values for the selected weights.
The search phase of RepTran.
NeuronScore
A score formed by multiplying normalized VDiff and MisAct. It highlights neurons whose variability shifts and whose activation is strong under the target fault.
It modulates weight-level suspiciousness.
Forward Impact and Gradient Loss
Forward Impact measures a weight’s influence on output, while Gradient Loss measures its loss-gradient contribution. Both are localization metrics from Arachne.
RepTran combines them with NeuronScore.
Repair Rate
The fraction of targeted misclassified samples that become correctly classified after repair. It is the main effectiveness metric, reaching 74.7% on average.
Used to compare RepTran with baselines.
Open Questions Unanswered questions from this research
- 1 It remains unknown whether RepTran scales to LLMs, multilingual tasks, cross-layer faults and real production failures. These settings require larger models, richer fault taxonomies and evaluation under distribution shift.
- 2 The method does not explain how to guarantee that a successful empirical patch is safe. Formal constraints, adversarial testing and long-term monitoring are needed to connect search-based repair with deployment assurance.
Applications
Immediate Applications
Autonomous-driving perception patching
An engineering team can collect examples of a recurring vehicle or pedestrian error, run RepTran on a ViT’s final FFN, and evaluate the patch on independent road-scene data. The expected outcome is fewer target errors with limited regression, subject to safety review.
Medical-imaging error correction
A device developer can target a lesion false-negative pattern using labeled correct and faulty scans while changing only a small FFN subset. Clinical review, out-of-distribution testing, calibration checks and regulatory validation remain prerequisites before deployment.
Long-term Vision
Auditable AI maintenance pipelines
A future system could connect fault monitoring, automatic localization, RepTran patch search, regression testing and versioned release management. This would make model maintenance more software-like, but formal safety, privacy, generalization and accountability remain major obstacles.
Abstract
To ensure the overall quality of AI-enabled software, not only traditional software components but also AI components need to be tested and repaired. Among AI components, Transformer models are increasingly integrated into software systems, which makes their misbehaviors critical. Although prior work in the software engineering community has proposed deep neural network (DNN) repair methods, most overlook Transformer-specific structures. We propose RepTran, a search-based repair method for Transformer models. It targets their feed-forward networks (FFNs), which play a central role in the architecture. RepTran identifies suspicious weights by combining two types of scores: a variance-based neuron score and an existing bidirectional score. It then iteratively optimizes these weights using differential evolution. Our evaluation includes 18 fault benchmarks constructed from CIFAR-100 and Tiny-ImageNet. We compare RepTran against three baselines: random weight selection, Arachne (a state-of-the-art DNN repair method), and ArachneW, which enables Arachne to control the number of selected weights. RepTran achieved an average repair rate of 74.7%, statistically outperforming random selection and Arachne across all benchmarks. Effect size analysis revealed that RepTran achieved higher repair rates than ArachneW regardless of the number of selected weights. These results suggest that RepTran is effective for enhancing the reliability of AI-enabled software.