Skip to content

A\(^2\)FM: An Adaptive Agent Foundation Model for Tool-Aware Hybrid Reasoning

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Unified agent foundation model bridging reasoning-only and tool-using LLM paradigms

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

A²FM unifies three execution modes—instant, reasoning, and agentic—within a single 32B backbone via a learned task-aware router, eliminating the longstanding split between reasoning-centric and tool-use-centric LLMs. Training follows a route-then-align SFT stage followed by Adaptive Policy Optimization (APO), a GRPO-based RL procedure with cost-regularized rewards. On a 32B backbone, A²FM achieves SOTA or near-SOTA on BrowseComp, AIME25, and HLE while cutting cost-of-pass by up to 45% versus single-mode baselines.

Problem

Reasoning-centric LLMs (e.g., o3, DeepSeek-R1) lack tool-use capability; agentic LLMs (e.g., Kimi K2, GLM-4.5) lag on multi-step logical tasks. Existing hybrid systems (e.g., Qwen3) train these capabilities in separate alignment stages with no unified inference-time controller, and neither family avoids wasteful computation on trivially simple queries. Prior "when-to-think" work operates text-only and ignores the cost of external tool invocations.

Method

Architecture: A shared Qwen2.5-32B-Instruct backbone with a self-adaptive router \(\pi_\text{route}(m \mid x)\) over \(\mathcal{M} = \{\text{instant}, \text{reasoning}, \text{agentic}\}\). Mode is emitted via <classification> tags; instant outputs <answer> directly, reasoning wraps CoT in <reasoning>, and agentic interleaves parallel <plan>, <tool_call> (MCP-format search/crawl/code tools), and <summary> tags with masked tool responses during training.

Stage 1 — Route-then-Align SFT: Distillation from mode-specialized teachers (DeepSeek-R1 for reasoning, DeepSeek-V3.1 for agentic/instant). Data curation uses difficulty-based sampling and classification-ambiguous query handling.

Stage 2 — Adaptive Policy Optimization (APO): Extends GRPO with: - Forced rollouts: \(\rho\) rollouts per mode via prefix injection, guaranteeing cross-mode exploration. - Adaptive rollouts: \(\gamma\) autonomous-routing samples that reward correct self-routing. - Composite reward: \(r_\text{total} = r_\text{acc} \times r_\text{adaptive} \times r_\text{format}\), where the adaptive term penalizes non-instant mode use on "easy" queries (those where instant mode empirically achieves accuracy \(\geq \tau\)): \(r_\text{adaptive} = 1 - p^\alpha\) if a costly mode is chosen, else \(1\). No KL divergence term; strictly on-policy.

Total rollout group size is \(G = 3\rho + \gamma\); at training time \(\rho = 3\), \(\gamma = 3\), \(\alpha = 2\).

Key Contributions

  • First hybrid model that jointly integrates task-aware routing and mode-specific trajectory generation under a unified backbone.
  • A two-stage training pipeline: (i) route-then-align SFT with mode-specialized teacher distillation; (ii) APO with adaptive sampling and multiplicative cost-regularized reward.
  • A third "instant mode" that bypasses both CoT and tool calls for simple queries, complementing the reasoning/agentic dichotomy.
  • Parallel agentic architecture with concurrent multi-tool execution via <plan>/<summary> scaffolding.
  • Empirical demonstration of Pareto-improving accuracy–efficiency trade-off vs. forced-mode baselines at 32B scale.

Results

  • BrowseComp: 13.4% (adaptive), 14.4% (forced agentic) — 2nd overall behind DeepDive (14.8%); surpasses OAgents (13.7%) and AFM-Search (11.1%).
  • GAIA: 60.7% (forced agentic), new SOTA; 57.3% (adaptive), second to OAgents (58.3%) which uses GPT-4.1.
  • XBench-DS: 56.0% (adaptive), highest among all baselines, +2.0 over AFM-Search (54.0%).
  • AIME24: 74.5%, new SOTA among 32B models; vs. ReTool 67.0%, AFM-Code 59.8%.
  • AIME25: 70.4%, second only to o1 (79.2%); vs. Claude 4 Sonnet 49.3%.
  • MATH500: 95.0%, comparable to o1 (96.4%).
  • HLE: 16.7% (adaptive), 20.6% (forced agentic); vs. GPT-4.1 5.8%, Qwen3-32B 8.3%.
  • SuperGPQA cost-of-pass: $0.00487/correct answer (adaptive) vs. $0.00889 (reasoning, −45.2%) and $0.00732 (agentic, −33.5%).

Limitations

  • Results are at 32B scale only; scaling behavior of APO and routing quality to smaller or larger models is not studied.
  • Forced-mode variants on some benchmarks (e.g., GAIA, BrowseComp) outperform the adaptive mode, indicating the router does not always make optimal decisions.
  • APO omits KL divergence regularization for training stability, potentially at the cost of divergence from the base model.
  • Evaluation uses a 103-item GAIA text-only subset and a 500-item HLE subset, which may not fully represent benchmark distributions.
  • Teacher distillation relies on closed or semi-open model outputs (DeepSeek-R1/V3.1); data construction details are largely deferred to appendices.
  • The adaptive reward threshold \(\tau\) and scaling factor \(\alpha\) are hyperparameters requiring tuning; their sensitivity across domains is not characterized.

Relevance to Agentic AI / LLM Agents

A²FM directly attacks a central architectural tension in agentic systems: how to decide, within a single model, when to reason internally versus invoke external tools versus respond immediately. The route-then-align + APO framework provides a principled RL recipe for this meta-decision, extending beyond prior "when-to-think" work (BPO, LHRM) by treating tool use as a first-class mode with its own cost model. The parallel agentic trajectory design (plan → parallel tool calls → dynamic summary) is a concrete template for high-efficiency multi-tool agents. For researchers building general-purpose LLM agents, the cost-of-pass framing and the multiplicative reward structure that jointly enforces correctness, mode appropriateness, and format compliance offer reusable design patterns.