DFAMS: Dynamic-Flow Federated Alignment Search

• DFAMS is a federated retrieval framework that uses dynamic information flow (DIF) to detect latent query intent and guide source selection.
• It employs Shapley value attribution and multi-prototype contrastive learning to align LLM internal representations with distributed knowledge bases.
• Experimental results show improvements up to 14.37% in classification accuracy, 5.38% in retrieval recall, and 6.45% in downstream QA accuracy over previous methods.

Searching arXiv for the DFAMS paper and closely related federated retrieval work. DFAMS, short for Dynamic-flow guided Federated Alignment based Multi-prototype Search, is a framework for Federated Retrieval (FR) in retrieval-augmented generation settings where external knowledge is distributed across multiple heterogeneous knowledge sources rather than concentrated in a single corpus (Yang et al., 28 Aug 2025). It is designed to improve retrieval quality for ambiguous queries, particularly in cross-domain scenarios, by exploiting dynamic information flow (DIF) inside LLMs to identify latent query intents and to construct semantically aligned knowledge partitions for routing and retrieval (Yang et al., 28 Aug 2025). The framework combines gradient signals from a few annotated queries, Shapley value-based attribution, an alignment module trained via multi-prototype contrastive learning, and an adaptive prototype-guided routing mechanism. Experimental results reported for five benchmarks show gains of up to 14.37% in knowledge classification accuracy, 5.38% in retrieval recall, and 6.45% in downstream QA accuracy over advanced FR methods (Yang et al., 28 Aug 2025).

1. Problem setting and motivation

Federated Retrieval is motivated by the observation that, in realistic retrieval-augmented generation deployments, the necessary evidence for answering a query may reside in multiple external knowledge sources, including medicine, enterprise, and science, and these sources may also be privacy-protected (Yang et al., 28 Aug 2025). In this setting, the retrieval problem is not limited to document ranking within one collection; it also includes source selection, source allocation, and the decision of whether retrieval is necessary at all.

DFAMS is introduced to address several limitations of prior FR approaches. The reported difficulties include ambiguous or compressed user queries that do not align well with knowledge-base structure, semantic ambiguity and latent user intent that are difficult to capture with surface-level rewrites or standard embedding similarity, and the complexity of cross-domain source selection and multi-source aggregation (Yang et al., 28 Aug 2025). The framework therefore targets the routing stage of FR as a representation-learning problem: instead of relying only on externally observable query text, it attempts to use internal LLM computation patterns as signals for intent recognition and knowledge partitioning.

This suggests a shift in emphasis from conventional retrieval heuristics toward model-internal routing cues. A plausible implication is that DFAMS treats the LLM not only as the downstream generator, but also as a structured sensor for latent knowledge demand.

2. Dynamic Information Flow as the basis for routing

The central organizing concept in DFAMS is Dynamic Information Flow (DIF). The framework is motivated by the hypothesis that LLMs dynamically activate different neural pathways depending on query semantics, intent, and knowledge requirements, and that these pathways encode signals relevant for federated routing but are not directly surfaced by standard embeddings or prompt engineering (Yang et al., 28 Aug 2025).

To operationalize this hypothesis, DFAMS begins with a probing dataset constructed from a small number of annotated queries. These queries are crafted with fixed instructions to explicitly elicit domain selection, with the goal of isolating domain-routing behavior for attribution (Yang et al., 28 Aug 2025). The probing dataset is written as

$$\mathcal{D}_{\text{probe}} = \{(x_i, \mathcal{K}_i)\}_{i=1}^{n_{\text{probe}}}.$$

Here, each $x_i$ is a crafted query and $\mathcal{K}_i$ is the correct knowledge-base label (Yang et al., 28 Aug 2025).

For neuron-level attribution, DFAMS uses a Shapley value-based attribution that combines gradient and Hessian information for the domain-selection loss. The importance score for parameter or neuron $\theta_j$ is given as

$$\phi_j = -g_j^{(\gamma)} \theta_j - \frac{1}{2} \omega_{jj}^{(j)} \theta_j^2 H_{jj}^{(\gamma)} - \frac{1}{2} \theta_j \sum_{k \ne j} \omega_{jk}^{(\mathcal{S})} H_{jk}^{(\gamma)} \theta_k.$$

In this expression, $g_j^{(\gamma)}$ denotes the supervised loss gradient, $H_{jk}^{(\gamma)}$ the Hessian entry, and $\omega_{jj}^{(j)}$, $\omega_{jk}^{(\mathcal{S})}$ coefficients for self and pairwise contributions (Yang et al., 28 Aug 2025).

After attribution, DFAMS selects the top $T$ layers by aggregate attribution and then the top-N neuron groups consisting of adjacent units with high attribution. For any input query $x_i$0, the framework constructs a DIF representation by concatenating the activations of these selected neuron groups:

$x_i$1

The resulting vector $x_i$2 is described as a compact, task-relevant embedding that captures both semantic intent and knowledge partition cues learned internally by the LLM (Yang et al., 28 Aug 2025).

3. Alignment and multi-prototype representation learning

The second major component of DFAMS is an alignment module trained with multi-prototype contrastive learning. Its purpose is to map the internal DIF space onto the external structure of federated knowledge bases, especially when source boundaries are overlapping, ambiguous, or fine-grained (Yang et al., 28 Aug 2025).

The DIF embedding $x_i$3 is first projected through a learned alignment function:

$x_i$4

Training then proceeds with two complementary contrastive objectives (Yang et al., 28 Aug 2025).

The first is an inter-KB contrastive loss, intended to pull together query representations associated with the same knowledge base and push apart those associated with different knowledge bases:

$x_i$5

Here, $x_i$6 denotes positives with the same KB label, $x_i$7 the candidate set, and $x_i$8 a temperature parameter (Yang et al., 28 Aug 2025).

The second is an intra-KB prototype contrastive loss. Each knowledge base is divided into multiple clusters by K-means, producing a set of prototypes $x_i$9 that represent subdomains. The loss is

$\mathcal{K}_i$0

In this formulation, $\mathcal{K}_i$1 denotes the most similar prototypes to $\mathcal{K}_i$2, $\mathcal{K}_i$3 the others, and $\mathcal{K}_i$4 the temperature (Yang et al., 28 Aug 2025).

The role of this dual objective is explicit in the reported interpretation: inter-KB alignment is used for correct source separation, while intra-KB alignment models the subdomain structure within each source (Yang et al., 28 Aug 2025). This suggests that DFAMS does not treat each knowledge base as semantically homogeneous. A plausible implication is that the use of multiple prototypes provides robustness to ambiguous or compositional queries that partially match several subdomains or sources.

4. Routing, triggering, and retrieval allocation

DFAMS uses the learned prototype structure at inference time through an adaptive prototype-guided routing module (Yang et al., 28 Aug 2025). For a new query embedding $\mathcal{K}_i$5, the system computes similarity scores to all prototypes,

$\mathcal{K}_i$6

Routing is then governed by two mechanisms. The first is adaptive triggering: if the maximum prototype similarity does not exceed a threshold $\mathcal{K}_i$7, the system abstains from retrieval (Yang et al., 28 Aug 2025). This is intended to avoid unnecessary retrieval for cases in which external knowledge is not confidently required.

The second is semantic routing. When retrieval is triggered, DFAMS selects the top-$\mathcal{K}_i$8 prototypes and allocates document budget across knowledge bases in proportion to the summed prototype similarities:

$\mathcal{K}_i$9

Here, $\theta_j$0 is the set of selected prototype indices, $\theta_j$1 is the knowledge-base identity associated with prototype $\theta_j$2, and $\theta_j$3 is the total retrieval budget (Yang et al., 28 Aug 2025).

This routing design supports three behaviors stated in the paper summary. First, it can trigger multiple sources when a query is semantically multi-domain. Second, it can select no retrieval when retrieval is not warranted. Third, it supports fine-grained source allocation rather than a hard single-source decision (Yang et al., 28 Aug 2025). In the context of FR, this is significant because the routing problem is not only about relevance ranking but about partition-aware evidence acquisition under uncertainty.

5. Architecture and operational pipeline

DFAMS is described as a four-part architecture (Yang et al., 28 Aug 2025).

Component	Input	Operation
DIF Extraction Module	Query $\theta_j$4	Selects relevant activations via gradient and Shapley analysis and outputs DIF embedding $\theta_j$5
Alignment & Prototyping Module	$\theta_j$6	Projects to semantic space and trains with joint inter- and intra-KB contrastive losses
Adaptive Prototype-Guided Routing Module	Query embedding	Computes similarities, decides whether to retrieve, and allocates quota among KBs
Retrieval & Generation Module	Routed retrieval results	Retrieves documents and feeds context to the LLM for final answer generation

Operationally, the framework begins by extracting a DIF embedding from the LLM using the preselected neuron groups. It then maps that representation into an aligned semantic space, where multiple prototypes model both knowledge-base identity and subdomain structure. Finally, routing converts prototype-level similarity into retrieval triggering and cross-source quota allocation, after which standard retrieval and generation proceed (Yang et al., 28 Aug 2025).

Within this pipeline, the distinctive methodological claim is that knowledge partitioning is aligned to the neural substrates the LLM itself uses for routing and intent recognition (Yang et al., 28 Aug 2025). This is presented as a way to obtain partitions that are simultaneously internally consistent and semantically meaningful. A plausible implication is that DFAMS attempts to reduce the mismatch between external corpus organization and internal model representations.

6. Empirical results, benchmarks, and interpretation

The reported experimental evaluation covers five benchmarks spanning in-domain settings (Wiki, Med, PEP) and out-of-domain settings (MMLU, MIRAGE), using multiple open-source LLMs including Qwen2.5-7B and LLaMA3.1-8B (Yang et al., 28 Aug 2025). The main aggregate findings are that DFAMS improves over prior FR methods by up to 14.37% in knowledge base classification accuracy, up to 5.38% in retrieval recall, and up to 6.45% absolute gain in end-to-end QA accuracy (Yang et al., 28 Aug 2025).

A sample result table reported for benchmark-level performance is as follows (Yang et al., 28 Aug 2025):

Method	Wiki Cls Acc	Med Cls Acc	PEP Cls Acc
Prompt	32.47	48.82	66.86
RAGRoute	76.07	69.04	51.47
DFAMS	85.03	71.81	82.85

The same summary also reports MMLU QA Acc = 86.17 and MIRAGE QA Acc = 79.88 for DFAMS, compared with 80.77 and 69.17 for Prompt, and 80.20 and 69.77 for RAGRoute (Yang et al., 28 Aug 2025).

The ablation findings are equally central to the interpretation of the method. Removing DIF-based modeling or multi-prototype alignment is reported to substantially reduce both partition accuracy and retrieval performance, and both the inter-KB and intra-KB contrastive objectives are described as crucial: the former for correct source selection and the latter for high-quality within-source retrieval (Yang et al., 28 Aug 2025). These results support the intended decomposition of the problem into coarse routing and fine-grained semantic localization.

The efficiency discussion is more cautious but notable. The framework is reported to be faster or comparable to multi-source baselines at inference despite querying the LLM at routing time, because it reduces redundant retrieval and can abstain from retrieval when appropriate (Yang et al., 28 Aug 2025).

7. Significance, scope, and interpretive issues

DFAMS occupies a specific position within federated retrieval research: it is neither a pure retrieval model nor merely a prompt-based router. Its defining claim is that latent query intent and knowledge-domain cues can be recovered from dynamic information flow inside LLMs, and that these cues can be converted into a usable representation for knowledge partitioning, routing, and source allocation (Yang et al., 28 Aug 2025).

A common misconception would be to equate DFAMS with a standard dense retrieval method augmented by clustering. The reported design is more specific: the clustering and prototypes are not learned directly from raw text embeddings alone, but from an alignment space derived from DIF embeddings selected through gradient signals and Shapley value-based attribution (Yang et al., 28 Aug 2025). Another possible misconception would be to treat the system as a hard classifier over knowledge bases. In fact, the routing rule explicitly allows multiple-source triggering and retrieval abstention, which makes the framework closer to a confidence-aware resource allocator than to a single-label router (Yang et al., 28 Aug 2025).

The principal scope of DFAMS, as presented, is complex FR scenarios involving ambiguous, cross-domain, or semantically compressed queries (Yang et al., 28 Aug 2025). The article’s experimental claims are confined to five benchmarks and the listed open-source LLMs, so broader generalization beyond those conditions should be treated as an inference rather than a reported fact. Still, the combination of DIF-based probing, contrastive alignment, and prototype-guided routing indicates a broader methodological direction in which internal LLM representations become first-class signals for retrieval-system design.

In summary, DFAMS is a federated retrieval framework that uses LLM-internal dynamic information flow to build a semantically aligned, multi-prototype routing space for distributed knowledge access. Its contribution lies in connecting interpretability-style neuron attribution, contrastive representation learning, and adaptive multi-source retrieval, with reported improvements in classification, recall, and downstream question answering in heterogeneous FR settings (Yang et al., 28 Aug 2025).