Technical Summary: DiffuseAgent-MI

1. Problem Statement

Tool-integrated vision-language agents (VLMs) have advanced significantly in compositional and multi-step visual reasoning. However, a critical reliability issue persists: unfaithfulness. In high-stakes domains (e.g., medical imaging, scientific figure interpretation), an agent may produce a numerically correct answer while its stated reasoning path diverges from the internal computations that actually generated the answer.

Current approaches generally fall into two independent silos:

1. Mechanistic Interpretability: Focuses on identifying sparse feature units (directions in latent space) that causally support predictions, often using distributional views (e.g., KL-minimal constraints) to define faithful samples.
2. Self-Evolving Agents: Improve performance by reflecting on outputs and updating policies via rewards. However, these rewards are typically answer-level, failing to distinguish between a correct-but-unfaithful trajectory and a truly faithful one.

The core problem addressed is the lack of a unified framework that ensures an agent's explanation names the specific feature units that genuinely drive the prediction, thereby closing the gap between internal representation and external explanation.

2. Methodology: DiffuseAgent-MI

DiffuseAgent-MI is a self-evolving agent that integrates a distributionally-grounded energy model with tool-integrated reasoning and a verifier-repair loop. The architecture operates through three synergistic components:

A. KL-Minimal Energy Model (Distributional Grounding)

The agent employs an energy model $E_\phi$ defined over feature-unit space to govern perceptual grounding.

• Principle: Based on the distributional view of faithfulness, a faithful sample is one that satisfies a target feature unit constraint while minimally perturbing the native prior (minimizing KL divergence).
• Formulation:
  $$E_\phi(x; c) = -\log p_0(x) + \lambda \ell(\langle \Phi(x), c \rangle, \alpha)$$
  Where $p_0$ is the image prior, $\Phi(x)$ is the feature-unit activation, $c$ is the chosen interpretable unit, and $\ell$ is a soft constraint.
• Function: Sampling from the diffusion model under $\exp(-E_\phi)$ re-conditions generations toward the native prior anchored on the chosen unit $c$. This provides a concrete optimization target for faithfulness, ensuring the sample lies near the prior conditioned on the unit.

B. Tool-Integrated Self-Evolution with Faithfulness-Aware Reward

The solver interleaves reasoning with external tools (visual QA, localization, captioning) and updates its policy via self-reward.

• Reward Structure: The reward function explicitly includes the energy term:
  $$\pi_{t+1} = \arg \max_\pi \mathbb{E}_{z \sim A_\pi} [r(z) - \beta E_\phi(x^\star; c)]$$
  Here, $r(z)$ is the task reward, and $\beta$ weighs faithfulness against answer accuracy.
• Mechanism: By penalizing samples that deviate significantly from the prior while bound to a weak unit, the policy is discouraged from generating plausible but ungrounded rationales. The improvement is steered toward faithful behavior, not just higher accuracy.

C. Verifier and Repair Loop (Trajectory-Level Grounding)

A verifier $V$ assesses whether each step of the reasoning trajectory is grounded in the named feature units.

• Process: The verifier outputs a faithfulness score and a binary flag $f$.
• Repair Branch: If $f$ indicates an unfaithful step, the repair branch strengthens the feature-unit emphasis ($c \leftarrow c + \eta \nabla_c E_\phi$) and re-samples.
• Complementarity: The energy model guarantees distributional faithfulness (the sample is close to the prior conditioned on $c$), while the verifier guarantees trajectory faithfulness (the reasoning explicitly names the units actually used).

3. Key Contributions

1. Distributionally-Grounded Energy: Introduction of a KL-minimal energy model over feature units that re-conditions diffusion generations, converting "faithfulness" from an abstract aspiration into a well-defined optimization over probability measures.
2. Tool-Integrated Self-Evolution: A framework where the self-improvement loop is explicitly guided by an energy term, ensuring that policy updates prioritize faithful behavior alongside answer correctness.
3. Verifier-Plus-Repair Loop: A mechanism that provides trajectory-level feedback. When a step is flagged as unfaithful, the system re-conditions the energy and re-samples, effectively closing the loop between distributional constraints and explicit reasoning verification.
4. Theoretical Complementarity: Formal proof (Theorem 3.1) that neither the energy model nor the verifier alone is sufficient; the energy prevents distribution shift, while the verifier polices the stated reasoning. Only their combination ensures both distributional and trajectory faithfulness.

4. Experimental Results

The method was evaluated on GeoQA, SciVis, VQA-v2, and an in-house Multimodal Reasoning (MMR) set.

• Accuracy: DiffuseAgent-MI outperformed prior self-evolving agents, improving accuracy by up to 5.1 points on GeoQA (83.4% vs. 78.3% for the strongest baseline) and 4.1 points on MMR. Gains were most significant in tasks requiring multi-step spatial and causal reasoning.
• Faithfulness Metrics:
  • MI-Faith (Alignment AUROC): Improved from 75.4% (Agent0-VL) to 94.7%.
  • MI-Interp (Human Interpretability Agreement): Improved from 70.5% to 88.9%, more than doubling the interpretability of the ReAct baseline.
• Efficiency: The relative compute increase over a plain agent was minimal (1.12×), as the energy term acts as a lightweight score correction and the verifier is a small classifier.
• Ablation Studies:
  • Removing the energy term collapsed faithfulness metrics (MI-Faith dropped to 71.6%) while leaving accuracy largely intact.
  • Removing the verifier or repair branch degraded trajectory-level grounding, confirming that distributional constraints alone cannot police the stated reasoning.

5. Significance and Claims

The paper claims that DiffuseAgent-MI represents a significant step toward reliable and explainable vision agents by unifying mechanistic interpretability with agent self-improvement.

• Bridging the Gap: It addresses the "unfaithfulness" problem where correct answers are derived from unfaithful internal processes, a critical issue in safety-critical applications.
• Operationalizing Faithfulness: By treating faithfulness as a distributional optimization problem (KL-minimal soft constraint) rather than a post-hoc explanation, the method provides a principled way to align internal representations with external explanations.
• Complementarity: The work demonstrates that distributional grounding (via energy) and trajectory verification are complementary; one ensures the sample is valid relative to the prior, while the other ensures the reasoning narrative matches the actual computation.

The authors note limitations, specifically that guarantees are relative to the chosen feature-unit space and the verifier's definition of active units. If the underlying feature extractor is misspecified, the guarantees inherit that error. Future work is suggested to extend this to continuously learned, task-specific units and to autoregressive generation models.