📄 DFAMS: Dynamic-flow guided Federated Alignment based Multi-prototype Search

Summary

Federated Retrieval (FR) sends a query to several separate knowledge bases instead of one central index. Current FR systems often pick the wrong source or return weak documents when a query is ambiguous or crosses domains. DFAMS addresses this using the Dynamic Information Flow (DIF), meaning the activation paths that form inside an LLM during inference.

The method works in two stages:

• Probing: gradients from a small set of annotated queries feed a Shapley value-based attribution step, which finds the neurons tied to user intent and subdomain boundaries.
• Alignment and routing: the extracted DIF features train an Align-MLP with multi-prototype contrastive learning. This models subdomains within each source and aligns related content across sources. At inference, Adaptive Prototype-Guided Routing decides whether retrieval is needed and which sources to query.

The authors also built a harder FR benchmark. On five benchmarks, DFAMS beats prior FR methods in source classification, recall and QA accuracy. This matters wherever knowledge is split across silos that cannot simply be merged, such as medical records, textbooks and papers.

Key points

• Reported gains over strong FR baselines reach up to +14.37% knowledge classification accuracy, +5.38% retrieval recall and +6.45% downstream QA accuracy.
• Shapley heatmaps were computed for Qwen2.5-7B, LLaMA3.1-8B, Qwen2.5-3B and Qwen2.5-0.5B. They show two peaks: shallow layers carry intent signals and deep layers carry domain signals, so DIF is extracted from the deep layers.
• A 2000-dim DIF representation outperforms both random features of the same size and full-layer 3584-dim features, even without training. On Wiki the gain is +14.6% accuracy.
• In the Wiki ablation, the full model reaches 85.03 accuracy and 53.83 recall. Removing inter-KB alignment lowers accuracy by 9.14 points, and removing Semantic Routing lowers recall by 5.16 points. Three prototypes per class and top-T = 3 work best.
• The new benchmark covers knowledge-free queries, multi-fragment queries within one source, and cross-source queries. Code is available at github.com/Artessay/DFAMS.