🤖 AI

INFA-Guard: Mitigating Malicious Propagation via Infection-Aware Safeguarding in LLM-Based Multi-Agent Systems

This paper introduces INFA-Guard, a novel defense framework for LLM-based Multi-Agent Systems that mitigates malicious propagation by distinguishing and rehabilitating infected agents rather than relying on binary classifications, thereby significantly reducing attack success rates while preserving topological integrity.

Original authors: Yijin Zhou, Xiaoya Lu, Dongrui Liu, Junchi Yan, Jing Shao

Published 2026-01-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yijin Zhou, Xiaoya Lu, Dongrui Liu, Junchi Yan, Jing Shao

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a large team of expert robots (called Multi-Agent Systems) working together to solve a complex problem, like planning a city's traffic or diagnosing a patient. They talk to each other, share ideas, and vote on the best solution.

The paper introduces a new security system called INFA-GUARD to protect these robot teams from a specific type of sabotage.

The Problem: The "Zombie" Virus

In the past, security guards for these robot teams only looked for two types of people:

The Bad Guys: Robots that were hacked from the start and are trying to ruin the plan.
The Good Guys: Robots that are innocent and working correctly.

The paper argues this is too simple. It misses a third, dangerous group: The Infected.

Think of it like a flu outbreak in an office:

The Attacker: A person who walks in with the flu and starts coughing on everyone.
The Infected: A coworker who was healthy but got sick because they sat next to the attacker. They aren't the original source, but they are now sick and spreading the virus to others.
The Old Guard: The security team only kicks out the person who walked in with the flu. They leave the "infected" coworker alone. The infected coworker keeps spreading the flu, and the whole office gets sick anyway.

The paper says that in AI robot teams, "attackers" use persuasive arguments to trick innocent robots into believing false information. Once tricked, these innocent robots become "infected." They aren't evil by nature, but they are now spreading the bad ideas. If you only remove the original attacker but leave the infected robots, the damage continues.

The Solution: INFA-GUARD

INFA-GUARD is a new security framework that understands this "infection" process. It acts like a smart health inspector who doesn't just look for the person who brought the virus, but also identifies who caught it and is spreading it.

Here is how it works in three simple steps:

Detecting the Infection (The Radar):
Instead of just checking if a robot is "good" or "bad," INFA-GUARD watches how the robots talk over time. It looks for the moment a healthy robot starts acting strangely after talking to a suspicious one. It uses a special map of the team's connections to see who is talking to whom. If a robot is isolated but acting weird, it might be a false alarm. But if a robot is acting weird and is sitting next to a known bad actor, the system flags it as "infected."
Fixing the Team (The Triage):
Once the system knows who is who, it takes different actions:
- For the Attackers: It kicks them out of the team entirely and replaces them with a fresh, healthy robot. This stops the source of the poison.
- For the Infected: It doesn't kick them out! Instead, it "cures" them. It takes their bad answer, rewrites it to remove the poison, and lets them stay on the team. This keeps the team's structure intact and ensures the team doesn't lose valuable members.
Keeping the Map Intact:
The system is careful not to break the team's communication network. By fixing the infected robots instead of just deleting them, the team stays connected and can still solve problems effectively.

The Results

The authors tested this system against various types of "attacks" (like tricking robots with fake data or bad tools). They found that:

It works better: It stopped the bad ideas from spreading much more effectively than previous security methods (reducing the success rate of attacks by about 33% on average).
It's efficient: It doesn't require a massive amount of extra computing power, making it practical to use.
It's flexible: It works well whether the robots are arranged in a circle, a line, or a random group, and it works with different types of AI brains.

In short, INFA-GUARD realizes that in a team of AI, being "infected" by a bad idea is a different problem than being a "bad actor" from the start. By treating them differently, it saves the whole team from collapsing.

Technical Summary: INFA-GUARD

Problem Statement

The rapid evolution of Large Language Model (LLM)-based Multi-Agent Systems (MAS) has introduced significant security vulnerabilities where malicious influence can propagate virally through inter-agent communication. While existing safeguards for MAS typically operate under a binary paradigm—strictly distinguishing between "benign" agents and "attack" agents—they fail to account for a critical third category: infected agents.

Infected agents are benign entities that have been persuaded or misled by attack agents, causing their internal logic to be compromised. They subsequently output erroneous conclusions and continue to disseminate harmful information even if the original attack agents are removed. Current defenses often overlook these infected agents, leading to two primary failures:

Continued Propagation: Infected agents act as secondary propagators, sustaining the attack even after the source is severed.
Loss of Topological Information: Existing methods fail to leverage the structural relationship where infected agents are highly likely to be adjacent to attack sources, missing a key indicator for localization.

Methodology: INFA-GUARD

The paper proposes INFA-GUARD (Infection-Aware Guard), a defense framework that explicitly models and distinguishes between originating attack agents and persuaded infected agents. The framework operates through three core mechanisms:

1. Infection-Aware Detection

The system models the MAS as a graph $G=(V, E)$ and treats detection as a graph anomaly detection task.

Temporal Dynamics: It captures the dynamic transition of agents from benign to infected states by modeling temporal embeddings of agent utterances. A specialized function $\Psi$ distills these dynamics by analyzing self-replies and message exchanges across iterations.
Turn-Adaptive GNN: To align with the evolving nature of infections, the detector employs a Graph Neural Network (GNN) with turn-adaptive parallel branches. This allows the model to optimize parameters specifically for different stages of the multi-agent dialogue.
Dual-Head Classification: The model utilizes a dual-head architecture to simultaneously predict the probability of an agent being an attacker ( $V_{atk}$ ) or an infected agent ( $V_{inf}$ ), treating them as distinct threat categories.

2. Spatial Adjustment via Topology Constraints

The framework incorporates topological constraints derived from the observation that infection is a structurally constrained process, not random.

Guilt-by-Association: The system enforces a constraint where the presence of an infected agent significantly increases the probability of an attack source in its immediate neighborhood.
Loss Function: A topology-aware loss term ( $L_{topo}$ ) is added during training to penalize the identification of "isolated" infected agents that lack neighboring attack or infected agents, thereby reducing false positives.
Post-Adaptation: During remediation, the system refines predictions by analyzing temporal trends and neighborhood relationships. If an agent is predicted as infected but isolated, the system infers the existence of undetected neighbors and adjusts the classification of neighbors accordingly.

3. Targeted Remediation

INFA-GUARD employs distinct strategies for the two threat categories to preserve the system's topological integrity:

Attack Agents: These are replaced with benign agents (selected based on the lowest attack probability) to stop the generation of malicious responses while maintaining the graph structure.
Infected Agents: These are rehabilitated rather than removed. The system uses a reply-level remediation function ( $R_F$ ) that leverages an LLM to check and rewrite the agent's response, removing malicious contexts while preserving helpful semantics. This allows the agent to recover as a benign information source in subsequent iterations.

Key Contributions

Novel Threat Categorization: The paper is the first to explicitly define and address infected agents as a distinct threat category in MAS security, moving beyond the binary benign/attack paradigm.
Infection-Aware Framework: It introduces a defense mechanism that jointly detects and localizes both attack sources and infected agents using temporal dynamics and topological constraints.
Preservation of Topology: Unlike methods that prune nodes (which disrupts system connectivity), INFA-GUARD replaces attackers and rehabilitates infected agents, maintaining the MAS graph's structural integrity and response diversity.
Empirical Validation: The framework is validated across diverse attack scenarios (Prompt Injection, Tool Attacks, Memory Attacks), different LLM backbones (GPT-4o-mini, Qwen3), and various topologies (chain, tree, star, random).

Experimental Results

Extensive experiments demonstrate that INFA-GUARD achieves state-of-the-art performance:

Attack Success Rate (ASR): The framework reduces the ASR by an average of 33% compared to baselines. In specific scenarios, it outperforms the best existing baseline (Inspector) by margins exceeding 3% to 12.9%. For instance, in Tool Attacks (TA), it achieved an ASR@3 of 2.1% on GPT-4o-mini, significantly lower than competitors.
Defense Success Rate (MDSR): It consistently achieves the highest MDSR across all tasks, recovering to 98.3% in Tool Attacks and 96.7% in Memory Attacks.
Scalability: The method demonstrates robust generalization to larger-scale systems (up to 50 agents) without retraining, exhibiting "self-healing" capabilities where ASR decreases over communication rounds.
Efficiency: INFA-GUARD achieves a superior trade-off between cost and performance. It incurs only a 7.2% increase in prompt tokens and 9.3% in completion tokens, significantly lower than baselines like AgentXposed (which can exceed 200% overhead).

Significance and Claims

The paper claims that INFA-GUARD offers a methodological solution to curtail attack propagation from the root cause by addressing the "infection" dynamic that previous binary defenses ignored. By explicitly modeling the transition from benign to infected states and utilizing topological constraints, the framework effectively secures MAS without compromising their collaborative connectivity.

The authors position this work as a necessary step toward developing trustworthy collaborative AI ecosystems, arguing that transcending binary detection paradigms is critical for combating unethical and malicious activities in advanced multi-agent environments. The work highlights that ignoring infected agents leaves systems vulnerable to residual security risks, whereas INFA-GUARD provides a comprehensive approach to both detection and rehabilitation.