🤖 machine learning

FedRef: Bayesian Fine-Tuning using a Reference Model to Mitigate Catastrophic Forgetting for Heterogeneous Federated Learning

FedRef is a Bayesian fine-tuning framework for heterogeneous federated learning that utilizes a server-side reference model and MAP-based regularization to mitigate catastrophic forgetting and reduce client-side computational overhead.

Original authors: Taehwan Yoon, Bongjun Choi, Wesley De Neve

Published 2026-04-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Taehwan Yoon, Bongjun Choi, Wesley De Neve

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are part of a massive, global cooking competition. Instead of everyone being in one kitchen, every chef is in their own home, using their own secret ingredients (this is Federated Learning—keeping data private by training locally).

Every few hours, all the chefs send their recipes to a Head Chef (the Server) to create one "Master Recipe."

The Problem: The "Identity Crisis" (Catastrophic Forgetting)

In a perfect world, the Master Recipe would just get better and better. But in reality, things get messy.

Imagine Chef A is making spicy pasta, and Chef B is making sweet dessert. When the Head Chef tries to combine them, the Master Recipe might suddenly become "Spicy Dessert"—a weird, confusing mess that forgets how to make good pasta and good dessert. This is Catastrophic Forgetting. Because every chef has different ingredients (Data Heterogeneity), the Master Recipe starts "drifting" wildly, losing its way and becoming a chaotic jumble of flavors. This is Unbounded Drift.

Existing solutions try to fix this by telling every individual chef, "Hey, don't stray too far from the Master Recipe!" But this makes the home chefs work much harder because they have to constantly check their progress against the Master Recipe while they are already busy cooking.

The Solution: FedRef (The "Golden Memory" Method)

The authors of this paper created FedRef. Instead of making the home chefs do extra work, they give the Head Chef a superpower: A Reference Model.

Think of the Reference Model as a "Golden Memory" or a "Best-of" compilation of all the successful Master Recipes from the past few rounds.

Here is how FedRef works using a simple analogy:

The Home Chefs (Clients): They just cook their dishes normally. They don't have to worry about extra rules or checking the Master Recipe. This keeps their "kitchen workload" very low.
The Head Chef (Server): When the Head Chef receives the new recipes, they don't just blindly mix them. Instead, they look at the Golden Memory (the Reference Model).
The Calibration: The Head Chef says, "I see these new recipes are very wild, but let's nudge them slightly back toward our Golden Memory so we don't forget the great things we learned yesterday."

By using this "Bayesian Fine-Tuning" (a fancy way of saying "making an educated guess based on past success"), the Head Chef acts like a stabilizer. They pull the recipe back toward a "safe zone" of flavors that have worked before.

Why is this better?

It’s Easier on the Workers: The home chefs (the devices like your phone or a hospital's computer) don't have to do extra math. They just cook and send the results.
It’s More Stable: Because the Head Chef is constantly anchoring the recipe to a "Golden Memory," the Master Recipe doesn't go haywire. It learns new things without "forgetting" the old ones.
It Works in Tough Conditions: Even if one chef is using weird, non-standard ingredients (like highly specialized medical data), the Master Recipe stays on track.

In short: FedRef is like a smart editor who looks at a bunch of messy drafts and, instead of just smashing them together, uses a "best-of" version of previous drafts to guide the final book into something coherent and brilliant.

Technical Summary: FedRef (Bayesian Fine-Tuning using a Reference Model)

1. Problem Statement

Federated Learning (FL) faces significant challenges due to data and system heterogeneity (non-IID data). These challenges manifest in two primary ways:

Catastrophic Forgetting: As the global model is updated sequentially across different client distributions, it tends to lose the knowledge acquired from previous distributions.
Unbounded Drift: The divergence between local client updates and the global objective leads to unstable model trajectories, which can degrade predictive performance and increase the risk of model-poisoning attacks.
Client-Side Constraints: Existing regularization methods (like FedProx or FedOpt) require clients to compute additional proximal terms, increasing their computational and energy overhead—a critical bottleneck for resource-constrained edge devices.

2. Methodology

The authors propose FedRef, a framework that shifts the burden of regularization from the clients to the server using Bayesian Fine-Tuning and a Reference Model.

Reference Model Construction: Instead of relying only on the most recent global model, the server constructs a reference model ( $\theta_{ref}$ ) by calculating a temporally weighted aggregate of the $\rho$ previous global models. This ensures the reference model captures a stable, smoothed representation of the model's historical trajectory.
MAP-based Regularization: FedRef treats the global model update as a Maximum A Posteriori (MAP) estimation problem. It integrates two types of regularization:
1. Temporal Consistency: An $L_2$ penalty on the difference between successive global models ( $\|\theta_g^r - \theta_g^{r-1}\|^2$ ).
2. Reference Anchoring: An $L_2$ penalty that keeps the global model close to the temporally smoothed reference model ( $\|\theta_{ref}^r - \theta_{ref}^{r-1}\|^2$ ).
Server-Side Optimization: Unlike FedProx, which forces clients to compute proximal terms, FedRef clients only perform standard local training. The server performs the complex Bayesian fine-tuning step (minimizing the combined loss of aggregated client losses and the reference constraints) to produce the updated global model.
Adaptive Scheduling: The strength of the reference-based regularization ( $\lambda_{ref}$ ) is adaptively increased over communication rounds to balance initial exploration with long-term stability.

3. Key Contributions

Mitigation of Unbounded Drift: The authors mathematically demonstrate that the reference-model prior acts as a "restoring force," reducing the Unbounded Drift Probability (UDP) compared to FedAvg.
Computational Efficiency: By moving the fine-tuning operations to the server, FedRef maintains the same client-side computational complexity as the baseline FedAvg, avoiding the 1–3% overhead seen in FedProx/FedOpt.
Robustness to Heterogeneity: The framework is specifically designed to stabilize the global model trajectory in highly non-IID environments (e.g., medical imaging).

4. Experimental Results

The method was evaluated on image classification (FEMNIST, CINIC-10) and medical image segmentation (FeTS2022):

Predictive Performance: FedRef achieved superior accuracy and F1-scores, particularly in the highly heterogeneous FeTS2022 dataset.
Catastrophic Forgetting: Using a custom metric $\zeta$ (which balances performance degradation and improvement), FedRef demonstrated significantly lower forgetting scores than FedAvg, FedProx, and FedOpt.
Convergence Speed: FedRef reached performance targets (accuracy and segmentation quality) at a rate comparable to or faster than existing methods, showing that regularization does not necessarily slow down training.
Stability: The global parameter trajectory was more monotonic and less oscillatory than baseline methods.

5. Significance

FedRef represents a paradigm shift in how regularization is applied in Federated Learning. By leveraging Bayesian Transfer Learning principles at the server level, it solves the dual problem of model instability and client-side resource exhaustion. This makes it a highly practical framework for real-world, large-scale FL applications—such as medical collaborations or mobile device networks—where data is highly skewed and client devices have limited processing power.