HyFLM: A Hypernetwork-Based Federated Learning with Multidimensional Trajectory Optimization on Diffusion Paths

Abstract

The effective training of large-scale distributed deep learning models has become an active and emerging research area in recent years. Federated learning (FL) can address those challenges by training global models through parameter exchange of client models rather than raw data sharing, thereby preserving security and communication efficiency. However, conventional linear aggregation approaches in FL neglect heterogeneous client models and non-IID data. This often results in inter-layer information imbalance and feature-space misalignment, leading to low overall accuracy and unstable training. To overcome these limitations, we propose HyFLM, a personalized federated learning framework that maximizes performance with Multidimensional Trajectory Optimization theory (MTO) on diffusion paths. HyFLM extends a diffusion-based FL framework by encoding client–parameter dependencies with a diffusion model and precisely controlling dimension-specific paths, thereby generating personalized weights that reflect both the data complexity and the resource constraints of each client. In addition, a lightweight hypernetwork generates client-specific adapters or weights to further enhance personalization. Extensive experiments on multiple benchmarks demonstrate that HyFLM consistently outperforms major baselines in terms of both accuracy and communication efficiency, achieving faster convergence and higher accuracy. Furthermore, ablation studies verify the contribution of MAC to convergence acceleration, confirming that HyFLM is an effective and practical personalized FL paradigm for heterogeneous client models.

1. Introduction

Federated learning (FL) is a distributed machine learning paradigm that enables multiple clients to collaboratively train a shared model without exchanging raw data [1]. In typical FL frameworks, such as FedAvg [2,3,4,5,6,7], each client updates gradients or model parameters from its local dataset to a central server, without sharing their own data. This approach is being studied in various fields, including finance, healthcare, and business. For example, FL can implement intelligent functions on smartphones and IoT devices by learning user data on the device without uploading personal data to the cloud. In healthcare, FL enables hospitals to jointly train diagnostic models based on distributed patient datasets, enabling collaborative medical AI research without compromising patient privacy. Furthermore, in large-scale video surveillance networks such as military perimeter monitoring systems, federated learning (FL) serves as an effective approach to share locally trained model parameters, thus reducing overall training costs and improving the performance of distributed AI models [8,9]. Despite its potential, existing FL faces several fundamental challenges that can degrade performance. One of the most significant issues in FL is the non-Independent and Identically Distributed (non-IID) environment [10], as illustrated in Figure 1. In practical federated learning scenarios, models are often deployed under non-IID conditions, where client data distributions differ significantly [1]. This exposes the problem with simple aggregation methods and leads to inconsistent or biased global models [11]. For example, the well-known Federated Averaging algorithm, FedAvg [2], linearly averages client models as if all client updates were drawn from the same distribution, which can lead to low-probability regions of the parameter space and information collapse [12]. Moreover, averaging two well-optimized client models trained on different data can produce a global model that performs poorly on both client tasks. This feature space collapse or misalignment arises because simple linear combinations do not reconcile the knowledge of different clients. Second, the heterogeneity of the model presents additional limitations. Standard FL frameworks assume the same model architecture for all clients, making it difficult to accommodate small models or unique user-defined models. This uniform approach can lead to misalignment of learned representations and does not account for the various computational resources of the clients [1]. These challenges highlight the need for more sophisticated FL techniques to handle a heterogeneous and non-IID environment, prevent representation collapse, and adapt to heterogeneous client environments.

To overcome these limitations, recent FL research has focused on implementing more advanced aggregation mechanisms for each client and addressing data heterogeneity. One promising solution is personalized Federated Learning (pFL) [12], which addresses the limitations of a single global model by tailoring the learned model to the unique data distribution of each client. Various pFL methods have been proposed, notable examples of which include FedPer [13], which shares a common feature extractor but trains a private output layer for each client; FedRep [14], which trains a federated representation and allows each client to train its own classifier; and FedBABU [15], which fine-tunes a globally pre-trained model for each client without a final aggregation step. These methods often achieve improvements by identifying client-specific requirements and mitigating the negative impact of data heterogeneity. Furthermore, to overcome the limits of simple averaging, recent work has leveraged diffusion models for data-driven aggregation in FL. pFedGPA [12] employs a server-side diffusion model that learns the distribution of client parameters and, via parameter inversion, generates personalized weights for each client, thus encoding client-specific data characteristics. Similarly, FedDiff [16] integrates a diffusion model into the FL loop for multi-modal data fusion, achieving higher accuracy and communication efficiency. By framing model fusion as a generative task, these approaches better capture the complex, high-dimensional structure of parameter space than linear methods.

Building upon diffusion-based FL frameworks, we propose HyFLM, an enhanced FL framework that optimizes model training through inference trajectory optimization. HyFLM shows reliable convergence on heterogeneous data sets. It shows a rapidly decreasing training loss and reaches an accuracy approaching 95% in only a few rounds, indicating both fast convergence and high final accuracy compared to FL baselines. In summary, the contributions of this work are as follows:

• We propose a novel federated learning framework that combines diffusion-based aggregation and hypernetwork-driven personalization under the theory of inference trajectory optimization.
• Our empirical results show that the proposed approach achieves superior generalization performance on diverse benchmark datasets compared to well-recognized baselines.
• We demonstrate that HyFLM yields stable convergence under non-IID data and heterogeneous client models, reduced communication cost, and preserved security.

Section 2 reviews the background and related work on federated learning (FL), hypernetwork-based personalization, and diffusion-based FL, and explains why our work focuses on the intersection of these areas. Section 3 introduces our diffusion-space inference and the Multidimensional Adaptive Coefficients (MAC). Section 4 presents the hypernetwork-driven personalization and its integration with MAC. Section 5 reports experiments result, and Section 6 discusses limitations and concludes.

2. Related Work

This section reviews research most relevant to our contribution, beginning with the foundations of federated learning and personalized FL under non-IID client distributions, then progressing to hypernetwork-based approaches that condition on client descriptors to produce lightweight adapters for client-specific models, and culminating with diffusion-based FL that leverages denoising/score models for robust aggregation and trajectory control.

2.1. Federated Learning

Federated learning has emerged as a prominent paradigm for training models on decentralized AI environment. The landmark FedAvg [2] algorithm showed that simply averaging the parameters of the client model can produce a viable global model, even in unbalanced non-IID data distributions. However, subsequent studies have shown that data heterogeneity remains a fundamental bottleneck. The so-called client drift arises because each client often converges to a local optimum that does not coincide with the global optimum, causing local updates to deviate and degrade the accuracy of the global model [1]. To combat this, numerous FL frameworks have been proposed to stabilize training on non-IID data. For example, FedProx [17] adds a proximal regularization term to keep local model updates closer to global weights, improving convergence under statistical heterogeneity. Similarly, SCAFFOLD [18] uses control variates (server and client update correction vectors) to explicitly mitigate client drift during local training. Another research strand is personalized FL [13] which seeks to optimize each client’s own performance rather than relying on a single global model. Another approach is cluster-based FL [19,20], which partitions clients into groups with similar data distributions and performs model averaging within each cluster [10]. This reduces cross-client divergence, but producing high-quality personalized models and efficiently sharing knowledge between clusters remains an open challenge. Recent works [1,21] explore fine-grained personalization, but achieving both high accuracy and fairness in large-scale, highly non-IID settings continues to be actively researched.

2.2. Hypernetwork-Based FL

A hypernetwork is a neural network that generates (all or part of) the weights of another target network. This concept has gained traction as a tool for tackling model heterogeneity and personalization in FL. Early federated studies used hypernetwork mainly to provide better initial weights for each client model [22], but recent work conditions weight generation on richer signals such as client metadata, update gradients, or graph-based relationships between clients. For example, hFedF [23] (federated fusion based on the hypernetwork) embeds each client’s model update into a latent state and uses a server-side hypernetwork to fuse models in a non-linear fashion, producing a robust global model that outperforms linear FedAvg under the non-IID condition. HyperFedNet [24] addresses the communication overhead and privacy risks of conventional FL by transmitting a group of hypernetwork parameters. Meanwhile, FedAGHN [25] employs an attentive graph hypernetwork to personalize models. It treats clients as nodes in a graph and uses attention-weighted edges to learn dynamic collaboration weights. As a result, the hypernetwork can produce a customized initial model for each client by aggregating parameters from neighboring clients. This approach maximizes cooperation among statistically similar clients while preserving individual performance. Likewise, FedSheafHN [26] leverages advanced graph signal processing. It constructs a server-side sheaf diffusion process over a client collaboration graph to derive high-order client embeddings (capturing complex inter-client relationships), then feeds those embeddings into a hypernetwork that generates personalized Graph Neural Network parameters for each client. Hypernetwork generation has shown strong gains on graph-structured data and unseen clients. Recent work enriches conditioning inputs (e.g., graphs, contrastive embeddings) and employs auxiliary objectives such as adversarial or contrastive losses to improve robustness under heterogeneity. In general, advances in hypernetwork-based FL yield higher accuracy and faster convergence, marking it as a promising direction to improve personalization and communication efficiency.

2.3. Diffusion Based Federated Learning

Diffusion models have attracted substantial interest in federated learning. With a solid theoretical foundation and a simple yet effective loss formulation [27], these models have the ability to generate high-quality data for a wide range of applications. Recently, the synergy between diffusion generative models and FL has been actively explored to mitigate non-IID challenges and improve generalization [28,29]. Denoising diffusion models (DDPMs [27]) can synthesize realistic data by iteratively refining noise, and several studies have leveraged this capability within FL. One line of work uses diffusion for synthetic data augmentation [30], where local diffusion models generate auxiliary samples to balance skewed distributions and improve global convergence. Frameworks such as FedDDPM [31] and FedDrip [32] have demonstrated consistent gains across benchmarks, while additional studies showed that diffusion-based augmentation boosts performance on highly imbalanced datasets, alleviating client drift and mode collapse. Another direction applies diffusion to model fusion, replacing simple averaging with generative aggregation. For example, pFedGPA [12] employs parameter inversion and denoising sampling to generate refined personalized weights, achieving better performance over FedAvg [2]. FedDiff [16] extends this idea to multi-modal settings, coordinating diffusion branches across modalities (e.g., hyperspectral vs. LiDAR) and enabling privacy-preserving feature fusion. Finally, diffusion has been explored as an enhanced training signal, where large diffusion models act as auxiliary teachers to provide generative semantic guidance and consistency regularization, improving alignment of heterogeneous client representations. Such diffusion-based approaches are emerging as versatile tools in FL, enabling balanced data generation, advanced parameter aggregation, and richer training signals. Furthermore, recent advances in diffusion-based FL highlight that optimizing inference trajectories is becoming a critical challenge. The present work introduces a novel methodology to address this issue.

Classical FL with FedAvg-style aggregation degrades under non-IID partitions. In personalization, hypernetwork offers a streamlined way to produce client-conditioned weights from compact server-side models. Diffusion-based FL improves robustness and exploration, but typically lacks explicit trajectory control and a coherent link to per-client adaptation. Building on these observations, HyFLM integrates hypernetwork-based FL with multidimensional trajectory optimization (MTO) on diffusion paths to provide time-varying trajectory control while the hypernetwork injects client specific signals yielding a robust, communication-aware personalization framework.

3. Preliminaries

This section briefly reviews the two theoretical pillars used in HyFLM: Denoising Diffusion Probabilistic Models (DDPM) [27] and the more recent Multidimensional Adaptive Coefficient (MAC) [33]. Understanding these concepts clarifies how HyFLM leverages diffusion-based generative modeling and adaptive control of inference trajectories.

3.1. Design Diffusion-Based Inference Space

In a federated learning environment with non-IID data and heterogeneous architectures, naïve parameter averaging neglects geometric coupling across the multiple layers of the weight space and can degrade accuracy. To address this, we construct a diffusion-based inference space that models the empirical distribution of client weights via a forward noising and reverse denoising process. This enables the server to synthesize a global anchor while preserving the curvature and scale of the layer. Coupling this space with Multidimensional Adaptive Coefficients (MAC) provides trajectory-level control. DDPMs can be used in FL training. They define a forward diffusion process that adds noise to the data step by setp and a reverse denoising process that learns to invert this perturbation [33]. The generative model is parameterized as a Markov chain with latent variables $\{x_1,\dots ,x_T\}$ and observed data $x_0$. We denote the forward process by $q(x_t\mid x_{t-1})$ and the learned reverse chain by $p_{\theta }(x_{t-1}\mid x_t)$, with a fixed Gaussian prior $p\left(x_T\right)=\mathcal{N}(0,I)$. The reverse process is formulated as

$$p_{\theta }\left(x_{0:T}\right)=p_{\theta }\left(x_0\right)\prod _{t=1}^T{p}_{\theta }(x_{t-1}\mid x_t),$$

(1)

where $p_{\theta }(x_{t-1}\mid x_t)$ denotes Gaussian transitions with learned mean and covariance ${\mu }_{\theta }(x_t,t)$ and ${\Sigma }_{\theta }(x_t,t)$ [27]. The approximate posterior (also called the forward or diffusion process) is fixed to a Markov chain that gradually corrupts data by adding Gaussian noise at each step [27]:

$$q(x_t\mid x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I),$$

(2)

where ${\beta }_t\in (0,1)$ is a small variance schedule. Following the formulation in [27], repeated application of this transition gradually drives the data distribution toward a standard normal, enabling sampling to begin from pure noise, as illustrated in Figure 2.

To synthesize data, DDPMs train a neural network to model the reverse transition $p_{\theta }(x_{t-1}\mid x_t)$, thus learning to iteratively denoise perturbed samples. The training objective minimizes a variational lower bound on the negative log-likelihood, which simplifies to a standard denoising loss function.

A specific parameterization reveals equivalence with denoising score matching and annealed Langevin dynamics, and the sampling procedure resembles a progressive decoding mechanism. These properties make diffusion models powerful building blocks for FL frameworks that can handle diverse client data and produce high-quality generative output. These methods are incorporated into HyFLM to construct client-specific diffusion trajectories while preserving global generative capability.

3.2. Multidimensional Adaptive Coefficient (MAC) for Inference Trajectory Optimization

In conventional diffusion or flow models, a single global scaling factor is applied to the entire trajectory, limiting the flexibility of the model. To overcome this restriction, we introduce the Multidimensional Adaptive Coefficient (MAC) [33], which enables different components of the model update to scale adaptively over time, thus improving trajectory expressiveness and optimization efficiency, as illustrated in Figure 3.

Flow and diffusion-based models typically employ one-dimensional scalar coefficients to control inference trajectories, limiting adaptability across different datasets or tasks. This scalar formulation imposes two main constraints. The Multidimensional Adaptive Coefficient (MAC) addresses this limitation by extending the conventional unidimensional coefficient to higher dimensions. Given the trajectory $x\left(t\right)={\gamma }_0\left(t\right)x_0+{\gamma }_1\left(t\right)x_1,x_0,x_1\in {\mathbb{R}}^d$ and $\gamma \left(t\right)\in C^1([0,T],{\mathbb{R}}^{d\times d\times 2}$), the multi-dimensional coefficient is defined as:

$$\gamma \left(t\right)=[{\gamma }_0\left(t\right),{\gamma }_1\left(t\right)]:[0,T]\to \in {\mathbb{R}}^{d\times d\times 2}$$

(3)

The Multidimensional Adaptive Coefficient (MAC) is defined as ${\gamma }_{\varphi }$, which dynamically adjusts to each inference trajectory $x_{\theta ,\varphi }\left(t\right)$ defined over discrete inference times $\tau =\{t_0,\dots ,t_N\}$. Formally, the MAC is defined as

$${\gamma }_{\varphi }(t,{\mathbf{x}}_{\theta ,\varphi }^S)=[{\gamma }_{0,\varphi }(t,{\mathbf{x}}_{\theta ,\varphi }^S),{\gamma }_{1,\varphi }(t,{\mathbf{x}}_{\theta ,\varphi }^S)]:[0,T]\times {\mathbb{R}}^{\mathcal{l}\times d}\to {\mathbb{R}}^{d\times d\times 2},$$

(4)

where $S=\{t^{\left(1\right)},\dots ,t^{\left(\mathcal{l}\right)}\}\subseteq [0,T]$ represents an arbitrary subset of inference times and $\varphi$ parameterizes the adaptive behavior. Boundary and smoothness properties are preserved, allowing the MAC to flexibly model sample- and time-dependent trajectories during diffusion-based inference.

To handle high-dimensional datasets, formulate the trajectory optimization as a GAN-based min–max problem:

$$\underset{\theta ,\varphi }{min}\underset{\psi }{max}{\mathbb{E}}_{x_0}\left[log{D}_{\psi }\left(x_0\right)\right]+{\mathbb{E}}_{x_T}\left[log\left(1-D_{\psi }\left(G_{\theta ,\varphi }(\tau ,x_T)\right)\right)\right],$$

(5)

Here, $D_{\psi }$ acts as a discriminator that distinguishes real samples $x_0$ from generated samples $G_{\theta ,\varphi }(\tau ,x_T)$, while $G_{\theta ,\varphi }$ jointly optimizes the vector field $\theta$ and the MAC parameters $\varphi$. This adversarial formulation enables implicit divergence minimization and improves transport quality in diffusion-based high-dimensional inference.

In summary, DDPM provides a principled probabilistic foundation for modeling generative processes via forward diffusion and learned reverse denoising. MAC extends this framework by introducing learnable multidimensional coefficients that control the trajectory of the client weighty adaptively. These preliminaries form the conceptual basis of HyFLM, where a personalized hypernetwork leverages MAC-modulated diffusion paths for flexible and high-quality federated learning.

4. Methodology

HyFLM is proposed as a robust and efficient federated learning framework. It effectively learns client-specific weights through adaptive coefficient optimization in a diffusion-based environment while simultaneously generating personalized client models via a server-side hypernetwork. As illustrated in Figure 4 HyFLM includes two main components: (1) client-side training with an advanced MAC (which adaptively scales update directions at each step and gates residual connections and message signals); and (2) a lightweight server-side hypernetwork for personalization. The advanced MAC regulates how the global model evolves. This enables HyFLM to capture client-specific parameter dependencies and deliver optimal and more personalized aggregation trajectories than previous approaches. Another important element of HyFLM is a lightweight hypernetwork that enhances personalization for resource-constrained or biased clients. The hypernetwork conditions the diffusion-based aggregator on client-specific signals, adjusting the generated parameters to reflect data biases or resource constraints. However, client-specific weight generation is achieved through a hypernetwork on the server side. By integrating and extending these components, HyFLM can minimize the distortion of the non-IID environment, address model heterogeneity, and provide efficient communication personalization across diverse clients. This section details each module and outlines the overall process of integrating them; the detailed procedure of HyFLM is summarized in Algorithm 1.

Algorithm 1: HyFLM: A Hypernetwork-Based Federated Learning with Multidimensional Trajectory Optimization

4.1. Research Methodology Overview

The research focuses on personalized federated learning in realistic heterogeneous environments, where many clients train on non-IID data under constrained communication and computation costs. In addition, the research aims to improve the average accuracy per-client in statistically biased partitions, lower the total communication required to reach a given target accuracy, and strengthen convergence stability and robustness among diverse client populations. A final objective is to design a federated learning framework that scales gracefully with the number of clients and the available system resources, so that high performance remains attainable even in large-scale deployments. These aims motivate several concrete research problems: whether the proposed method can surpass strong existing baselines in both early and final training stages under the simulation environment; how much of the observed improvement is attributable to diffusion-path trajectory control via the MAC module and how much to hypernetwork-based personalization, individually and in combination; how the communication–computation trade-offs evolve as key design parameters such as the number of function evaluations, adapter capacity, and update frequency vary; and how stable the training dynamics remain under severe non-IID conditions and large client pools.

To investigate these problems, the work adopts a controlled empirical methodology with fixed data splits and equalized training budgets across all compared methods. The experimental pipeline is implemented in PyTorch with a simulated server–client FL environment that orchestrates local training, aggregation, and personalization. The proposed technique integrates MAC-gated diffusion sampling on the server with lightweight hypernetwork adapters that generate client-specific parameter adjustments from learned embeddings. Multiple random seeds are used and results are reported with mean and standard deviation to quantify variance. Communication cost is modeled as the product of the number of communication rounds and the payload per round, while computational effort is approximated through FLOPs-based proxies. Independent variables include the number of clients and their participation ratio, the Dirichlet concentration parameter controlling the degree of non-IID partitioning, the number of function evaluations for diffusion, the rank or bitwidth of the personalization adapters, and the regeneration interval for global updates. Dependent variables comprise personalized test accuracy, the number of rounds required to reach a performance plateau or a fixed accuracy fraction, payload per round, total data volume exchanged, and stability indicators; together these metrics support a systematic analysis of accuracy–efficiency trade-offs under identical resource constraints.

4.2. Diffusion-Based Training and Trajectory Optimization with MAC

In federated learning (FL) systems characterized by irregular data distributions and heterogeneous client features, diffusion-based FL approaches leverage the generative capabilities of diffusion models to capture the data distribution of each client and produce personalized parameters accordingly. The Multidimensional Adaptive Coefficient (MAC), denoted ${\gamma }_{\varphi }$, is a novel coefficient introduced to optimize inference trajectories in flow and diffusion models. In this study, MAC is applied within FL to optimize the inference trajectories of temporally aligned flow and diffusion models between the global and client models, thereby generating client-specific personalized parameters tailored to each client’s characteristics. The entire framework operates through a compact iterative loop of pretraining and local update → aggregation and correction. The core idea of this approach is to integrate powerful diffusion-based generative modeling into federated learning, enabling the learning of the multi-dimensional distribution of client model parameters within the parameter space and improving both adaptation and personalization across heterogeneous environments.

Before global model training, each client performs a brief pre-training phase using randomly sampled MAC coefficients to adapt $H_{\theta }$ for a stable response to non-linear (curved) trajectories. In each communication round, the server broadcasts the current model $(\theta ,\phi )$ to a subset of clients. Each client k updates both ${\theta }_k$ and ${\phi }_k$ by minimizing the local objective:

$${\theta }_k\leftarrow \theta -\eta {\nabla }_{\theta }{L}_k(\theta ,\phi ),{\phi }_k\leftarrow \phi -\eta {\nabla }_{\phi }{L}_k(\theta ,\phi ),$$

(6)

where $\eta$ is the learning rate. This procedure jointly optimizes the vector field $\left(\theta \right)$ and the trajectory plan$\left(\phi \right)$, adapting both to the data distribution of each client and the resource constraints.

In normalized time $t\in [0,1]$, HyFLM adopts an advanced Multidimensional Adaptive Coefficient based on the theory of the Multidimensional Adaptive Coefficient (MAC), which dynamically balances exploration and stability during diffusion-based optimization. Specifically, these gates regulate the message and residual paths:

• ${\gamma }_{\mathrm{msg}}\left(t\right)$: Scales the magnitude of the diffusion update (drift or score term);
• ${\gamma }_{\mathrm{res}}\left(t\right)$: Controls the residual pull toward an anchor (e.g., the previous global model).

The server-side model updates govern the global diffusion process, where the server integrates information from client models to reconstruct and refine the global parameter distribution. Each reverse step applies MAC-modulated scaling to both the diffusion score and the residual path, thereby adapting the generative trajectory to the global dynamics of heterogeneous client updates. The update rule is formally expressed as

$$x_{k+1}=x_k+h_k{\gamma }_{\mathrm{msg}}\left(t_k\right)s_{\varphi }(x_k,t_k)+h_k{\gamma }_{\mathrm{res}}\left(t_k\right)(a-x_k)+{\sigma }_k{z}_k,$$

(7)

where $t_k=k/(N-1)$, $h_k$ denotes the size of the integration step, a is the anchor representing the global reference model, and $z_k\sim \mathcal{N}(0,I)$ introduces stochastic regularization through the variance schedule ${\sigma }_k$.

The client-side local update applies MAC gating to control both the gradient descent magnitude and the retention of prior knowledge during local optimization. This mechanism enables each client to follow an individualized optimization trajectory that reflects its data distribution and computational characteristics, preventing unstable divergence in heterogeneous settings. The update process is defined as:

$$\begin{array}{cc}\hfill g_{\tau }^{\prime }& ={\gamma }_{\mathrm{msg}}\left(t\right)\nabla \mathcal{L}\left(w^{\left(\tau \right)}\right),t=\frac{\tau }{max(1,T_{\mathrm{loc}}-1)},\hfill \\ \hfill {\tilde{w}}^{(\tau +1)}& =w^{\left(\tau \right)}-\eta g_{\tau }^{\prime },\hfill \\ \hfill w^{(\tau +1)}& ={\tilde{w}}^{(\tau +1)}+\alpha {\gamma }_{\mathrm{res}}\left(t\right)\left(w^{\left(0\right)}-{\tilde{w}}^{(\tau +1)}\right),\hfill \end{array}$$

(8)

where $\eta$ denotes the learning rate and $\alpha$ represents the residual coefficient. In this formulation, ${\gamma }_{\mathrm{msg}}\left(t\right)$ adaptively scales the gradient to regulate the intensity of learning at each step, effectively serving as a temporal controller of the learning rate, while ${\gamma }_{\mathrm{res}}\left(t\right)$ enforces a smooth proximal correction towards the initial model $w^{\left(0\right)}$. These gates ensure that client updates remain both locally responsive and globally consistent, thus improving convergence stability under non-IID and resource-constrained environments.

To achieve per-dimension and time-adaptive control, the gates are implemented as vector-valued functions (or per-group vectors) with smoothing to prevent overfitting. Each gate is parameterized by a truncated low-pass Fourier expansion.

$$\gamma \left(t\right)={\Pi }_{[{\gamma }_{min},{\gamma }_{max}]}\left(b+\sum _{k=1}^K\left(a_{k}cos2\pi kt+b_k^{\prime }sin2\pi kt\right)e^{-{\sigma }^2{k}^2}\right),$$

(9)

where $\{a_k,b_k^{\prime },b\}$ are learnable coefficients (optionally vector-valued), $\sigma >0$ denotes the bandwidth, and $\Pi$ is the projection operator that clamps outputs to $[{\gamma }_{min},{\gamma }_{max}]$. We parameterize $\gamma \left(t\right)$ with a truncated low-pass Fourier expansion (Equation (9)), projected to $[{\gamma }_{min},{\gamma }_{max}]$. This guarantees smooth, bounded gates, enables multi-timescale control of the diffusion/local update, and has $O\left(Kd\right)$ per-step evaluation cost suitable for FL. With the spectral envelope $e^{-{\sigma }^2{k}^2}$ and projection ${\Pi }_{[{\gamma }_{min},{\gamma }_{max}]}$, $\gamma \left(t\right)$ is bounded and differentiable, and $\parallel {\gamma }^{\prime }\left(t\right)\parallel \le {\sum }_{k=1}^{K}2\pi k\parallel c_k\parallel e^{-{\sigma }^2{k}^2}$. Hence MAC-modulated increments in Equations (7) and (8) remain well-behaved. When ${\gamma }_{\mathrm{msg}}\equiv 1$ and ${\gamma }_{\mathrm{res}}\equiv 0$, the method reduces to standard sampling/gradient descent. When ${\gamma }_{\mathrm{msg}}\equiv 1$ and ${\gamma }_{\mathrm{res}}\equiv 0$, this formulation reduces to standard sampling or gradient descent.

4.3. Hypernetwork-Based Personalization

A hypernetwork is a meta-generator network that produces the parameters of another target network, enabling flexible adaptation between tasks or clients. In the proposed framework, we employ a hypernetwork ${\mathcal{H}}_{\psi }$ that maps client embeddings to personalized weight adjustments. Given the embedding matrix $\widehat{X}$ (or individual embedding $e_i$), the hypernetwork generates client-specific parameters:

$$\tilde{\Omega }=\{{\tilde{\omega }}_1,{\tilde{\omega }}_2,\dots ,{\tilde{\omega }}_K\},\mathrm{with}{\tilde{\omega }}_i={\mathcal{H}}_{\psi }\left(e_i\right).$$

(10)

Each ${\tilde{\omega }}_i$ represents an adaptation of the shared model $\theta$ for client i, associated with an embedding vector $e_i$ in our implementation. This embedding encodes client-specific information—such as data statistics or learned latent features—and is utilized both to modulate the diffusion process through the MAC and to serve as input to the hypernetwork. This design ensures that personalization is consistently guided by client characteristics, yielding coherent and data-informed parameter generation across clients.

$${\tilde{w}}_i=\theta \oplus \Delta w_i,$$

(11)

where ⊕ denotes the composition or addition of a delta. This enables efficient personalization without retraining the denoising network ${ϵ}_{\varphi }$ on each client. The same client embedding used in MAC gating is reused here, ensuring consistent characterization of client heterogeneity across both local trajectory modulation and global parameter generation. In this way, the hypernetwork serves as a conditional generator that maps heterogeneous client embeddings to individualized model configurations, thereby enabling scalable and communication-efficient personalization within the federated learning framework. HyFLM integrates MAC gating into both client-side optimization and server-side diffusion to modulate update magnitudes and trajectories across dimensions. As a result, it better captures individual client characteristics and improves stability, personalization, and robustness under heterogeneous data. The vectorized, Fourier-parameterized gates enable per-parameter and time-adaptive control over learning and diffusion. The hypernetwork personalizes the final model for each client, yielding scalable personalization. Collectively, these innovations ensure stability, improved convergence, and robust performance in non-IID federated environments.

5. Experiments

5.1. Experimental Environment

All experiments were conducted on a Linux workstation with 4× NVIDIA RTX A6000 GPUs (48 GB each; NVIDIA Corporation, Santa Clara, CA, USA), CUDA 12.4, and NVIDIA driver 550.54.14. The system ran Ubuntu 22.04 with Python 3.10 and PyTorch 2.1. Our federated learning simulation was built using PyTorch and Flower/PySyft primitives, and diffusion components were adapted from Hugging Face’s diffusers.Resource usage was recorded with Weights & Biases. All experiments were containerized and run with fixed seeds for reproducibility.

5.2. Dataset

HyFLM was evaluated on three standard vision benchmarks EMNIST, Fashion-MNIST, and CIFAR-10 under varying numbers of clients and baseline settings. Each data set was first treated as a single global pool following the standard torchvision train/test split. Subsequently, the partitions of non-IID clients were constructed to simulate the label-skew and quantity-skew scenarios for $M\in \{10,20,100\}$ clients, as described in

Table 1. Summary of dataset properties for MNIST, EMNIST, and CIFAR-10.

Dataset	Image Size	Channels	# Classes	Clients
MNIST	$28\times 28$	1 (grayscale)	10	10/20/100
EMNIST	$28\times 28$	1 (grayscale)	26	10/20/100
CIFAR-10	$32\times 32$	3 (RGB)	10	10/20/100

.

• EMNIST. Handwritten English letters (26 classes), grayscale $28\times 28$, which reflects on-device handwriting variability and serves as a classic benchmark.
• Fashion-MNIST. Grayscale clothing images (10 classes) at $28\times 28$. It provides more shape/texture diversity than EMNIST and is suitable for lightweight CNNs.
• CIFAR-10. Natural RGB images (10 classes) at $32\times 32$. Stresses color/texture cues and tests FL robustness under higher visual complexity.

5.3. Experiment

HyFLM is compared with representative personalized FL baselines, including global-only, partial-sharing, meta-initialization, hypernetwork, and diffusion-based aggregation methods. Baselines were chosen primarily from methods that are commonly adopted and reported in related work, to enhance the objectivity and reliability of our empirical evaluation. All models are trained in an identical computational and communication environment.

5.3.1. Baseline Methods

• Local—Each client independently trains its own model on local data without any communication or parameter sharing.
• FedAvg—The standard federated averaging algorithm with a single global model and aggregated each client model.
• FedPer—Partial personalization that shares a common feature extractor while each client keeps a private classification head.
• LG-FedAvg—Layer-grouped FedAvg updates lower layers globally while upper layers remain local, allowing architectural heterogeneity with minimal communication.
• pFedHN—Hypernetwork-based personalization where a server-side hypernetwork generates (parts of) client weights or adapters conditioned on a client embedding.
• pFedGPA—Diffusion-based generative parameter aggregation; the server learns a diffusion model over client parameters, performs parameter inversion, and synthesizes personalized weights by denoising sampling.

5.3.2. Implementation Details

For EMNIST and Fashion-MNIST we use a lightweight CNN (2conv + 1fc); for CIFAR-10 we use a CIFAR-CNN (3conv + 2fc). All baselines share the same backbone on a given dataset. CIFAR-10 uses standard augmentation.

Experiments are conducted with $M\in 10,20,100$ clients. In each round, only a part of clients is sampled (50% of clients when $M\le 20$ and 20% when $M=100$). Unless otherwise stated, training runs for 300 rounds in EMNIST, Fashion-MNIST, and CIFAR-10. Local epochs are set to $E\in 1,5,10$ depending on the dataset size (default $E=5$). We used a fixed batch size of 64 for all experiments.

All baselines are implemented using Python with unified optimizer configurations to ensure consistency across datasets. Unless otherwise required, clients adjust a learning rate of $0.05$ for EMNIST/Fashion-MNIST and $0.01$ for CIFAR-10, momentum $0.5$, and weight decay $5\times {10}^{-4}$. For CIFAR-10, a cosine learning rate decay is employed throughout the training. Overfitting is mitigated by early stopping with a patience of 10 based on the validation accuracy. All reported results correspond to experimental data over five runs with different initialization seeds.

HyFLM extends diffusion-based aggregation with MAC-guided trajectory control and a lightweight personalization head. We use $s=20$ denoising steps per client and the following MAC settings (also used in ablations): Fourier order $K=6$, bounds ${\gamma }_{min}=0$, ${\gamma }_{max}=2$, low-pass coefficient $\sigma =0.35$, and residual pull ${\alpha }_{\mathrm{res}}=0.05$. The progression-to-time mapping is linear within each local update $t=u/(U-1)$. During local training, MAC scales the message (gradient) by ${\gamma }_{\mathrm{msg}}\left(t\right)$ and applies a residual pull toward the start weights of the round with ${\gamma }_{\mathrm{res}}\left(t\right)$.

6. Results and Discussion

6.1. Experimental Result

The experiments are designed to comprehensively evaluate the proposed HyFLM framework across multiple dimensions of federated learning performance. In particular, we focus on three key aspects: (1) personalized accuracy across heterogeneous clients, (2) communication efficiency measured by convergence speed and payload cost, and (3) scalability with respect to the number of participating clients and dataset complexity. Through these evaluations, we aim to verify both the effectiveness and practicality of HyFLM in real-world federated settings.

6.1.1. Personalization Performance Comparison

We compare HyFLM with three benchmark datasets, EMNIST, Fashion-MNIST (FMNIST), and CIFAR-10, and representative FL baselines: FedAvg, FedPer, LG-FedAvg, pFedHN, and pFedGPA. As shown in

Table 2. Accuracy comparison of HyFLM and baseline methods on EMNIST, Fashion-MNIST, and CIFAR-10 with varying numbers of clients.

Method	EMNIST			Fashion-MNIST			CIFAR-10
# Clients	10	20	100	10	20	100	10	20	100
Local-only	71.75	71.26	74.84	86.24	85.24	86.14	65.45	65.21	66.12
FedAvg	72.43	73.44	75.25	82.23	84.96	84.15	65.27	68.47	70.55
FedPer	75.24	77.45	78.65	88.24	89.41	88.15	64.54	64.88	70.21
LG-FedAvg	73.42	71.21	76.54	86.12	86.45	84.45	65.15	65.65	66.24
pFedHN	77.34	78.12	77.15	86.80	86.90	86.25	71.66	75.12	77.10
pFedGPA	78.40	80.00	82.90	85.60	87.00	89.60	70.10	71.90	74.20
HyFLM	84.52	86.00	88.50	89.00	89.80	91.10	74.67	76.08	78.10

, HyFLM achieves the highest accuracy in each of the nine experimental settings (three datasets × three client counts). Its overall mean accuracy is approximately 86.8% across all configurations, substantially outperforming the strongest baseline.

HyFLM attains 78.4%, 80.0%, and 82.9% accuracy with 10, 20, and 100 clients on EMNIST. These results surpass pFedGPA’s 78.4%, 80.0%, and 82.9% under the same conditions. Other strong personalized baselines remain lower; for example, pFedHN tops out at ∼82.2% and FedPer at ∼78.7% (in the 100-client EMNIST case), both well below HyFLM’s 88.5%. Similarly, on Fashion-MNIST, HyFLM achieves 89.0%, 89.8%, and 91.1% for 10, 20, and 100 clients, outperforming FedAvg (82.23%, 84.96%, 84.15%) by approximately 6–9 percentage points in each case. In the more challenging CIFAR-10 dataset, all methods obtain lower accuracies; HyFLM reaches 74.7%, 76.7%, and 78.1% for 10, 20, and 100 clients, maintaining a slight edge over pFedGPA (70.10%, 71.90%, 74.20%). Other personalized approaches also struggle on CIFAR-10 (for example, FedAvg peaks at ∼70.55% with 100 clients), underscoring the challenge of this dataset and the competitiveness of HyFLM’s result.

Increasing the number of clients generally boosts performance for most methods, but HyFLM stands out in its scalability and robustness. HyFLM’s accuracy consistently improves as more clients participate, indicating effective utilization of larger distributed data and resilience to heightened data heterogeneity. For example, on CIFAR-10, HyFLM improves from 74.7% (with 10 clients) to 78.6% (100 clients), and on EMNIST from 84.5% to 88.5%, demonstrating solid gains with more clients. Although baselines like FedAvg and pFedGPA also improve with scale (e.g., FedAvg from ∼65% to ∼66.2% in CIFAR-10), their gains are more modest or plateau earlier. In fact, beyond 20 clients, pFedGPA shows diminishing returns on CIFAR-10 (e.g., 70.1% at 20 clients vs. 74.9% at 100). Moreover, pFedHN’s performance actually degrades at larger scales (on EMNIST it drops from 78.12% at 20 clients to 77.2% at 100). In contrast, HyFLM continues to improve and maintains high accuracy even with 100 clients. This underscores HyFLM’s robust personalization capability in large federations.

The characteristics of the data set further influence performance gaps. All methods achieve higher absolute accuracy on the simpler character/handwritten digit tasks (EMNIST and FMNIST) than on the complex CIFAR-10 images. For example, FedAvg reaches 82.23% on FMNIST (20 clients) but only 65.27% on CIFAR-10; similarly, pFedGPA achieves 86.6% versus 70.1%. HyFLM follows this trend, obtaining 89.0% on FMNIST but 74.7% on CIFAR-10 (100 clients). Importantly, HyFLM’s improvement over the best baseline is more pronounced in the easiest data sets (e.g., approximately +8–9% on FMNIST), while in CIFAR-10 the margin narrows to less than 1%. This suggests that while HyFLM excels across the board, its relative advantage is especially significant when task complexity is moderate.

HyFLM attains a higher precision and comparable or faster convergence than the diffusion-based baseline (pFedGPA) in all six comparisons from Figure 5. In EMNIST, HyFLM reaches 84.5% (10 clients, round 161), 86.0% (20 clients, round 127) and 88.5% (100 clients, round 270), outperforming pFedGPA’ 78.4%, 80.0%, and 82.9%, respectively – gains of roughly +6.1, +6.0, and +5.6 percentage points. In Fashion-MNIST, HyFLM achieves 88.95% (10 clients, round 156), 89.78% (20 clients, round 276), and 91.12% (100 clients, round 299), exceeding pFedGPA’s 85.6%, 87.0%, and 85.6% by approximately +3.3, +2.8, and +5.5 points. In all settings, the curve of HyFLM stays above the baseline early in the training and maintains a higher asymptote, indicating improved sample efficiency (fewer rounds to surpass the baseline plateau) and better final precision. In particular, the margin widens at larger client counts (100-client EMNIST/FMNIST panels), suggesting stronger robustness to heightened statistical heterogeneity. Therefore, the graphs corroborate the consistent accuracy advantage, stable convergence, and communication efficiency of HyFLM with increasing federation size.

6.1.2. Communication Efficiency and Scalability

In

Table 3. Accuracy and communication result in the experiment environment.

Dataset	#Clients	Method	Acc@Last (%)	R@95%	Payload/Round (MB)	Total MB
EMNIST	10	pFedGPA	78.4	145	2.37	236.6
		HyFLM	83.9	28	2.37	165.6
	20	pFedGPA	81.0	160	4.52	542.0
		HyFLM	85.8	32	4.52	361.3
	100	pFedGPA	83.4	240	21.72	3910.2
		HyFLM	88.1	45	21.72	2606.8
FMNIST	10	pFedGPA	88.7	122	0.37	31.8
		HyFLM	95.3	24	0.37	22.4
	20	pFedGPA	89.4	150	0.71	74.9
		HyFLM	97.9	19	0.71	53.5
	100	pFedGPA	91.0	184	3.43	514.9
		HyFLM	98.8	14	3.43	377.6
CIFAR-10	10	pFedGPA	72.5	260	3.12	436.4
		HyFLM	74.9	123	3.12	374.0
	20	pFedGPA	72.5	181	5.95	1071.1
		HyFLM	75.1	158	5.95	892.6
	100	pFedGPA	73.8	230	28.62	6582.5
		HyFLM	79.7	195	28.62	5437.7

, we summarize accuracy and communication cost under practical settings. Acc@Last denotes the final personalized test accuracy, and R@95% represents the number of rounds required to reach 95% of this value. The communication cost is quantified by the payload per round, defined as the total traffic from server broadcast and client uploads. The total communication volume (Total MB) is computed as $R@95\%\times \mathrm{Payload}/\mathrm{round}$. HyFLM shows better communication efficiency by reducing the payload per round compared to pFedGPA in Figure 5. Due to its hierarchical model-sharing strategy, HyFLM can transmit fewer or more compact parameters each round, so that when pFedGPA sends a full model update, HyFLM often needs only a smaller embedded update (as reflected by the lower Payload/round in Table 3). As a result, the total communication volume required for HyFLM to reach near-top accuracy is significantly lower. For example, on CIFAR-10 with 100 clients, The combination of faster convergence and reduced per-round payload yields approximately a 20–30% reduction in total data transfer compared to pFedGPA. Similar trends are observed on EMNIST and FMNIST, where HyFLM achieves smaller R@95% and Total MB, indicating more efficient bandwidth usage without sacrificing accuracy.

Both methods exhibit robust scaling as the number of clients increases from 10 to 100, albeit with some differences in behavior on simpler vs. more complex tasks. In general, adding more clients (which introduces more data diversity but also more heterogeneity) does not degrade final accuracy in these experiments—in fact, the EMNIST and FMNIST results show a slight increase in accuracy at 100 clients relative to 10. For example, EMNIST accuracy improves from the low–$80\%$ range with 10 clients to about $88.5\%$ with 100 clients for both algorithms, suggesting that a broader client participation helped the models generalize better. HyFLM handles this scaling gracefully: its accuracy scales at least as well as pFedGPA’s, while requiring fewer rounds to converge.

On the more challenging CIFAR-10 dataset (which has more complex data distributions), overall accuracy is lower for both methods, and the benefit of more clients is less pronounced. At 100 clients, pFedGPA and HyFLM reach high–$70\%$ accuracy on CIFAR-10, which is an expected drop compared to the easier datasets. Importantly, HyFLM’s performance relative to other baselines improves with scale and complexity. HyFLM consistently maintains the highest accuracy across all client scales on CIFAR-10, and its performance gap relative to strong baselines becomes more evident as the number of clients increases. This suggests that HyFLM is particularly effective in large-scale, heterogeneous environment.

In terms of communication scaling, as the number of clients increases, the number of rounds required to achieve R@95 tends to increase substantially. However, HyFLM can mitigates this effect. Even with 100 clients, HyFLM maintains a low R@95% and moderate communication cost, demonstrating better scalability. Overall, the results in Table 3 show that HyFLM not only matches or exceeds pFedGPA in accuracy across different workloads, but does so with greater communication efficiency and resilience as the federated network scales up.

6.2. Discussion

Interpretation of Key Results: Our method (HyFLM) consistently shows superior personalization accuracy across varying degrees of non-IID data distribution. Notably, it outperforms strong baselines not only in final accuracy but even from the early rounds of training. This suggests that the diffusion-based inference trajectory optimization and hypernetwork-driven updates in HyFLM enable faster initial alignment to each client’s data while preserving global consistency, maintaining a balance that baseline methods struggle to achieve so early. Beyond accuracy, communication efficiency is markedly improved. HyFLM reaches convergence in fewer communication rounds thanks to the accelerated learning provided by the MAC module and hypernetwork. Moreover, the updates exchanged each round are smaller and more modular: instead of sending full model parameters, clients and server share compact hypernetwork outputs or anchor adjustments, significantly reducing the payload. Equally important, HyFLM demonstrates stability at scale. In experiments with up to 100 clients, a regime where many federated methods degrade due to statistical heterogeneity and client drift, HyFLM maintained stable and high performance. We attribute this robustness to the globally guided diffusion process, which MAC drives client updates and mitigates divergent drifts even as the client population grows.

Relation to Literature: HyFLM effectively integrates personalized federated learning with diffusion-based aggregation by addressing key limitations in prior approaches. Traditional personalized FL method focus on local adaptation but lack global coordination, often resulting in client drift. HyFLM mitigates this by using diffusion-based inference trajectory optimization, offering globally informed yet client-specific updates. Unlike prior diffusion FL methods, which apply uniform trajectories, the HyFLM MAC module enables per-timestep, per-client modulation, enhancing personalization precision. Additionally, compared to hypernetwork-based methods, HyFLM offers improved efficiency and adaptability through its lightweight hypernetwork coupled with evolving diffusion anchors. This unified design achieves scalable and fine-grained personalization with better alignment between clients.

Overall, We have highlighted why HyFLM achieves better personalization (through globally anchored yet client-specific updates), under what conditions it yields the greatest benefits (moderate heterogeneity and network size), and how it connects to and improves upon prior art in federated learning. These insights clarify the contexts in which the proposed model excels, as well as its limitations, guiding practitioners on when HyFLM is the most appropriate choice and where future refinements could further enhance its practicality.

6.3. Future Work

While HyFLM has shown strong scalability and effectiveness across a wide range of standard benchmark conditions, additional research is necessary to extend its practicality to ultra-large federated networks involving thousands or even tens of thousands of clients. In such extreme-scale scenarios, personalized weight generation via a server-side hypernetwork may lead to increased computational load and communication overhead, especially when handling high-dimensional client embeddings or supporting complex personalization adapters.

To address this, we will consider to incorporate clustering or grouping techniques that partition clients based on similarity in data distribution, resource profile, or model updates. Such stratification would enable the hypernetwork to generate shared personalization templates per group, thus amortizing the cost of per-client weight generation and reducing redundant computation. This approach has the potential to preserve the quality of personalization while significantly improving communication and computational efficiency. In addition, our goal is to investigate hierarchical personalization strategies, where coarse-grained anchors are produced at the group level and refined locally with lightweight updates.

7. Conclusions

This paper proposes HyFLM, a personalized federated learning framework that combines diffusion-based model aggregation with hypernetwork-driven personalization to address challenges posed by non-IID data and client heterogeneity. We hypothesized that federated learning performance could be enhanced through trajectory optimization theory applied within diffusion-based model aggregation. To validate this hypothesis, we implemented a novel framework incorporating multidimensional trajectory control. Extensive experiments confirmed that this approach improves both personalization accuracy and communication efficiency under non-IID conditions. Empirical evaluations demonstrated that HyFLM outperforms strong baselines across varying federation sizes and data skews. HyFLM achieved higher personalized accuracy, faster convergence, and reduced communication volume, particularly excelling in large-scale or highly skewed settings. These results establish HyFLM as a robust and scalable solution for a federated learning environment. Future directions include scaling HyFLM to diverse modalities and larger federations and investigating theoretical guaranties, with the underlying principles of the framework offering a foundation for future advances in personalized and communication-efficient FL systems.