Technical Summary: Federated Nested Learning (FedNL)

Problem Statement

Federated Learning (FL) for Large Language Models (LLMs) faces two persistent, intertwined challenges: data heterogeneity (Non-IID) and long-tail distributions. In realistic scenarios, client data distributions are highly skewed (e.g., medical vs. coding data), and critical knowledge often resides in the long tail, which is easily overshadowed by head classes during global aggregation.

Existing approaches primarily focus on regularizing or personalizing static model weights (e.g., FedProx, FedSSI). While effective at mitigating catastrophic forgetting, these methods treat the global model as a container of static knowledge. Consequently, when deployed to a client with unseen, highly heterogeneous data, the model lacks test-time plasticity—the ability to adapt to the current context without gradient updates. This leads to suboptimal performance on domain-specific tasks where nuances depend on context.

Methodology: Federated Nested Learning (FedNL)

The paper proposes Federated Nested Learning (FedNL), a framework that reframes FL from aggregating static knowledge to collaboratively learning optimization rules. Drawing inspiration from Nested Learning (NL) theory, FedNL treats inference not as a passive retrieval process but as a high-frequency "inner-loop" training process where the model actively compresses context into a transient memory state.

Core Architecture

FedNL reformulates FL as a three-level nested optimization system:

1. Level 0 (Outer Loop - Server): Collaborative Generalization. The server aggregates meta-rules (parameters $\theta$) rather than memory content. These rules govern how memory is constructed and updated.
2. Level 1 (Intermediate Loop - Client): Rule Learning. During local training, clients optimize $\theta$ (e.g., LoRA adapters, gating mechanisms) to ensure the inner loop produces useful memory states for downstream tasks. This involves backpropagating through the memory trajectory.
3. Level 2 (Inner Loop - Client): Test-Time Adaptation. During inference, clients use the global rules $\theta$ to update a private, context-aware memory state $S_t$ online via a Delta Rule. This state is transient and never leaves the device.

Technical Implementation

• Titans-Based Linear Attention: FedNL leverages the Titans architecture, which utilizes a linearized attention mechanism. The memory state $S_t$ is updated via a Delta Rule:
  $$S_t = S_{t-1} + \eta(v_t - S_{t-1}k_t)k_t^\top$$
  This update is mathematically equivalent to a single step of gradient descent on an associative memory objective, effectively treating inference as an online optimization process.
• Decoupling Knowledge and Rules: The model consists of a frozen backbone (providing general linguistic capabilities) and trainable meta-parameters ($\theta$). The meta-parameters include Low-Rank Projections (LoRA) for query/key/value generation and a Memory Gating mechanism ($\alpha$) to balance static and dynamic attention.
• Efficiency: By aggregating only the small set of meta-rules (adapters) rather than full model weights or memory states, FedNL drastically reduces communication overhead. Furthermore, the linear attention mechanism maintains constant inference memory complexity ($O(1)$) regardless of sequence length, unlike standard attention which requires $O(T)$ KV cache.

Key Contributions

1. Theoretical Reframing: The paper introduces the Nested Learning perspective to FL, formalizing the problem as the collaborative training of optimization rules rather than static representations. This provides a theoretical basis for addressing Non-IID issues via test-time adaptation.
2. The FedNL Framework: A practical algorithm integrating Titans-based linear attention into the FL pipeline. By treating the Delta Rule as an online gradient descent step, it enables Zero-Shot Test-Time Adaptation on unseen domains without heavy computational lifting.
3. Empirical Superiority: Experiments demonstrate that FedNL achieves competitive performance in short-context reasoning and superior gains in long-context retrieval and streaming diagnostics compared to strong federated baselines.

Experimental Results

The authors evaluated FedNL on two federated settings: a five-client Non-IID MMLU split and a seven-client Non-IID Needle In A Haystack (NIAH) long-context benchmark.

• Non-IID MMLU (Short-Context Reasoning):

  • On the Qwen2.5-1.5B backbone, FedNL achieved the highest average accuracy (58.88%), slightly outperforming the best baseline, FedSSI (58.70%).
  • Significant gains were observed in STEM (+2.0 pp) and Math/CS (+3.8 pp) domains.
  • On the Llama-3.2-1B backbone, FedNL reached 42.64% vs. 42.10% for FedSSI.
  • Crucially, FedNL minimized the "aggregation drop" (performance loss when moving from local to federated models), keeping it near zero across domains, whereas static baselines showed significant degradation.

• Long-Context Retrieval (NIAH):

  • FedNL achieved the highest average accuracy (29.7%) across seven retrieval templates, outperforming FedALA (28.6%).
  • Gains were most pronounced in MK-NIAH, MV-NIAH, and UUID code tasks.
  • Streaming Cross-Entropy (CE) Diagnostics: In a 16K token streaming probe, FedAvg's normalized CE increased by 3.1% as the context unfolded (accumulating uncertainty), whereas FedNL decreased by 2.1%, indicating the recurrent memory state successfully absorbed useful context over long streams.
  • Memory Efficiency: FedNL maintained constant inference memory, avoiding Out-Of-Memory errors at 16K tokens that affected static-attention baselines.

• Communication Efficiency:

  • FedNL reduced per-round communication by approximately 350x compared to FedAvg (0.26 MB vs. 22.5 MB for the trainable parameters in the tested setup).

Significance and Claims

The paper claims that FedNL establishes a new direction for Federated Learning where models evolve from a repository of static knowledge into a paradigm of continuous, context-aware learning.

• Paradigm Shift: By decoupling general linguistic capabilities (frozen backbone) from memory construction rules (trainable adapters), FedNL allows diverse clients to share "learning capabilities" rather than conflicting "knowledge states."
• Robustness to Heterogeneity: The global meta-parameters do not need to encode the conflict between disparate domains (e.g., Code vs. Medical). Instead, they encode a universal rule: "If context involves X, update memory to store Y." This alleviates the Non-IID conflict in the parameter space.
• Zero-Shot Adaptation: The framework enables models to adapt to specific local distributions during the first few tokens of inference purely by executing learned memory update rules, solving the "cold-start" problem often faced by Test-Time Training (TTT) methods.

The authors note that while results are consistent at the 1B–1.5B scale, extending FedNL to larger foundation models remains a necessary next step to validate the generality of memory-rule aggregation.