Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MergeGuard: Efficient Thwarting of Trojan Attacks in Machine Learning Models

Soheil Zibakhsh Shabgahi, Yaman Jandali, Farinaz Koushanfar

TL;DR

The paper tackles Trojan backdoor attacks in large neural networks, especially transformers, by introducing MergeGuard, a post-training, regularization-based method that linearizes inter-layer activations to merge consecutive linear layers. This approach yields a shallower, compression-aware model that retains accuracy while substantially reducing attack success rates across multiple trojan methods. Empirical results on CIFAR-10 with ViT and CNN backbones show MergeGuard delivering consistent ASR reductions, significant runtime speedups, and up to 15% parameter and 14% MAC reductions, outperforming traditional fine-tuning and other defenses. The method is architecture-agnostic, scalable to vision transformers and CNNs, and accompanied by open-source code to support reproducibility and further research.

Abstract

This paper proposes MergeGuard, a novel methodology for mitigation of AI Trojan attacks. Trojan attacks on AI models cause inputs embedded with triggers to be misclassified to an adversary's target class, posing a significant threat to model usability trained by an untrusted third party. The core of MergeGuard is a new post-training methodology for linearizing and merging fully connected layers which we show simultaneously improves model generalizability and performance. Our Proof of Concept evaluation on Transformer models demonstrates that MergeGuard maintains model accuracy while decreasing trojan attack success rate, outperforming commonly used (post-training) Trojan mitigation by fine-tuning methodologies.

MergeGuard: Efficient Thwarting of Trojan Attacks in Machine Learning Models

TL;DR

The paper tackles Trojan backdoor attacks in large neural networks, especially transformers, by introducing MergeGuard, a post-training, regularization-based method that linearizes inter-layer activations to merge consecutive linear layers. This approach yields a shallower, compression-aware model that retains accuracy while substantially reducing attack success rates across multiple trojan methods. Empirical results on CIFAR-10 with ViT and CNN backbones show MergeGuard delivering consistent ASR reductions, significant runtime speedups, and up to 15% parameter and 14% MAC reductions, outperforming traditional fine-tuning and other defenses. The method is architecture-agnostic, scalable to vision transformers and CNNs, and accompanied by open-source code to support reproducibility and further research.

Abstract

This paper proposes MergeGuard, a novel methodology for mitigation of AI Trojan attacks. Trojan attacks on AI models cause inputs embedded with triggers to be misclassified to an adversary's target class, posing a significant threat to model usability trained by an untrusted third party. The core of MergeGuard is a new post-training methodology for linearizing and merging fully connected layers which we show simultaneously improves model generalizability and performance. Our Proof of Concept evaluation on Transformer models demonstrates that MergeGuard maintains model accuracy while decreasing trojan attack success rate, outperforming commonly used (post-training) Trojan mitigation by fine-tuning methodologies.
Paper Structure (21 sections, 11 equations, 4 figures, 2 tables)

This paper contains 21 sections, 11 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Trojan Attacks
  4. Types of Attacks
  5. Mitigation Strategies
  6. Methodology
  7. Threat Model
  8. Attacker
  9. Defender’s Objective
  10. Key Intuition
  11. MergeGuard
  12. Formal Definition
  13. Compression Ratio
  14. Convolutional Layers
  15. Alternative Activation Functions
  16. ...and 6 more sections

Figures (4)

  • Figure 1: Clean Image
  • Figure 2: Trigger Embedded Image
  • Figure 4: Illustration of the proposed MergeGuard methodology for mitigating trojans in neural networks. (a) Identification of a potentially compromised layer suspected to contain a trojan. (b) Application of fine-tuning with the regularization strategy described in Section \ref{['subsec:method']}, designed to incrementally guide the activation towards linearity. (c) Post-tuning, the activation function approximates an identity function. (d) Simplification of the network architecture by merging two consecutive fully connected layers into a single, reduced layer, effectively eliminating the intermediary layer.
  • Figure 5: Computational time for mitigation strategies on ViT.