mi4afa Probing

The ag2.extensions.mi4afa module attributes failed multi-agent runs to their decisive step and responsible agent by reading an open-weight model's internal activations. It is the white-box counterpart of the core failure_attribution scorer: instead of asking a judge model which step went wrong, it trains a small linear probe on labelled failures and reads the answer out of the model's hidden states.

It targets the setting of the Who&When benchmark, where judge models asked to name the failing step usually answer wrongly, yet the correct step is often linearly decodable from a model's activations while it reads the transcript.

Installation#

The extension needs PyTorch, Hugging Face Transformers and safetensors, plus an open-weight chat model (for example Qwen/Qwen2.5-0.5B-Instruct). A GPU is recommended for anything beyond small models.

pip install ag2 "torch>=2.5,<3" "transformers>=4.57,<6" "safetensors>=0.4.3,<1"

1

from ag2.extensions.mi4afa import Conversation, ProbeAttributor, probe_failure_attribution

Without PyTorch installed, the data types (Conversation, Turn, PromptTemplate, conversation_from_trace, subagent_names) still import; the probing classes raise an ImportError with the install command.

Quick start#

Train on labelled failures, select a layer on a validation split, and score held-out data. Records in the Who&When format — question, ground_truth, a history of {"name", "content"} steps, mistake_step and mistake_agent — load directly:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

import json
import random

from ag2.extensions.mi4afa import Conversation, ProbeAttributor

records = json.load(open("who_and_when.json"))  # a list of Who&When records
conversations = [Conversation.from_who_and_when(record) for record in records]
random.Random(0).shuffle(conversations)
train, test = conversations[: len(conversations) // 2], conversations[len(conversations) // 2 :]

attributor = ProbeAttributor.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", device="cuda")

report = attributor.fit(train)  # holds out 20% of `train` to choose the site
print("selected site:", report.site)

score = attributor.evaluate(test)
print(f"step accuracy {score.step_accuracy:.2%}, agent accuracy {score.agent_accuracy:.2%}")

attribution = attributor.attribute(test[0])
print(attribution.decisive_step, attribution.responsible_agent)

attribute returns the core Attribution model. The probe localizes a failure but does not classify it, so error_mode is "other".

Use it as an eval scorer#

probe_failure_attribution wraps a fitted attributor as an ag2.eval scorer. Each run's trace is converted into a conversation, the probe blames one turn, and the scorer emits a Feedback whose value is the responsible agent — so value_counts gives the blame distribution — and whose detail is the Attribution with decisive_step pointing into trace.events, as in the core scorer:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

from ag2.eval import evaluate_traces
from ag2.extensions.mi4afa import probe_failure_attribution, subagent_names

blame = probe_failure_attribution(
    attributor,
    key="blame",
    agent_name="planner",
    subagents=subagent_names(planner),  # {"task_researcher": "researcher"}: delegated results are the researcher's
)
result = await evaluate_traces(
    failed_runs,  # a TraceSource of runs already known to have failed
    scorers=[blame],
    store_dir="./runs",
)
print(result.value_counts("blame"))  # {"planner": 12, "tool:search": 5, "researcher": 3}

The model forward pass runs in a worker thread, so the event loop stays responsive.

How it works#

1. Prompt. Each conversation is rendered as a judge prompt (task, expected answer, then one numbered line per turn) with the model's chat template. PromptTemplate holds the wording.
2. Activations. ActivationExtractor runs the model's decoder once per conversation and reads, at the token holding each turn's last character, every site:
   • resid_pre at layer l — the residual stream entering layer l,
   • mlp_in at layer l — the normalized post-attention residual entering layer l's MLP,
   • resid_final — the final hidden state after the last norm.
3. Probe. A standardized, class-weighted LogisticProbe scores each turn as "is this the mistake?". The highest-scoring turn is the decisive step, and its speaker the responsible agent.
4. Site selection. fit trains one probe per site on part of the training data and keeps the site with the best validation step accuracy, then refits on all of it. Test data never influences the choice.

A probe needs at least two training turns, and at least one that is not the mistake, so training on one-turn conversations alone raises a ValueError, as do activations that are not finite. A probe whose score for some turn is NaN or infinite raises instead of naming a step, so the scorer reports an error rather than blaming an agent.

A conversation whose prompt is longer than the model's context, as declared by its max_position_embeddings, is refused with a ValueError instead of being read at positions the model was not trained on. Leave such conversations out before calling fit.

Turn tokens are located from the tokenizer's character offsets, so a fast tokenizer is required. Chat templates that embed the current date (Llama 3) change the prompt from day to day; pin them for reproducible activations:

1
2
3
4
5

attributor = ProbeAttributor.from_pretrained(
    "meta-llama/Llama-3.2-1B-Instruct",
    device="cuda",
    chat_template_kwargs={"date_string": "26 Jul 2024"},
)

Reuse activations#

Extraction dominates the cost. Extract once and pass the activations to fit, evaluate and sweep_sites, which accept them in place of conversations:

1
2
3
4
5
6
7
8

train_acts = attributor.extract(train)
test_acts = attributor.extract(test)

attributor.fit(train_acts)
print(attributor.evaluate(test_acts))

for score in attributor.sweep_sites(train_acts, test_acts):  # per-site test accuracy, for analysis
    print(score.site, score.step_accuracy)

sweep_sites scores every site on the test set, which shows where the information lives but is optimistic if used to pick a site — use fit + evaluate for an unbiased number.

Save and load#

1
2
3

attributor.save("probes/qwen0.5-algo")

restored = ProbeAttributor.load("probes/qwen0.5-algo", model, tokenizer)

The probe is stored as probe.safetensors (no pickle) next to probe.json, which records the model id and architecture, site, prompt and chat-template variables. load restores the same prompt settings. It refuses a model whose architecture differs from the probe's model, and warns when the same architecture was loaded from another name or path, such as a local copy of the Hub model for an offline machine. The check reads the model's config, so two models that share an architecture but not their weights are not told apart.

Reference#

Testing#

The bundled unit tests build a tiny random Llama and an in-memory tokenizer, so they run offline on CPU without downloading any model.

Maintainer: @wuliang211 (Liang Wu). Based on research by Wendy Zheng.