Atomic Information Flow: A Network Flow Model for Tool Attributions in RAG Systems

TL;DR

This paper introduces Atomic Information Flow (AIF), a graph-based model that decomposes tool outputs and LLM calls into atomic units to enable precise attribution and explainability in complex RAG systems.

Contribution

We propose AIF, a novel network flow model that decomposes information into atoms for granular attribution and train a lightweight model to approximate critical information flow.

Findings

AIF enables granular attribution metrics for RAG systems.

Training on AIF signals significantly improves the accuracy of a small language model in identifying critical information.

The AIF approach bridges the performance gap between small and large language models in information attribution.

Abstract

Many tool-based Retrieval Augmented Generation (RAG) systems lack precise mechanisms for tracing final responses back to specific tool components -- a critical gap as systems scale to complex multi-agent architectures. We present \textbf{Atomic Information Flow (AIF)}, a graph-based network flow model that decomposes tool outputs and LLM calls into atoms: indivisible, self-contained units of information. By modeling LLM orchestration as a directed flow of atoms from tool and LLM nodes to a response super-sink, AIF enables granular attribution metrics for AI explainability. Motivated by the max-flow min-cut theorem in network flow theory, we train a lightweight Gemma3 (4B parameter) language model as a context compressor to approximate the minimum cut of tool atoms using flow signals computed offline by AIF. We note that the base Gemma3-4B model struggles to identify critical information with \textbf{54.7\%} accuracy on HotpotQA, barely outperforming lexical baselines (BM25). However, post-training on AIF signals boosts accuracy to \textbf{82.71\%} (+28.01 points) while achieving \textbf{87.52\%} (+1.85\%) context token compression -- bridging the gap with the Gemma3-27B variant, a model nearly $7\times$ larger.