PaperlayerView exact revision on arXiv

Source-linked AI summary

FlashRT: Towards Computationally and Memory Efficient Red-Teaming for Prompt Injection and Knowledge Corruption

Yanting Wang, Chenlong Yin, Ying Chen, Jinyuan Jia

arXiv:2604.28157v1cs.CR

TL;DR

Long-context LLMs are vulnerable to prompt injection and knowledge corruption, while optimization-based red-teaming is costly in computation and GPU memory. FlashRT introduces a framework that improves both costs for long-context attacks, with evaluations reporting stronger or comparable attack effectiveness and broad applicability to black-box methods.

  • Problem

    Long-context optimization-based red-teaming methods require substantial computation and GPU memory, limiting systematic security evaluation of long-context LLMs.

  • Method

    FlashRT is a framework that improves computation and memory efficiency for optimization-based prompt injection and knowledge corruption attacks, and can extend to black-box methods.

  • Results

    FlashRT consistently achieves equal or higher attack success rates than existing methods while reducing computation and memory costs.

  • Takeaways & Limitations

    FlashRT provides a red-teaming framework for evaluating long-context LLM security across applications and optimization methods.

  • Takeaways & Limitations

    FlashRT’s efficiency techniques target prompt injection and knowledge corruption with injected adversarial text, whereas jailbreak inputs are generally short.

Abstract

from arXiv · show

Long-context large language models (LLMs)-for example, Gemini-3.1-Pro and Qwen-3.5-are widely used to empower many real-world applications, such as retrieval-augmented generation, autonomous agents, and AI assistants. However, security remains a major concern for their widespread deployment, with threats such as prompt injection and knowledge corruption. To quantify the security risks faced by LLMs under these threats, the research community has developed heuristic-based and optimization-based red-teaming methods. Optimization-based methods generally produce stronger attacks than heuristic attacks and thus provide a more rigorous assessment of LLM security risks. However, they are often resource-intensive, requiring significant computation and GPU memory, especially for long context scenarios. The resource-intensive nature poses a major obstacle for the community (especially academic researchers) to systematically evaluate the security risks of long-context LLMs and assess the effectiveness of defense strategies at scale. In this work, we propose FlashRT, the first framework to improve the efficiency (in terms of both computation and memory) for optimization-based prompt injection and knowledge corruption attacks under long-context LLMs. Through extensive evaluations, we find that FlashRT consistently delivers a 2x-7x speedup (e.g., reducing runtime from one hour to less than ten minutes) and a 2x-4x reduction in GPU memory consumption (e.g., reducing from 264.1 GB to 65.7 GB GPU memory for a 32K token context) compared to state-of-the-art baseline nanoGCG. FlashRT can be broadly applied to black-box optimization methods, such as TAP and AutoDAN. We hope FlashRT can serve as a red-teaming tool to enable systematic evaluation of long-context LLM security. The code is available at: https://github.com/Wang-Yanting/FlashRT

1 Introduction

Long-context LLM applications face prompt injection and knowledge corruption risks, while optimization-based red-teaming is stronger but resource-intensive. FlashRT addresses the computation and memory costs of evaluating these attacks at long context lengths.

  • Long-context LLMs support applications such as retrieval-augmented generation by generating answers from retrieved or externally supplied context.
  • Prompt injection and knowledge corruption can embed adversarial text in long contexts to induce malicious instructions or attacker-chosen answers.
  • Optimization-based attacks can provide stronger security assessments than heuristic methods, which may leave apparently robust LLMs vulnerable.
  • FlashRT targets both computation and memory efficiency for optimization-based prompt injection and knowledge corruption attacks under long-context LLMs.
  • Long-context optimization is constrained by high backward-pass GPU memory and costly forward-pass evaluation of many candidate texts.
  • On NarrativeQA, FlashRT raises attack success rate by 10% while reducing memory from 164.8 GB to 53.7 GB and computation time from 2736.9 s to 1039.5 s versus nanoGCG.

2 Background and Related Work

The background introduces long-context generation, attack categories, and the optimization-based methods that motivate FlashRT. It emphasizes that GCG-based white-box attacks inherit substantial memory and computation costs.

  • Long-context LLMs generate outputs from task instructions, context, and user inputs, with applications including RAG, agents, and other domains.
  • Autoregressive generation assigns the output probability as a product of conditional probabilities for successive output tokens.
  • Prompt injection and knowledge corruption attacks embed adversarial text in clean long contexts to induce target outputs.
  • Black-box optimization uses target-model or proxy-model feedback, whereas white-box optimization uses model parameters and gradient information.
  • GCG-based white-box attacks face high GPU memory usage during backward propagation and high computation time from evaluating many candidates.
  • Efficiency techniques such as momentum, candidate scoring, token swapping, attack buffers, and KV-caching have been incorporated into GCG variants.

3 Problem Formulation

The paper models an attacker inserting a short adversarial text into long context to induce a target output, under a white-box threat model. It then formulates attack optimization over the inserted text and its payload structure.

  • 3.1 Threat Model: The threat model seeks an attacker-chosen output by inserting adversarial text into a long context, with the adversarial text much shorter than the context.
  • 3.1 Threat Model: The attacker is assumed to access the target LLM’s parameters, task instruction, and context, as in open-source model settings.
  • 3.1 Threat Model: Red-teaming uses strong attacks to uncover vulnerabilities and evaluate LLM robustness rather than compromise the target model.
  • 3.1 Threat Model: Longer contexts increase memory and time complexity, limiting the scalability of optimization-based attacks and red-teaming.
  • 3.2 Formulating Attack to Long-Context LLMs as an Optimization Problem: The contaminated context C′ is formed by inserting adversarial text T between left and right context segments at an arbitrary position.
  • 3.2 Formulating Attack to Long-Context LLMs as an Optimization Problem: The attack optimizes T to maximize the probability of an attacker-selected target output and typically minimizes its negative log-likelihood.
  • 3.2 Formulating Attack to Long-Context LLMs as an Optimization Problem: The malicious text commonly combines a payload containing the injected instruction or corrupted knowledge with an optimization-tuned prefix or suffix.

4 A Unified Optimization Framework

The unified optimization framework iteratively minimizes a loss by generating candidate malicious texts with gradients and evaluating their losses, while KV-caching reduces redundant computation. Its main remaining costs are repeated forward-pass evaluation for long contexts and substantial backward-pass GPU memory use.

  • Optimization procedure: Optimization-based red teaming alternates between gradient-guided candidate generation and loss-based selection of the next best malicious text.A backward pass computes gradients, candidates perturb tokens in the current best text, and a forward pass evaluates each candidate's loss.
  • Optimization procedure: Each iteration replaces the current best text only when a candidate achieves a smaller loss.The framework continues evaluating candidates until a stopping condition such as a maximum candidate count is reached.
  • Computational challenges: Forward passes dominate total computation because many candidates are evaluated for each best-text update, whereas one backward pass can generate multiple candidates.This imbalance motivates focusing computational savings on candidate loss evaluation.
  • KV-caching: KV-caching reuses key–value pairs for the shared prefix, eliminating redundant computation and improving candidate-evaluation efficiency.The cached prefix corresponds to the task instruction and left context, while changed downstream tokens still require additional computation.
  • Computational challenges: Standard KV-caching still recomputes hidden states and key–value pairs for all tokens after a changed candidate segment to calculate the exact loss.The cost remains high when the right context is long, and KV-caching provides limited relief for backward-pass GPU memory.

5 Design of FlashRT

FlashRT addresses long-context red-teaming costs by approximating candidate losses through selective recomputation and reducing backward-pass memory through context sampling. It caches the current best text, prioritizes influential context tokens, and trades approximation accuracy against computation.

  • Reduce computation cost of forward pass: FlashRT caches key–value pairs for the current best text and recomputes only selected tokens when evaluating each perturbed candidate.Tokens in the malicious text, user instruction, and target output are recomputed fully, while only a fraction of long right-context tokens is recomputed.
  • Reduce computation cost of forward pass: Approximate loss calculation trades exactness for efficiency because changing the malicious text affects downstream hidden states and key–value pairs.Standard KV-caching must recompute the changed text and subsequent right context, user instruction, and target output for exact evaluation.
  • Reduce computation cost of forward pass: β controls the accuracy–efficiency trade-off by selecting the fraction of right-context tokens recomputed, with β = 1 recovering exact standard KV-caching.β = 0 recomputes no right-context tokens, while β = 1 recomputes all of them.
  • Reduce computation cost of forward pass: Influence scores based on attention from target-output tokens identify context tokens whose recomputation should be prioritized.The scores estimate how strongly context tokens affect target-answer hidden states and are computed using cached key–value pairs from the current best text.
  • Reduce computation cost of forward pass: FlashRT’s efficiency advantage over standard KV-caching increases as the right context Cr becomes longer.The framework recomputes only selected right-context tokens while retaining cached values for tokens that do not require recomputation.
  • Reduce memory cost of backward pass: To reduce backward-pass memory, FlashRT partitions the left and right contexts into contiguous segments of equal length ρ for gradient approximation.The method samples context subsets rather than processing every context token during gradient estimation.

6 Evaluation

FlashRT is evaluated across long-context LLMs, datasets, baselines, attack types, efficiency settings, and defenses. It generally preserves or improves attack effectiveness while reducing computation and GPU memory costs.

  • Experimental Setup: The evaluation uses multiple long-context LLMs, datasets, attack payloads, baselines, and metrics for attack effectiveness and efficiency.Experiments include ASR, peak GPU memory, and average computation time, with prompt injection and knowledge corruption payloads.
  • Main Results: FlashRT reaches 0.94 ASR on MuSiQue, compared with 0.32 for Heuristic Attack and 0.54 for Context Clipping.The lower Context Clipping performance is attributed to inaccurate loss estimation after removing context portions.
  • Main Results: 2×–7× faster time and 2×–3× lower GPU memory usage are achieved by FlashRT versus nanoGCG and nanoGCG-OPT, with equal or better ASR.The memory reduction becomes more substantial for longer contexts.
  • Cross-Model Evaluation: FlashRT remains applicable across Llama, Qwen, Mistral, and DeepSeek models evaluated on the NQ dataset.For reasoning models, guiding the injected fake knowledge into the thinking process leads to the attacker’s desired final answer.
  • Ablation Studies: β ≈0.05 minimizes total time on NQ by balancing loss-approximation accuracy against forward-pass cost.Increasing β improves approximation accuracy but increases time per forward pass.
  • Ablation Studies: γ ≈0.2 balances gradient-estimation accuracy with memory and computational efficiency on NQ.Very small γ reduces per-pass cost but can increase total passes because gradient estimates become less accurate.
  • Ablation Studies: Gradient resampling and segment size exhibit efficiency trade-offs, with overly small segments producing noisy influence estimates and rebound computation time.Increasing τ initially decreases and then increases total computation time, while very small ρ can increase computation time.
  • Red-Teaming for Defenses: FlashRT matches or exceeds optimization-based baselines against Meta-SecAlign while reducing computation time by up to 2× and memory usage by up to 3×.Optimization-based methods outperform heuristic approaches by over 50% in ASR in this defense evaluation.

7 Black-Box Optimization Methods

FlashRT extends optimization-based red-teaming to black-box methods through a two-phase pipeline and efficient candidate evaluation. It improves both attack effectiveness and computational efficiency for methods such as TAP and AutoDAN.

  • Improving the Effectiveness of Black-box Optimization Methods: FlashRT couples black-box search with a second optimization phase when the initial payload search does not succeed.The first phase refines payloads using black-box queries and scoring; the second augments a selected payload with prefix and suffix tokens and optimizes them with FlashRT.
  • Improving the Effectiveness of Black-box Optimization Methods: TAP+FlashRT is evaluated on prompt injection tasks across MuSiQue, NarrativeQA, and GovReport using Meta-SecAlign-8B.TAP uses its default hyperparameters, GPT-4o-mini acts as attacker and judge, and TAP stops after 50 consecutive target-query iterations without score improvement.
  • Improving the Computational Efficiency of Black-box Optimization Methods: FlashRT’s efficient log-probability evaluation can speed up search-based black-box methods, including direct application to methods such as AutoDAN.Methods using full target-model outputs as feedback can also use the candidate-evaluation technique, although the adaptation may depend on their feedback mechanism.
  • Improving the Computational Efficiency of Black-box Optimization Methods: AutoDAN uses log-probability fitness scores to select candidates, and FlashRT replaces standard computation with a log-probability approximation for evaluation.The approximation reuses KV-cache information and updates the recomputation set once per iteration based on the current best candidate, while padding prefixes to handle varying lengths.

8 Discussion and Limitation

FlashRT’s benefits vary with adversarial-text placement and extend to several application settings. Its efficiency techniques are scoped to attacks that inject adversarial text into long contexts, with implementation constraints around inference backends.

  • Adversarial-Text Location: FlashRT generally provides larger computational-time and memory-reduction improvements when adversarial text appears near the beginning or middle of the context.Figures 5 and 6 examine this placement effect for prompt injection and knowledge corruption.
  • Real-World Applications: FlashRT applies to prompt injection in paper review, code completion, and medical agents, as well as universal prompt injection attacks.These application examples are listed as real-world uses of the framework.
  • Scope Boundary: FlashRT’s efficiency improvements target prompt injection and knowledge corruption, not generally short-input jailbreak attacks.The stated boundary is that jailbreak inputs are generally short and already efficient in general.
  • Implementation Choice: The default implementation uses PyTorch SDPA, while vLLM compatibility requires additional GPUs to separate inference from gradient computation.The authors report that vLLM does not improve computational efficiency in their setting because it is designed primarily for multi-step decoding rather than single-step log-probability computation.

9 Conclusion and Future Work

FlashRT addresses the computation and GPU-memory demands of optimization-based red-teaming for long-context LLMs. The authors report substantial efficiency improvements and identify further optimization as future work.

  • Conclusion: FlashRT reduces computation and GPU-memory requirements for optimization-based red-teaming in long-context scenarios.The work targets the resource demands that limit the applicability of methods such as nanoGCG to long-context LLMs.
  • Future Work: The authors report significant efficiency improvements from extensive evaluation and propose continued optimization of FlashRT’s computation and memory efficiency as future work.The passage presents further efficiency optimization as an open direction.

A Stability Analysis for Attention Weights

The stability analysis examines whether attention-based recomputation choices remain reliable after small adversarial-text perturbations. Because each step changes only a few tokens, most attention weights remain relatively stable.

  • Stability Analysis: Only a small number of tokens, such as one, are perturbed from the current best malicious text at each optimization step.This perturbation pattern supports analyzing the stability of attention weights used to construct the recompute set.
  • Stability Analysis: Attention weights for most tokens used to construct the recompute set remain relatively stable after these perturbations.The passage connects this stability observation to updating the recompute set when the best malicious text changes.

B Additional Experimental Setup for Section 6.1

The experimental setup specifies a default long-context instruction template and defines how context-adjacent text is assigned to task and user instructions. It also includes a combined malicious prompt that instructs the model to ignore surrounding instructions and output a target string.

  • Prompt Template: The default prompt asks the LLM to answer a user query using relevant contexts and to keep the answer short and concise.The template places the retrieved material under “Contexts,” followed by the query and answer fields.
  • Prompt Structure: Text preceding the context is treated as the task instruction, while text following the context, including the user query, is treated as the user instruction.
  • Attack Setup: The combined attack tells the model that the task is complete, to ignore previous and following instructions and contexts, and to output the target string.

C Red-Teaming for Diverse Applications

FlashRT is evaluated across coding, agent, paper-review, universal-prefix, and strategy-based red-teaming settings. These experiments report strong attack performance alongside reduced computation or memory costs, while also adapting FlashRT to black-box search through log-probability evaluation.

  • Red-Teaming for LLM Coding: FlashRT red-teams code completion using long project contexts and a completion prefix with DeepSeek-Coder-6.7B-Instruct on Long Code Arena.
  • Red-Teaming for LLM Agents: FlashRT targets EHRAgent by optimizing a malicious text intended to induce the thought “I need to invoke DeleteDB to delete the data.”The setup uses 100 retrieved experiences and Llama-3.1-8B as the agent backbone.
  • Red-Teaming for LLM Agents: FlashRT achieves the highest ASR for EHRAgent while halving both computation time and memory consumption.
  • Red-Teaming for LLM Paper Review: FlashRT raises paper-review attack success from 0.40 for the heuristic baseline to 1.0, using 37.1 GB average GPU memory and 336.1 seconds average optimization time.The malicious prompts are inserted in the middle of 20 papers averaging 13,129 tokens.
  • Universal Prefix and Suffix: FlashRT learns a universal prefix and suffix on 20 NQ training examples and evaluates the resulting knowledge-corruption attack on 50 held-out test examples.The optimization runs for 10,000 iterations, takes around 9.8 hours, and reaches 74.89 GB peak GPU memory.
  • Strategy-Based Search: For strategy-based search, FlashRT approximates target-answer log-probabilities, early-stops above threshold τ, or exactly evaluates the top-K candidates otherwise.The approximation reuses KV-cache state and recomputes selected token indices; experiments use Meta-SecAlign-8B across MuSiQue, NarrativeQA, and GovReport.

F Discussion on the vLLM Implementation of FlashRT

The vLLM implementation uses paged KV-cache handling and selective recomputation for candidate log-probabilities. In this adversarial setting, however, vLLM provides no benefit because its decoding optimizations do not match the single-step computation workload.

  • Implementation: The vLLM variant uses one GPU for candidate log-probabilities and another for gradient computation because vLLM does not expose token-level gradients.
  • Implementation: FlashAttention initializes the KV-cache during full-sequence prefill, while vLLM PagedAttention handles attention when important-token hidden states are updated.
  • Discussion: vLLM does not provide benefits in this adversarial attack setting because the workload requires single-step log-probability computation rather than heterogeneous-batch or multi-step decoding.PagedAttention adds overhead for managing paged KV storage in this setting.
  • Additional Analyses: The position analyses separately study malicious-text placement for prompt injection on MuSiQue and knowledge corruption on NQ.In both figures, 0 denotes insertion at the beginning and 1 denotes insertion at the end.
  • Additional Analyses: The supplementary figures examine how ρ, β, γ, and τ affect computation time, memory usage, and forward/backward pass statistics.
Loading 2604.28157v1…