Improving the Robustness of Transformer-based Large Language Models with Dynamic Attention

Improving the Robustness of Transformer-based Large Language Models with Dynamic Attention

Speakers: Lujia Shen

Conference: NDSS Symposium

YouTube: (no public video)

Overview

Transformer-based large language models (LLMs) like BERT and GPT have revolutionized natural language processing (NLP), achieving unprecedented performance across a myriad of tasks, from text generation to classification. However, their widespread adoption is hampered by a critical vulnerability: their susceptibility to textual adversarial attacks. Maliciously crafted input perturbations, often imperceptible to humans, can easily mislead these models, leading to severe consequences such as the generation of harmful content, misclassification of toxic comments, or incorrect responses in critical applications. Existing defense mechanisms, including computationally intensive adversarial training and certified robust approaches that often degrade performance, have struggled to provide scalable and effective solutions, particularly for the enormous scale of modern foundation models.

This talk introduces dynamic attention, a novel method specifically designed to enhance the inherent robustness of transformer architectures against diverse adversarial attacks. Developed by Lujia Shen, this approach addresses the limitations of prior work by directly integrating dynamic modeling into the attention mechanism, a core component of transformers. Dynamic attention aims to mitigate the impact of adversarial examples without requiring costly adversarial training, prior knowledge of downstream tasks, or significant computational overhead. Its design focuses on being compatible with various transformer architectures and can be synergistically combined with other defense modules to achieve even higher levels of protection.

The research presented offers a significant step towards building more resilient and trustworthy LLMs. By targeting the attention mechanism—a key area where adversarial examples exert their influence—dynamic attention provides a practical and effective defense strategy. The talk details comprehensive evaluations across different model types, tuning strategies, and attack scenarios, demonstrating its ability to reduce attack success rates, increase attacker queries, and maintain model stability, even against advanced adaptive and backdoor attacks.

Background

[▶ Watch: Transformer vulnerabilities and dynamic attention proposal (0:00)]()

The evolution of NLP has been profoundly shaped by the transformer architecture, introduced by Vaswani et al. in 2017. This architecture, with its innovative scaled dot-product attention mechanism, efficiently captures complex inter-word relationships, paving the way for powerful LLMs such as OpenAI's GPT series and Google's BERT. These models typically fall into three categories: encoder-only (like BERT), decoder-only (like GPT), and encoder-decoder (like T5). They are initially pre-trained on vast, unlabeled text corpora and then adapted to specific downstream tasks through strategies like full parameter fine-tuning, parameter-efficient prefix-tuning, or prompt-tuning. Despite their prowess, this very sophistication makes them vulnerable to subtle manipulations.

Textual adversarial attacks exploit these vulnerabilities by perturbing input texts in ways that are semantically preserved for human understanding but cause misclassification or erroneous generation by the model. These attacks are generally categorized by the granularity of manipulation: character-level, word-level, and sentence-level. This research primarily focuses on word-level attacks, which are often more challenging for both humans and deep neural networks (DNNs) to detect. Key word-level attacks evaluated include:

• TextBugger 31: A multi-level attack combining character and word-level perturbations.
• TextFooler 24: A black-box attack that identifies crucial words and replaces them with synonyms.
• PWWS 49: Similar to TextFooler, but uses counter-fitted word vectors for synonym selection.

Efforts to improve model robustness against such attacks broadly fall into two categories:

• Empirical Adversarial Training (AT): This involves augmenting training data with adversarial examples. While effective, AT is computationally expensive, especially for large foundation models, requires knowledge of downstream tasks, and struggles with generalization across all attack types. For large pre-trained models, AT necessitates repeated queries and updating all model parameters, which can hinder generalization.
• Dynamic Modeling: This is an emerging area where neural networks alter their parameters or structure during inference. Previous works like Stochastic Activation Pruning (SAP) 10 and defensive dropout 58 introduce randomness to hidden features. However, these methods often lead to information loss and do not specifically target the attention mechanism, which is a critical point of vulnerability in transformers.
• Certified Robust Training: These approaches offer provable robustness guarantees by verifying a lower bound of the robust space. While theoretically appealing, they frequently degrade performance on original tasks, are difficult to generalize in discrete NLP spaces, and are computationally prohibitive for large language models, often yielding only trivial bounds.

The presented work aims to bridge these gaps by introducing dynamic modeling directly into the attention mechanism, addressing the shortcomings of existing defenses and providing a more focused, efficient, and effective solution for transformer robustness.

Key Findings

[▶ Watch: Shortcomings of current NLP adversarial defense mechanisms (2:00)]()

Before designing dynamic attention, the researchers conducted three empirical studies to better understand the behavior of adversarial texts and their impact on the transformer's attention mechanism. These insights formed the foundational principles for their proposed defense:

1. Tokens with High Attention in Adversarial Texts Differ from Original Texts: Analysis revealed a significant disparity in how attention is distributed in adversarial examples compared to clean inputs. In clean texts, tokens receiving high attention are typically task-relevant, contributing meaningfully to the model's decision. However, in adversarial examples, high attention is often assigned to task-irrelevant tokens, misleading the model. Quantitatively, on datasets like Amazon and Twitter, the average prediction confidence difference between merged sentences formed by attentive tokens from adversarial texts versus their original clean counterparts was substantial (0.3713 for Amazon, 0.5305 for Twitter). These differences were significantly larger than those observed with random masking, strongly indicating that adversarial examples exploit the attention mechanism by misdirecting its focus.

1. Replacing Adversarial Attention with Original Attention Restores Correct Classification: To confirm the attention mechanism's role in adversarial misclassification, an experiment was conducted where the attention maps of adversarial texts were replaced with those of their corresponding clean versions. The results were striking: a high percentage of previously misclassified adversarial texts were correctly classified. Specifically, for fine-tuned models, an average of 88% of adversarial texts were rectified; for prompt-tuned models, this figure rose to 97%; and for prefix-tuned models, it was 80%. This compelling evidence directly implicated the attention mechanism as the primary target of adversarial attacks and demonstrated that rectifying it could effectively restore correct model behavior. This finding directly inspired the design principle of modifying the attention mechanism itself to mitigate attacks.

1. Most Adversarial Examples Are Inherently Unstable: Previous studies have highlighted a trade-off between transferability and imperceptibility in adversarial examples. This research further demonstrated that adversarial examples in NLP exhibit low transferability. By training two models on the same dataset, generating adversarial texts from one, and attempting to transfer them to the other, the researchers found low success rates. For instance, TextFooler's adversarial transfer rate on Amazon was only 41%, and TextBugger's on Twitter was 29%. This inherent instability means that adversarial examples are not consistently effective if the target model's parameters or internal states change. This observation provided the motivation for incorporating dynamic modeling into the defense, making it harder for attackers to consistently fool a model whose internal state is constantly shifting.

These three key findings collectively underscore that adversarial attacks on transformers often manifest as misdirected attention, that correcting this misdirection can restore model accuracy, and that the inherent instability of adversarial examples can be exploited to build more robust defenses.

Technical Deep Dive

[▶ Watch: Empirical study: Adversarial texts shift attention to irrelevant tokens (4:00)]()

Based on the empirical insights, the proposed dynamic attention mechanism is designed with two core modules: attention rectification and dynamic modeling. The overarching goal is to mitigate the effect of mistakenly assigning high attention values to irrelevant tokens, which leads to misclassification, and to exploit the inherent instability of adversarial examples by dynamizing the model.

Defense Overview

Dynamic attention operates by selectively masking or weakening the attention of tokens that receive unusually high attention values, thereby preventing the model from over-attending to potentially erroneous or adversarial tokens. Unlike traditional dropout, which randomly drops neurons, dynamic attention's intervention is more targeted, based on the computed attention values. To introduce dynamism, the number of modified tokens varies across transformer layers and with each inference run. This mechanism is designed to be compatible with existing transformer architectures and can be combined with other robustness modules, such as dropout, for enhanced defense.

Attention Rectification

In a transformer's multi-head attention mechanism, input sequences are transformed into query (Q), key (K), and value (V) matrices. Each attention head calculates an attention matrix, $A_t = \text{softmax}(\frac{Q_t K_t^T}{\sqrt{d}})$, where $d$ is the dimension of the key vectors. These attention maps are then summed across all heads to form a global attention map, $A = \sum_t A_t$. The total attention value for each key token $j$ is subsequently calculated as $A_s = \sum_i A_{i,j}$.

The core of attention rectification lies in identifying a set of "attentive tokens" for modification. The selection process is tailored to the specific NLP task:

• Text Classification: For tasks like sentiment analysis or spam detection, adversarial words are often assigned unusually high attention. Therefore, the top m tokens with the highest attention values are selected for modification.
• Text Generation: In tasks like machine translation or summarization, adversarial attacks may target less significant words. Here, the top m_a tokens are kept unchanged, while a subsequent range of tokens (from m_a to m_b) are selected for modification.

Once the set of target tokens T is identified, their attention values are adjusted by a **reduction factor *β***, a value between 0 and 1. This yields a rectified attention map $A'$. If β = 0, the attention is completely masked; if β = 1, it remains unmodified. This rectified attention map is then used to compute the hidden representations, influencing how information flows through the subsequent layers.

Dynamic Modeling

To prevent attackers from learning and exploiting a static set of modified tokens, dynamic modeling is introduced. The number of tokens to modify, m (or m_b for generation tasks), becomes a random variable.

• Text Classification: For each input text, m is a random integer drawn from a predefined range, typically proportional to the text's length (e.g., 10% to 20% of text length). This ensures that m changes with every run and in every transformer layer.
• Text Generation: m_a (the number of unchanged top tokens) is fixed, while m_b (the upper bound for the range of modified tokens) is a random variable that also changes across layers and runs.

This dynamic nature introduces an element of unpredictability, making it significantly harder for adversarial examples to consistently fool the model, as the model's internal state and the specific tokens being rectified change with each inference.

System Design and Architecture

The dynamic attention mechanism is seamlessly integrated into the self-attention layers of the transformer architecture. It functions as a modular component that intercepts the standard attention calculation. After the multi-head attention mechanism computes initial attention scores, the dynamic attention module takes over. It calculates the global attention map, identifies attentive tokens based on task-specific rules (top m for classification, m_a to m_b for generation), and then applies attention rectification using the β factor. The dynamic modeling aspect ensures the randomized selection of m or m_b for each layer and run, making the defense adaptive.

The rectified attention scores are then used to compute the weighted sum of value vectors, which are subsequently concatenated and projected to form the hidden representations. These representations then proceed through the standard Add & Norm and Feed Forward layers, followed by another Add & Norm and Dropout layer, before being passed to the next transformer block or the final output layer. This modular design facilitates easy integration into existing encoder-only, decoder-only, or encoder-decoder transformer models without requiring major architectural changes. It can also function as a supplementary module, complementing other defense mechanisms like dropout, information bottleneck, or adversarial training to form a more robust "fusion model."

Implementation Details

Experiments were conducted across various models, tasks, attacks, and defense strategies to thoroughly evaluate dynamic attention:

• Models: BERT-large (encoder-only) for classification (fine-tuned, prefix-tuned, prompt-tuned), GPT-2 (decoder-only) for classification, and T5-base (encoder-decoder) for text generation (machine translation, summarization).
• Hyperparameters for Dynamic Attention:
• Classification (BERT): For fine-tuned and prefix-tuned BERT, β = 0 (complete masking) and m ranged from 10%-20% of text length. For prompt-tuned BERT, β = 0.4 and m ranged from 20%-40%.
• Text Generation (T5): β = 0.6, m_a = 0.1, and m_b ranged from 30%-50% of text length. Dynamic attention was only enabled in the encoder for T5, as enabling it in both encoder and decoder degraded performance.
• Attack Methods: TextBugger, TextFooler, and PWWS were used. For classification, attackers aimed for high-confidence misclassifications (prediction confidence 0.6). For text generation, the goal was to minimize the BLEU score of the generated text.
• Defenses (Baselines): Compared against the original model, Dropout 58 (10% rate at test time), Empirical Adversarial Training (AT) using A2T 68, and Information Bottleneck (IB) 72 (using FreeLB adversarial training). Fusion models combining dynamic attention with these baselines were also evaluated.
• Datasets:
• Classification: Amazon (3.6M product reviews), Twitter (77k toxic comments), Enron (28k spam emails). Yelp (560k business reviews) was used for dataset shift evaluation.
• Text Generation: TED Talk (English-French, English-German translation), Gigaword (summarization).
• Metrics: Attack Success Rate (ASR) and Attack Query (Query) for classification; BLEU score for text generation.

Demo / Proof of Concept

[▶ Watch: Empirical study: Correcting attention maps restores classification (4:00)]()

The talk extensively details the evaluation and results of the dynamic attention mechanism across various threat models and datasets, presenting quantitative data in what would typically be tables in a paper. While a direct "demo" or "proof of concept" in the form of a live demonstration of the system was not explicitly mentioned in the transcript, the comprehensive experimental results serve as a robust empirical proof of concept for the effectiveness of dynamic attention. The evaluation covers how the system performs against different attacks, models, and tasks, providing strong evidence of its practical utility.

Defensive Implications

[▶ Watch: Empirical study: Adversarial examples exhibit inherent instability (4:00)]()

The findings from the comprehensive evaluation highlight several critical implications for defenders seeking to secure large language models against adversarial attacks:

1. Directly Address Attention Mechanism Vulnerabilities: Dynamic attention demonstrates that targeting the core attention mechanism is a highly effective defense strategy. By dynamically masking or weakening the influence of highly attentive, potentially adversarial tokens, defenders can significantly reduce the success rates of various word-level attacks (TextBugger, TextFooler, PWWS). This is particularly important because adversarial examples often mislead models by misdirecting attention to task-irrelevant tokens.

1. Increased Attack Cost and Detectability: Dynamic attention dramatically increases the number of queries required for attackers to generate successful adversarial examples. For fine-tuned BERT with TextFooler, dynamic attention increased queries from 379.42 to 744.54 on the Amazon dataset. This makes attacks more resource-intensive and provides more opportunities for detection through query monitoring.

1. Robustness Against Transfer Attacks: The dynamic nature of the defense makes adversarial examples less transferable. Under dynamic transfer attacks (where the attacker uses an identical dynamic surrogate), dynamic attention reduced the ASR by over 50% for fine-tuned BERT with TextFooler (from 52.90% to 24.80%). Against static transfer attacks (where the attacker uses a static surrogate), dynamic attention consistently outperformed dropout. This implies that even if attackers train models locally, their adversarial examples will likely fail against a dynamically defended online model.

1. Synergistic Defense with Existing Methods: Dynamic attention is highly compatible and synergistic with other defense mechanisms. Combining it with dropout (a fusion model) consistently yielded the best performance across various threat models and datasets, often achieving the lowest ASR and highest query counts. For instance, on fine-tuned BERT with TextBugger, the fusion model achieved an ASRQ of 26.83% with 56.04% ASRS, requiring 323.03 queries. Furthermore, dynamic attention significantly improved robustness when integrated with Adversarial Training (AT) and Information Bottleneck (IB), even in cases where AT alone performed poorly. This means defenders don't need to replace existing defenses but can augment them with dynamic attention for superior protection.

1. Task-Agnostic and Cost-Effective: A key advantage is that dynamic attention does not require downstream task knowledge or incur additional training costs, making it highly practical for large foundation models. It operates during inference, making it a "plug-and-play" solution.

1. Suitability for Text Generation Tasks: Unlike dropout, which often degrades performance in text generation due to information loss, dynamic attention improved the BLEU scores of adversarial texts in machine translation and summarization tasks. This is because its "weakening" mechanism preserves more information than dropout's aggressive "masking," making it a more suitable defense for generation-focused LLMs.

1. Enhanced Stability and Robustness Preservation: Dynamic attention models produce more stable outputs for adversarial examples (lower standard deviation of output confidence) compared to dropout, which tends to introduce excessive randomness. Crucially, dynamic attention preserves approximately 98% of the original model's robustness space, outperforming other dynamic modeling methods that often diminish robustness rapidly.

1. Mitigation of Advanced Attacks: The defense proves effective against backdoor attacks like BadNets, reducing ASR from 100% to less than 10%, by identifying and masking trigger tokens. While less effective against more elusive POR attacks, it still showed a reduction in ASRS for several triggers. Moreover, it demonstrated resilience against sophisticated adaptive attacks where attackers attempt to manipulate attention weights, confirming that direct manipulation of attention in discrete NLP space remains challenging for adversaries.

1. Performance Under Dataset Shift: Dynamic attention maintains its effectiveness even when the evaluation dataset differs from the training dataset (e.g., Amazon-trained, Yelp-evaluated), significantly increasing attack queries and reducing ASRs under these realistic scenarios.

In essence, defenders should consider integrating dynamic attention into their transformer-based NLP systems, particularly in security-critical applications. Its ability to provide robust protection against a wide range of attacks, increase attacker effort, and maintain model stability, all without significant overhead or task-specific training, makes it a valuable addition to the LLM security toolkit.

Key Takeaways

• Targeted Attention Defense: Dynamic attention is a novel, transformer-specific defense that directly modifies the attention mechanism to mitigate adversarial attacks by dynamically masking or weakening highly attentive, potentially malicious tokens.
• Empirically Proven Effectiveness: Extensive evaluations demonstrate significant reductions in Attack Success Rates (ASR) and substantial increases in Attack Queries across diverse models (BERT, GPT-2, T5), tasks (classification, generation), and word-level adversarial attacks (TextBugger, TextFooler, PWWS).
• Enhanced Stability and Robustness: Dynamic attention produces more stable outputs for adversarial examples compared to dropout and preserves approximately 98% of the original model's robustness space, outperforming other dynamic modeling methods.
• Synergistic with Existing Defenses: It effectively combines with other defense modules like dropout, adversarial training, and information bottleneck to form fusion models, achieving superior overall robustness.
• Adaptive and Transfer Attack Resilience: The dynamic nature of the defense makes adversarial examples less transferable and effectively defends against sophisticated adaptive attacks and backdoor attacks (like BadNets).
• Practical and Scalable: Dynamic attention requires no downstream task knowledge, no additional training costs, and is compatible with various transformer architectures, making it a highly practical and scalable solution for securing large language models.

About the Speaker(s)

Lujia Shen is the speaker for this presentation at the NDSS Symposium. The transcript and metadata provided do not contain further details about their title, affiliation, or specific background beyond their name.