🤖 machine learning

Personalized Additive Modeling for Multi-level Federated Learning

This paper proposes Federated Multi-level Additive Modeling (FeMAM), a dynamic and cost-efficient framework that addresses complex non-IID data heterogeneity in federated learning by constructing personalized predictors through the additive composition of global, subgroup, and client-specific models that adaptively grow and prune during training.

Original authors: Shutong Chen, Guodong Long, Tianyi Zhou, Jie Ma, Jing Jiang, Chengqi Zhang

Published 2026-06-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shutong Chen, Guodong Long, Tianyi Zhou, Jie Ma, Jing Jiang, Chengqi Zhang

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 trying to teach a group of 50 different people how to cook.

In a standard classroom (traditional AI), the teacher gives everyone the exact same recipe. If the students have very different tastes or ingredients, the recipe won't work well for anyone.

In a "Federated Learning" scenario, the students are in their own kitchens. They can't share their secret ingredients (data) because of privacy rules. They can only send their ideas back to the teacher to update the main recipe.

The problem is that these 50 students are all different. Some are Italian, some are vegan, some are professional chefs, and some are beginners. A single recipe is too simple, and giving everyone a totally unique recipe is too chaotic.

The Problem: The "One-Size-Fits-None" Dilemma

Existing methods try to solve this by either:

Grouping them: Putting all Italians in one group and all Vegans in another. But what if an Italian student is also a beginner? The groups are too rigid.
Personalizing everyone: Giving everyone a unique recipe. But this ignores the fact that they might share some common knowledge (like how to chop an onion).

The paper argues that real life is more complex. Similarity exists at multiple levels:

Level 1 (Global): Everyone knows how to boil water.
Level 2 (Subgroup): Everyone in the "Italian" group knows how to make pasta sauce.
Level 3 (Individual): You specifically know how to add a secret family spice.

Most current AI systems can only handle two of these levels at once. They get stuck when the structure is messy and multi-layered.

The Solution: FeMAM (The "Lego Tower" Approach)

The authors propose a new method called FeMAM (Federated Multi-level Additive Modeling).

Think of FeMAM as building a Lego tower where each layer adds a little more detail:

The Base Layer: Everyone starts with the same basic block (the Global Model). This covers the "boiling water" part.
The Middle Layers: As the students train, the system automatically figures out who needs extra help. It adds a "Pasta" block for the Italian group and a "Vegan" block for the others. These are shared among subgroups.
The Top Layer: Finally, it adds a tiny, custom "Secret Spice" block just for you.

The Magic Trick:
Instead of building a new, giant tower from scratch for every student, FeMAM stacks these blocks.

You take the Global Block.
You add the Group Block on top.
You add your Personal Block on top.
Result: Your final tower is the sum of all these parts.

How It Stays Efficient (The "Grow and Prune" Strategy)

You might think, "Wait, if everyone has 5 different blocks, isn't that too much data to send back and forth?"

The paper says no, thanks to two smart rules:

One Step at a Time (Boosting): The system doesn't build the whole tower at once. It builds the base, waits until it's stable, then adds the next layer. It only sends the new layer back to the teacher. This keeps the communication cost low, just like sending a single postcard instead of a whole book.
The "Stop" Button (Pruning): Not everyone needs 5 layers.
- If a student is a simple beginner, the system realizes, "Hey, the base layer is enough for you!" and stops adding more blocks.
- If a student is a complex expert, the system keeps adding layers until it's perfect.
- This means the system adapts to the student. It doesn't force a complex tower on a simple user, saving them time and storage.

What They Found

The researchers tested this on image datasets (like recognizing cats and dogs) and large language models (like AI chatbots).

Versatility: Whether the students were grouped by simple categories or had messy, complex differences, FeMAM worked better than methods that tried to force everyone into fixed groups or give everyone a totally unique model.
Smart Adaptation: In the experiments, the system automatically decided that some users needed 1 layer, some needed 2, and some needed 5. It didn't guess; it learned what worked best for each person.
Cost: Even though it uses multiple "layers" of models, it didn't cost much more to run than standard methods because it only communicated one layer at a time.

The Bottom Line

FeMAM is like a smart, flexible teacher who realizes that learning happens in layers. It starts with what everyone shares, adds what specific groups share, and finally adds what makes you unique. It builds this knowledge step-by-step, only adding what is necessary, making it perfect for a world where everyone is different but still connected.

Technical Summary: Personalized Additive Modeling for Multi-level Federated Learning

Problem Statement
Contemporary Federated Learning (FL) faces a critical challenge: balancing the need for generalization with user-specific personalization in the presence of highly heterogeneous client data. Standard FL assumptions of Independent and Identically Distributed (IID) data often fail to capture the complex, structured non-IID patterns found in real-world scenarios. Existing methods, including Clustered FL and Personalized FL (PFL), are typically designed for restricted heterogeneity settings, such as a fixed number of clusters or a rigid global-plus-local structure. These approaches struggle when client relationships exhibit multi-granular, implicit structures (e.g., coarse OS-level similarities, subgroup application patterns, and fine-grained activity traits) because they require prior knowledge of task boundaries or commit to a fixed hierarchy that may not match the underlying data distribution.

Methodology: Federated Multi-level Additive Modeling (FeMAM)
The authors propose FeMAM, a flexible framework that models FL from a multi-level non-IID perspective. Instead of a fixed hierarchy, FeMAM learns a hierarchy of shareable models across multiple granularities (from global to client-specific) and constructs personalized predictors via additive composition.

Multi-level Additive Architecture: The system organizes knowledge into $L$ levels. At each level $l$ , the server maintains $K_l$ models representing knowledge at that granularity. A client's personalized predictor is formed by additively composing one selected model from each level:
$f_i(X_i) = \sum_{l=1}^{L} \sum_{k=1}^{K_l} r_{i,k}^{(l)} f(X_i; \Theta_{k}^{(l)})$
where $r_{i,k}^{(l)}$ is an assignment indicator. This structure allows for coarse-to-fine refinement, where higher levels provide shared structure and lower levels act as residual corrections.
Boosting-like Optimization: To avoid the cost of jointly optimizing all levels and the need to pre-specify the hierarchy depth, FeMAM employs a sequential, boosting-like training strategy. It learns levels one by one, freezing previously learned levels and optimizing only the current level to minimize the residual error. This ensures training stability and allows the method to adapt to the data without committing to a fixed structure upfront.
Client-side Structure Pruning: To ensure cost-efficiency and prevent overfitting, FeMAM implements a pruning mechanism. After training a new level, a client accepts it only if it reduces the validation error. This results in an "on-demand" personalization depth where clients with simpler data distributions may stop at earlier levels, while those with complex heterogeneity utilize deeper hierarchies.
Cost-Adaptive Design: Despite maintaining multiple models, FeMAM remains communication-efficient. During training, only the current level's model is communicated per client per round, keeping communication costs comparable to single-model FL. At inference, clients store only the models they have accepted, which is typically fewer than the maximum hierarchy depth due to pruning.

Key Contributions

Formulation: The paper formulates federated personalization under multi-level non-IID heterogeneity, addressing the structural complexity induced by modern AI personalization needs.
Mechanism: It introduces a multi-level additive personalization mechanism that constructs client predictors via residual-style composition across levels, eliminating the need for explicit task boundaries.
Algorithm: The authors propose a "grow-and-prune" training pipeline that progressively expands the hierarchy while pruning unnecessary levels per client, achieving adaptivity without a fixed structure.
Theory: A convergence analysis is provided, showing that the sequence of objectives is monotone non-increasing, similar to gradient boosting.
Empirical Validation: Extensive experiments demonstrate that FeMAM consistently outperforms representative clustered and personalized FL baselines across diverse non-IID settings.

Experimental Results
The authors evaluated FeMAM on image classification datasets (Tiny ImageNet, CIFAR-100) and a large-scale foundation model scenario (Qwen3-4B with LoRA on FLAN-T5 tasks).

Performance: FeMAM consistently achieved the highest test accuracy and Macro-F1 scores across three non-IID partition types: Cluster-wise, Client-wise, and Multi-level. It outperformed strong baselines like FedAvg, Ditto, FeSEM, FedCAM, and pFedMoE.
Adaptability: In multi-level and client-wise non-IID settings, FeMAM automatically utilized deeper hierarchies (up to 5 levels), whereas in Cluster-wise settings, it pruned back to fewer levels, demonstrating its ability to match the underlying data structure.
Foundation Models: In the large-scale experiment, FeMAM achieved the lowest average loss and highest ROUGE-1 score, outperforming methods that rely on simple global-local designs.
Efficiency: The method showed competitive wall-clock runtime (faster than Ditto) and communication efficiency. Pruning ensured that the average number of stored models per client remained low (e.g., ~2.4 in Cluster-wise settings), reducing storage costs compared to methods like FedEM.

Significance and Claims
The paper claims that FeMAM offers a robust solution for the "tension between generalization and personalization" in modern FL. Its primary significance lies in its ability to approximate diverse, complex non-IID structures without requiring prior knowledge of the number of clusters or the depth of the hierarchy. By treating personalization depth as client-dependent and using an additive, boosting-style approach, FeMAM adapts to the data rather than forcing the data into a pre-defined model structure. The authors position this as a step toward handling the structured heterogeneity inherent in large-scale, multi-capability AI systems, such as foundation models with client-side adaptation.

Limitations and Future Work
The authors modestly note that their current evaluation of level-wise components is limited to a representative instantiation (FedAvg, K-means clustering, and local fine-tuning). They acknowledge that a systematic study of alternative level-wise methods (e.g., IFCA, FedPer) and their cross-level interactions remains an open direction. Additionally, while foundation model experiments were conducted, the scale was smaller than realistic deployments, suggesting future work should evaluate FeMAM on substantially larger and more diverse task sets.