---
title: 'DFAMS: Dynamic-Flow Federated Alignment Search'
url: https://www.emergentmind.com/topics/dfams
type: topic
---

# DFAMS: Dynamic-Flow Federated Alignment Search

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 [2508.20353]. It is designed to improve retrieval quality for **ambiguous queries**, particularly in **cross-domain scenarios**, by exploiting **dynamic information flow (DIF)** inside large language models to identify latent query intents and to construct semantically aligned knowledge partitions for routing and retrieval [2508.20353]. 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 [2508.20353].

## 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** [2508.20353]. 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** [2508.20353]. 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 [2508.20353].

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 [2508.20353]. 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 [2508.20353].

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 [2508.20353].

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$, the framework constructs a DIF representation by concatenating the activations of these selected neuron groups:
$$
\mathbf{z} = \mathrm{Concat}(\{h_{\ell_i}^{(g_i)}\}_{i=1}^{P}).
$$
The resulting vector $\mathbf{z}$ is described as a compact, task-relevant embedding that captures both **semantic intent** and **knowledge partition cues** learned internally by the LLM [2508.20353].

## 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 [2508.20353].

The DIF embedding $\mathbf{z}$ is first projected through a learned alignment function:
$$
\mathbf{r} = g_{\text{align}}(\mathbf{z}).
$$
Training then proceeds with two complementary contrastive objectives [2508.20353].

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:
$$
\mathcal{L}_{\text{CL}} = -\sum_{i} \frac{1}{|P(i)|} \sum_{p \in P(i)}
\log \frac{\exp(\mathbf{r}_i^\top \mathbf{r}_p/\tau_{cl})}
{\sum_{a \in A(i)} \exp(\mathbf{r}_i^\top \mathbf{r}_a/\tau_{cl})}.
$$
Here, $P(i)$ denotes positives with the same KB label, $A(i)$ the candidate set, and $\tau_{cl}$ a temperature parameter [2508.20353].

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 $\{\boldsymbol{\mu}_m\}_{m=1}^{M}$ that represent subdomains. The loss is
$$
\mathcal{L}_{\text{PCL}} = -\sum_i \frac{1}{|C(i)|} \sum_{m \in C(i)}
\log \frac{\exp(\mathbf{r}_i^\top \boldsymbol{\mu}_m/\tau_{pcl})}
{\sum_{j \in AC(i)} \exp(\mathbf{r}_i^\top \boldsymbol{\mu}_j/\tau_{pcl})}.
$$
In this formulation, $C(i)$ denotes the most similar prototypes to $\mathbf{r}_i$, $AC(i)$ the others, and $\tau_{pcl}$ the temperature [2508.20353].

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 [2508.20353]. 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 [2508.20353]. For a new query embedding $\mathbf{q}$, the system computes similarity scores to all prototypes,
$$
s_i = \mathrm{sim}(\mathbf{q}, \boldsymbol{\mu}_i).
$$

Routing is then governed by two mechanisms. The first is **adaptive triggering**: if the maximum prototype similarity does not exceed a threshold $\tau$, the system **abstains from retrieval** [2508.20353]. 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-$N$ prototypes and allocates document budget across knowledge bases in proportion to the summed prototype similarities:
$$
w_k = \left\lfloor \frac{\sum_{i \in \mathcal{I},\,k_i=k} s_i}
{\sum_{k'} \sum_{i \in \mathcal{I},\,k_i=k'} s_i} \cdot T \right\rfloor.
$$
Here, $\mathcal{I}$ is the set of selected prototype indices, $k_i$ is the knowledge-base identity associated with prototype $i$, and $T$ is the total retrieval budget [2508.20353].

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 [2508.20353]. 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 [2508.20353].

| Component | Input | Operation |
|---|---|---|
| DIF Extraction Module | Query $x$ | Selects relevant activations via gradient and Shapley analysis and outputs DIF embedding $\mathbf{z}$ |
| Alignment & Prototyping Module | $\mathbf{z}$ | 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 [2508.20353].

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** [2508.20353]. 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** [2508.20353]. 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** [2508.20353].

A sample result table reported for benchmark-level performance is as follows [2508.20353]:

| 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 [2508.20353].

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 [2508.20353]. 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 [2508.20353].

## 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** [2508.20353].

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** [2508.20353]. 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 [2508.20353].

The principal scope of DFAMS, as presented, is **complex FR scenarios** involving **ambiguous**, **cross-domain**, or **semantically compressed** queries [2508.20353]. 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 [2508.20353].

Source: https://www.emergentmind.com/topics/dfams