Improving Adversarial Robustness via Activation Amplification and Attenuation

Technical Summary: Improving Adversarial Robustness via Activation Amplification and Attenuation

Problem Statement
Deep Neural Networks (DNNs) are vulnerable to adversarial examples—inputs with imperceptible perturbations designed to mislead model predictions. This vulnerability is often attributed to the presence of "non-robust features" within the network's internal representations. While prior defense strategies have attempted to mitigate this via pruning, masking, or feature recalibration, these approaches often face trade-offs. Some methods discard non-robust features entirely, potentially losing useful information, while others introduce complex plug-in modules with significant parameter overhead and reduced interpretability. Furthermore, existing activation scaling techniques designed for Out-of-Distribution (OOD) detection often rely on non-differentiable operations (e.g., top-k selection) or exponential formulations that can cause gradient saturation, making them unsuitable for end-to-end adversarial training.

Methodology: Activation Amplification and Attenuation (A3)
The authors propose A3, a lightweight, plug-in module designed to enhance adversarial robustness through a controlled rescaling mechanism. A3 is integrated into the backbone network (e.g., after specific ResNet blocks) and operates on intermediate activations ($z$).

1. Learnable Masking and Scaling:
   Instead of hard pruning, A3 computes a channel-wise binary mask ($m$) using a learnable projection and a Gumbel-Softmax trick to ensure differentiability. This mask identifies channels based on the distribution of original activation magnitudes.
   The module calculates two statistics from the activations: the total magnitude ($s_1$) and the magnitude of the masked subset ($s_2$). Using these, it derives two scaling factors ($f_m$ and $f_{1m}$) via a logarithmic formulation:
   $$f_m = \frac{\ln(1 + s_2)}{\ln(1 + s_1)}, \quad f_{1m} = 1 - f_m$$
   This logarithmic approach is chosen over exponential formulations to ensure stable gradients and prevent numerical instability during adversarial training.

2. Dual Modes (Amplification and Attenuation):
   A3 utilizes the same learnable parameters to produce two distinct outputs by flipping the sign of the scaling operation:

   • Attenuation ($z_{att}$): $z \cdot [1 - \text{Scale}(z, m)]$. This mode suppresses the influence of non-robust features.
   • Amplification ($z_{amp}$): $z \cdot [1 + \text{Scale}(z, m)]$. This mode exaggerates the activation patterns, effectively serving as a "negative reference."

3. Novel Objective Functions:
   The training process leverages both modes to define new loss functions:

   • Cross-Entropy Ranking Loss ($L_{rank}$): A hinge loss that enforces the attenuated activations to yield a lower cross-entropy loss than the amplified activations. This creates a competitive optimization process where the model learns to degrade predictions in amplification mode while improving them in attenuation mode.
   • Contrastive Logit Loss ($L_{cl}$): A triplet-style loss that minimizes the distance between clean and adversarial logits in the attenuation mode while maximizing the distance between attenuated and amplified logits.

The final objective combines the standard adversarial training loss (e.g., AT, TRADES, MART) with these proposed components. During inference, only the attenuated mode ($z_{att}$) is used.

Key Contributions

• A3 Module: A lightweight plug-in that improves robustness with minimal computational and memory overhead, relying on a simple scaling mechanism rather than additional network capacity.
• Dual-Mode Mechanism: A novel approach that jointly learns to amplify and attenuate signals using the same parameters. The paper demonstrates that learning to degrade predictions in amplification mode simultaneously improves robustness in attenuation mode.
• Differentiable Scaling: The replacement of non-differentiable top-k operations (common in OOD detection) with a learnable, differentiable mask and logarithmic scaling to avoid gradient obfuscation and ensure stability.
• New Loss Functions: The introduction of contrastive and ranking losses that utilize amplified activations as negative references to guide the robustness learning process.

Experimental Results
Extensive experiments were conducted on CIFAR-10, CIFAR-100, and Tiny ImageNet using ResNet-18 and WideResNet-34-10 backbones, integrated with AT, TRADES, and MART training strategies.

• Performance: A3 consistently outperformed existing plug-in defenses (e.g., CAS, CIFS, FSR, FPCM, EMFF) across various attack types (FGSM, PGD-20/100, C&W, AutoAttack). For instance, on ResNet-18 with AT on CIFAR-10, A3 achieved a 4.99% improvement in ensemble accuracy over the vanilla backbone.
• Efficiency: A3 introduces negligible overhead. Compared to the baseline, it adds only ~0.27M parameters and ~0.0021G FLOPs for ResNet-18, significantly less than other masking-based methods like FTA2C.
• Ablation Studies:
  • Mode Analysis: The attenuation mode significantly outperforms the amplification mode and hard suppression strategies, confirming that selective attenuation is superior to complete feature discarding.
  • Loss Components: Both the ranking loss and contrastive loss are essential; removing them or replacing the hinge formulation with standard cross-entropy leads to performance degradation.
  • Scaling Formulation: The proposed logarithmic scaling factors ($f_m$) outperformed linear, quadratic, and exponential alternatives.
  • Distribution Shift: Analysis of activation distributions shows that A3 induces a clear separation between clean and adversarial examples in the attenuation mode, whereas the amplification mode causes distributions to overlap, validating the role of amplification as a negative reference.

Significance
The paper claims that A3 offers a promising direction for building robust and interpretable models by focusing on feature-based defense with few trainable parameters. By demonstrating that learning to degrade predictions in one mode enhances robustness in another, the work challenges the notion that defenses must solely focus on suppression. The method achieves state-of-the-art robustness among plug-in modules while maintaining high clean accuracy and introducing negligible computational cost, making it a practical and efficient solution for securing DNNs against adversarial attacks.