⚡ electrical engineering

Divergence-based Safety Measure for Large Language Models via Rational Inattention

This paper proposes a novel divergence-based safety measure for large language models under embedding-input attacks, which leverages the equivalence between transformer attention and rational inattention to quantify worst-case output distribution shifts while ensuring stealthiness.

Original authors: Anh Tung Nguyen, Quanyan Zhu

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Anh Tung Nguyen, Quanyan Zhu

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 you have a very smart, well-read robot assistant (a Large Language Model, or LLM) that helps you make important decisions. You ask it questions, and it gives you answers. But what if a sneaky hacker could whisper tiny, almost invisible changes into the robot's "ears" (its input data) to trick it into giving you a dangerous or wrong answer?

This paper proposes a new way to measure how safe a robot assistant is against these sneaky tricks. Here is the breakdown using simple analogies:

1. The Core Problem: The "Whisper" Attack

Think of the robot's brain as a library of text. A hacker doesn't need to break down the door; they just need to whisper a tiny, almost imperceptible change to the words the robot is reading.

The Goal: The hacker wants the robot to change its answer (e.g., from "The door is open" to "The door is locked") without you noticing.
The Challenge: How do we know if a robot is "safe"? We need a way to measure the worst-case scenario: How much can a hacker change the robot's mind before we catch them?

2. The New "Safety Meter": The Rationality Check

The authors created a new safety meter. Instead of just looking at the final answer, they look at how the robot thinks.

The Analogy: The Overworked Manager
Imagine a manager (the robot) who has to read a huge pile of emails (the input text) to make a decision.

Rational Inattention: Humans can't read every single word of every email perfectly. We have to be "rationally inattentive." We skim the most important parts and ignore the rest to save energy.
The Robot's "Attention": The robot does something similar. It uses a mechanism called "Transformer Attention" to decide which words in the sentence are most important.
The Connection: The paper argues that the robot's way of picking important words is mathematically identical to how a human manager decides what to focus on when they are tired or busy.

The Meter:
The authors built a "Rationality Meter" that checks: Is the robot focusing its attention in a normal, human-like way, or is it getting confused?

If the robot is acting normally, the meter reads low (safe).
If the robot is being tricked, its attention gets scattered, and the meter goes up.

3. The "Stealth" Rule

The safety test has a strict rule: The attack must be stealthy.

Imagine a thief trying to steal a painting. If they wear a bright neon suit, they get caught immediately.
The "Safety Measure" asks: What is the biggest change the thief can make to the robot's answer while still wearing a "camouflage suit" that keeps the Rationality Meter from screaming?

The paper calculates the Kullback–Leibler (KL) Divergence. In simple terms, this is a score that measures how different the robot's "tricked" answer is from its "normal" answer.

High Score: The robot's mind was easily twisted by a sneaky attack. (Unsafe!)
Low Score: The robot resisted the trick, or the trick was too obvious to be hidden. (Safer!)

4. The Experiment: GPT-2 vs. GPT-Neo

The authors tested this safety meter on two famous robot brains: GPT-2 and GPT-Neo.

They gave both robots the same "whisper" attacks (the same budget of invisible changes).
They watched how much the robots' answers changed while trying to stay hidden from the Rationality Meter.

The Result:
The two robots reacted differently.

GPT-2 was more easily tricked. Even with the "camouflage," its answers shifted significantly. It had a "leakier" safety profile.
GPT-Neo was harder to trick. It kept its answers closer to normal even when attacked.

Summary

This paper doesn't just say "hackers are bad." It builds a thermometer for AI safety.

It treats the AI's attention mechanism like a human's limited focus.
It uses that to create a "lie detector" that watches for subtle confusion.
It measures how much an AI's mind can be twisted by a hidden attack before the "lie detector" goes off.
It proves that some AI models are naturally more "tough" against these hidden tricks than others.

The ultimate goal is to help engineers pick the safer robot for important jobs and to build better defenses against future hackers.

Technical Summary: Divergence-based Safety Measure for Large Language Models via Rational Inattention

Problem Statement
The integration of Large Language Models (LLMs) into decision-support pipelines introduces system-level safety concerns, particularly regarding malicious perturbations of input tokens that alter model outputs and influence downstream human decisions. While quantitative safety measures are demanded from a systems and control perspective to support model comparison, runtime monitoring, and parameter fine-tuning, limited progress has been made in developing such measures for LLMs. Specifically, there is a need for a metric that quantifies the worst-case output perturbation caused by input attacks while ensuring the attacks remain stealthy to internal monitoring mechanisms.

Methodology
The paper proposes a divergence-based safety measure that quantifies the worst-case Kullback–Leibler (KL) divergence between the output distributions of a clean LLM and an LLM subjected to embedding-input attacks. The core methodological innovation lies in constructing a "stealthiness constraint" by leveraging the equivalence between transformer attention mechanisms and Rational Inattention (RI) models of human decision-making.

Rational Inattention (RI) Framework: The authors revisit the RI framework, where agents trade off expected utility against the information cost of processing data. The optimal action policy in RI is characterized by a softmax function derived from an optimization problem involving Shannon mutual information.
Structural Connection: The paper establishes that the attention weights in transformer-based LLMs (computed via a scaled dot-product softmax) share the same functional form as the optimal action policy in the RI framework. This allows the interpretation of transformer attention as a solution to an RI-type optimization problem where the model allocates limited information-processing capacity to context tokens.
Bounded Rationality Measure: Based on this connection, the authors define a head-wise "bounded rationality measure" ( $\Delta V^{(m)}(\tau)$ ). This metric represents the difference between the $\tau$ -bounded rationality value (the optimal value of the attention optimization problem with a temperature parameter $\tau$ ) and the full rationality value (the limit as $\tau \to 0$ ). This measure serves as an internal monitoring output indicating the "rationality level" of the LLM's attention allocation.
Safety Measure Formulation: The proposed safety measure is formulated as an optimization problem: maximizing the KL divergence between the clean and attacked output distributions, subject to two constraints:
- Energy Constraint: The total energy of the attack signals injected into embedding tokens is bounded ( $\sum \|\zeta_i\|^2 \leq \rho$ ).
- Stealthiness Constraint: The weighted sum of the bounded rationality measures across all attention heads must remain below a specified trust tolerance ( $\epsilon_T$ ). This ensures the attack does not significantly alter the model's internal attention-based monitoring signal.

Key Contributions
The paper presents three primary contributions:

Structural Correspondence: It demonstrates that transformer attention weights are the unique solution to an RI-type optimization problem, establishing a formal link between LLM mechanisms and bounded-rational decision-making models.
Head-wise Bounded Rationality Measure: It introduces a novel metric derived from the RI framework to quantify the rationality of attention allocations in LLMs.
Divergence-based Safety Measure: It utilizes the bounded rationality measure as a monitoring output to define a safety metric that quantifies the worst-case stealthy output shift. The paper also analyzes "perfectly undetectable attacks" (attacks that do not alter the monitoring output) and derives an upper bound for the safety measure using a Bregman-divergence argument, assuming the output score map is Lipschitz continuous.

Experimental Results
The proposed measure was evaluated on two pretrained causal language models: GPT-2 and GPT-Neo-125M.

Setup: A controlled candidate-prediction task was used where ambiguous sentences required predicting the next token (an adjective) from a restricted pool. Attacks were generated using projected gradient ascent to maximize the KL divergence of the output distribution while adhering to energy and stealthiness constraints.
Findings: The experiments revealed distinct safety profiles for the two models. GPT-2 exhibited several high-divergence attacks that remained within the stealthy region (satisfying the trust tolerance), whereas high-impact attacks on GPT-Neo-125M generally resulted in larger monitoring residuals. This demonstrates that the proposed measure can effectively distinguish between model-level safety profiles under identical attack budgets and trust tolerances.

Significance and Claims
The paper claims that this work fills a gap in developing quantitative safety measures for LLMs from a systems and control perspective. By interpreting transformer attention through the lens of rational inattention, the authors provide a principled mechanism for constructing internal monitoring outputs. The proposed divergence-based safety measure offers a way to quantify the vulnerability of LLMs to stealthy embedding-input attacks. The authors position this contribution as a step toward equipping LLMs with interpretable performance measures analogous to robustness, resilience, and safety measures found in traditional control systems. The work does not claim to solve all LLM security issues but rather provides a specific metric to distinguish model safety profiles and analyze worst-case stealthy impacts.