MATA: A Trainable Hierarchical Automaton System for Multi-Agent Visual Reasoning

TL;DR

MATA tackles the opaque and hallucination-prone reasoning of vision-language models by casting visual reasoning as a hierarchical finite-state automaton whose top-level transitions are learned by a trainable hyper agent. Each top-level state corresponds to a small, rule-based agent, enabling reliable micro-control, while a shared memory records an auditable execution history. The hyper agent is supervised via a transition-trajectory dataset (MATA-SFT-90K) derived from transition trees, enabling robust cross-task policy learning and generalization across GQA, OK-VQA, and referring expression benchmarks. Across multiple datasets, MATA achieves state-of-the-art performance by enabling collaboration and competition among complementary agents, reducing hallucinations and providing interpretable reasoning traces that can transfer across domains. The work advances interpretable, modular visual reasoning with practical efficiency by combining a learned high-level policy with deterministic, verifiable agent behavior, and provides datasets and code for further research.

Abstract

Recent vision-language models have strong perceptual ability but their implicit reasoning is hard to explain and easily generates hallucinations on complex queries. Compositional methods improve interpretability, but most rely on a single agent or hand-crafted pipeline and cannot decide when to collaborate across complementary agents or compete among overlapping ones. We introduce MATA (Multi-Agent hierarchical Trainable Automaton), a multi-agent system presented as a hierarchical finite-state automaton for visual reasoning whose top-level transitions are chosen by a trainable hyper agent. Each agent corresponds to a state in the hyper automaton, and runs a small rule-based sub-automaton for reliable micro-control. All agents read and write a shared memory, yielding transparent execution history. To supervise the hyper agent's transition policy, we build transition-trajectory trees and transform to memory-to-next-state pairs, forming the MATA-SFT-90K dataset for supervised finetuning (SFT). The finetuned LLM as the transition policy understands the query and the capacity of agents, and it can efficiently choose the optimal agent to solve the task. Across multiple visual reasoning benchmarks, MATA achieves the state-of-the-art results compared with monolithic and compositional baselines. The code and dataset are available at https://github.com/ControlNet/MATA.

Figures (6)

• Figure 1: Overview of MATA.(a) Linear pipelines (previous methods) execute modules in a fixed, manually designed order. (b) MATA organizes agents as states in a hyper automaton. A trainable hyper agent learns high-level transitions between agents (blue arrows), enabling collaboration and competition, while each agent runs a small rule-based sub-automaton for reliable micro-control (black arrows). (c) To train the hyper agent, we expand a transition-trajectory tree per image-query, score the leaves using task metrics, and convert each node’s snapshot into a supervised pair current memory $\rightarrow$ best next state for supervised finetuning (SFT), forming MATA-SFT-90K.
• Figure 2: Pipeline of MATA. A trainable hyper agent reads a snapshot of the shared memory, predicts the next state with an LLM State Controller. Its decision (blue arrows) routes control among agent states in the hyper automaton: Oneshot Reasoner, Stepwise Reasoner, and Specialized Agent. Each agent runs a rule-based sub-automaton that iterates until return to the hyper automaton. All agents read/write an append-only Shared Memory, enabling the hyper agent to access the current context for choosing the optimal next state. Lifecycle states Initial and Failure are shown outside the agents (see \ref{['sec:hyper_automaton']} for details).
• Figure 3: Results of different LLM sizes. Accuracy versus the model size (in billions of parameters) of the hyper agent's LLM state controller. Left: GQA; right: OK-VQA. X-axis: LLM size; Y-axis: accuracy.
• Figure 4: Results of different numbers of sub-agents. X-axis: number of sub-agents; Y-axis: accuracy in GQA.
• Figure 5: Details of agents in MATA. Each block shows the sub-automaton executed when the hyper automaton transits into that agent. Black arrows indicate the normal paths; red arrows show local error-correction paths. Persistent failures transition to Failure state of the hyper automaton (omitted for clarity).