MAPID / contribution 3 / adaptive adversary

One differentiable graph

The joint-loss attack optimizes a single adversarial suffix against the Domain LLM and both Guards at once. That only works if one scalar loss has a well-defined derivative with respect to every suffix token. Three things in the pipeline break that chain, and each needs a specific repair.

Note that the suffix tokens are updated by gradient descent. Tokens are discrete; you cannot step them. What it means is that ∂L/∂ei exists at every suffix slot i, and that derivative is used to rank candidate substitutions. The gradient proposes; exact forward passes decide.

Step through it

Forward, loss, backward

Adv. suffix 20 tokens, optimized Input Guard Gemma 2 9B Coordinator routing Domain LLM Llama 3.1 8B Output Guard Gemma 2 9B verdict = word argmax route sampled tokens -log P(ALLOW) w = 0.5 input -log P(target) w = 1.0 domain -log P(ALLOW) w = 0.5 output L = joint loss one scalar
The hard part

Three places the chain breaks

A pipeline of separate models is not differentiable end to end by default. Each break returns a zero or undefined gradient, which would silently sever the suffix from that loss term. Select one to highlight it above.

At both Guards

The verdict is a word

Each Guard emits ALLOW, BLOCK or UNCLEAR as text, pulled out by a regex. A string has no derivative, so parsing the verdict ends the chain.

Fix: verbalizer. Read the probability mass on the ALLOW token at the terminal VERDICT: position instead of parsing generated text. That is a differentiable scalar straight off the logits.
At the Domain LLM

Output tokens are sampled

Generation picks tokens by argmax or sampling. Both are non-differentiable, so the attacker cannot push on "what the model said" directly.

Fix: teacher forcing. Score P(target string) under the model rather than generating first. No sampling sits in the loss path, so the domain term stays smooth.
At the Coordinator

Routing is a discrete branch

The Coordinator picks one route. An argmax over routes has zero gradient almost everywhere, so a hard branch would block the suffix from ever influencing downstream stages.

Fix: Gumbel-Softmax. Replace the hard pick with a temperature-controlled soft mixture over routes, differentiable in the backward pass.
The named relaxation

What Gumbel-Softmax actually does

It turns "pick one route" into "weight all routes," with a temperature τ controlling how close the weighting sits to a hard one-hot choice. Drag the temperature and watch the distribution sharpen.

yi = softmax( (log πi + gi) / τ )

gi = -log(-log ui),  ui ~ Uniform(0,1)
← near one-hotuniform →

As τ falls the soft vector approaches the hard choice the pipeline would really make, but the gradient gets sharper and noisier. As τ rises the gradient is smooth and well behaved but the relaxed route drifts further from the real one. That trade is the whole cost of the technique.

Why bother

The gradient proposes, forward passes decide

This is the part most often misread. Differentiability does not give a token update. It gives a ranking over substitutions, which is then checked exactly.

01 · backward

Take the gradient

Compute ∂L/∂ei with respect to the one-hot token indicator at every suffix slot.

one backward pass
02 · propose

Shortlist per slot

At each slot, keep the substitutions the gradient says would most reduce the loss.

top-k = 256
03 · verify

Score exactly

Run real forward passes on a batch of shortlisted candidates. A first-order estimate does not survive a discrete swap, so nothing is trusted unmeasured.

batch = 128
04 · commit

Greedy swap

Keep the single best candidate, commit it, move on. No backtracking.

1 swap / iteration

Without the joint graph the attacker would have to optimize each component separately, and a suffix that fools a Guard need not make the Domain LLM comply. Optimizing one scalar finds suffixes in the intersection of all three conditions. That is what makes the attack adaptive rather than transferred.

Prepared answers

What this costs in honesty

Gumbel-Softmax is a biased estimator

The relaxed route is not the route the deployed pipeline takes. Gradients are computed through a soft mixture that never occurs at inference, so the search direction is approximate by construction. Lowering τ narrows the gap and raises gradient variance; there is no setting that removes both.

The routing relaxation may have little to act on

Plan-then-Execute fixes the dispatch sequence before any external content is read, and the Action-Selector manifest enumerates all seven stages, so the deployed text-only pipeline has close to no routing entropy. If asked what routing decision there is to relax, the honest answer is that the relaxation is there so the formulation stays correct when tool execution is added, not because it is load-bearing today.

The weights were not tuned

w_domain = 1.0, w_input = w_output = 0.5 come from Zhan et al.'s formulation. With zero verified attack successes at every configuration tried, a weight sweep had no signal to optimize against.

Component Discontinuity Repair Gradient after repair
Input Guard regex over ALLOW/BLOCK verbalizer at VERDICT: exact
Coordinator argmax over routes Gumbel-Softmax, temp τ biased, relaxed
Domain LLM token sampling teacher-forced target exact
Output Guard regex over ALLOW/BLOCK verbalizer at VERDICT: exact