HyperPrism: An Adaptive Non-linear Aggregation Framework for Distributed Machine Learning over Non-IID Data and Time-varying Communication Links

Authors: haizhou  du
yijian  chen
ryan  yang
yuchen  li
linghe  kong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 17:02:31

Abstract: While Distributed Machine Learning (DML) has been widely used to achieve decent performance, it is still challenging to take full advantage of data and devices distributed at multiple vantage points to adapt and learn, especially it is non-trivial to address dynamic and divergence challenges based on the linear aggregation framework as follows: (1) heterogeneous learning data at different devices (i.e., non-IID data) resulting in model divergence and (2) in the case of time-varying communication links, the limited ability for devices to reconcile model divergence. In this paper, we contribute a non-linear class aggregation framework HyperPrism that leverages distributed mirror descent with averaging done in the mirror descent dual space and adapts the degree of Weighted Power Mean (WPM) used in each round. Moreover, HyperPrism could adaptively choose different mapping for different layers of the local model with a dedicated hypernetwork per device, achieving automatic optimization of DML in high divergence settings. We perform rigorous analysis and experimental evaluations to demonstrate the effectiveness of adaptive, mirror-mapping DML. In particular, we extend the generalizability of existing related works and position them as special cases within HyperPrism. Our experimental results show that HyperPrism can improve the convergence speed up to 98.63% and scale well to more devices compared with the state-of-the-art, all with little additional computation overhead compared to traditional linear aggregation.

Main paper: 
 introduction: 
 The proliferation of edge devices, such as mobile phones, wearable devices and unmanned aerial vehicles, has resulted in an exponential growth in diverse data types (e.g., images, sound and text). Addressing this surge necessitates advanced techniques capable of accurately, quickly, and practically processing vast amounts of data through efficient and scalable algorithms. Distributed Machine Learning (DML) has gained significant traction in recent years as a strategy to minimize data transfer and computational costs by bringing computation closer to the data sources. It has become a natural solution for scaling up machine learning while preserving data privacy in various domains (e.g., large language model training [39, 33] , autonomous driving [60, 37] , military applications [56] , web search [29, 30, 28] and recommendation [32, 27, 31] ).While existing DML attempts have made significant progress, two technical barriers remain in realistic scenarios. First, data heterogeneity is one of the most critical concerns. Specifically, in a DML system where mobile devices serve as computing nodes, the data generated is typically non-independent and identically distributed (non-IID), which means the data distribution may be unbalanced across different classes or categories. Consequently, the model may exhibit bias toward the majority class or dominant data patterns, leading to suboptimal performance on the minority class or rare patterns. Second, traditional DML methods usually deliver poorer performance under the time-varying communication condition. Specifically, devices holding critical data may unexpectedly go offline or become out of range during the training, causing changes in the communication links. These interruptions can cause information loss in the aggregation, making local models no longer interchangeable and drifting away from the global model. We can summarize these powers of deviation and drift as "divergence forces", which drastically slows down the convergence and significantly impacts the efficiency and effectiveness of the training process in DML.In order to tackle the above issues, we present a novel decentralized DML framework HyperPrism, which utilizes mirror descent [3] and employs adaptive mapping functions to project models into a mirror space, with both aggregation and gradient steps then carried out in the mapped space. Moreover, HyperPrism leverages the concept of "weighted power means" (WPM) as the aggregation function, raising the model parameters to the power of p, and uses HyperNetworks [12] to adaptively adjust the power degree p. In summary, the main contributions are summarized as follows:• We study the problem of divergence forces in decentralized DML, where we particularly focus on two technical barriers due to data heterogeneity (i.e., non-IID data) and time-varying communication links. To the best of our knowledge, it is the first work to simultaneously address the challenges of non-IID data and time-varying communication links in realistic DML scenarios. • We propose a non-linear class aggregation DML framework HyperPrism based on Kolmogorov Means, which can also be seen as mapping models to mirror space for aggregation, and we instantiate the means with an adaptive p power function to enhance the convergence speed and scalability. HyperPrism achieves superior performance in its dependence on the number of devices m and the power degree p, upgrading from m √ m to m p √ m. This achievement also reduces the optimality gap [2] from √ m to p √ m. • We conduct rigorous analysis and prove that the loss bound of HyperPrism is O(( m P +2 T )1 P +1 ). In cases where few communication epochs can occur (i.e., T ≤ m), employing a larger value of p yields improvements over traditional linear aggregation. Our theoretical results are consistent with state-of-the-art bounds in distributed gradient/mirror descent and single-device mirror descent.• We carry out comprehensive experiments to assess the performance of HyperPrism framework. The experimental results demonstrate that HyperPrism achieves a remarkable acceleration in convergence speed with improvements of up to 98.63%. Moreover, HyperPrism also shows increased scalability in settings characterized by time-varying communication. 
related work: 
 Decentralized DML with Data Heterogeneous. Numerous studies have addressed non-IID data using linear solutions, such as local fine-tuning of a global model [4, 9, 8, 10] , personalization in Federated Learning (FL) as a meta-learning objective [15] , knowledge distillation [61] , and prototype aggregation [53] . Furthermore, [17] . De Vos et al. [5] introduce Epidemic Learning in decentralized learning, leading to faster convergence and improved performance by randomly changing topologies. Nedic et al. [44, 42] tackle the DML with topology dynamicity from the consensus perspective. They propose a model aggregation method for agents in a time-varying network topology to collaboratively solve a convex objective function, with a convergence guarantee. Moreover, future generation DML systems [20, 19, 7] also consider these time-varying communication settings as an important research area, where the naturally time-varying connectivity among pairs of mobile devices or mobile devices and edge servers will be dictated by their physical proximity.In this work, we focus on solving two barriers due to data heterogeneity (i.e., non-IID data) and time-varying communication links and conduct rigorous analysis to show that prior methods suffer from lower accuracy, higher loss, and slower convergence speed over time-varying networks than under a fixed topology. Additionally, their methods use linear averages for aggregation, which our framework encompasses as a particular case.3 Problem Formulation and Preliminaries 
problem formulation: 
 We consider decentralized DML with m devices over time-varying communication links represented by a directed graphG(t) = (V, ε(A (t) )), where V denote vertices andEQUATIONmm ] is the weight matrix of the topology graph. Then the ε(A (t) ) denotes the set of directed time-changing edges between vertexes. In each time interval, we assume that the communication links between devices are symmetrical and restricted, causing random disconnections and reconnections. Each device i has its own private dataset, denoted by D i . Let w = (w 1 , w 2 , . . . , w m ) denote the collection of all local models, where w i is the model held at device i. Then, Device i constructs its local loss function asf i (w) = E ζi∼Di [F(w i ; ζ i , G(t))].Then, devices aim to solve the following optimization problem collaboratively:EQUATIONwithout sharing local data (i.e., without revealing f i ).In recent years, model merging has developed non-linear aggregation methods. One motivation for non-linear aggregation is the decreased variance in the original parameter space. HyperPrism is designed to combat stale gradients from diverged models with synchronization.In existing work, synchronization is achieved with an aggregation mechanism, usually a mean or median, with clipping. Inspired by similarities between Mirror Descent [46] and Quasi-Arithmetic Means [16] , also known as Kolmogorov Means, of the form f -1 1 n n k=1 f (x k ) , HyperPrism maps models to the dual domain before averaging to better align with the geometry of the objective function. In implementation, we focus on the special case of ϕ(w) = 1 p+1 ∥w∥ p+1 , transforming models as w → w p , where HyperPrism therefore replaces linear means with the following weighted power mean:EQUATIONwhere a (t) ij is the aggregation weights of device j in device i at t round. The WPM is a special case of the general strategy of averaging in a mirror descent dual space, but it is particularly useful due to its ease of computation and convexity guarantees. In general, the aggregation step of HyperPrism is the following Equation (3):EQUATIONBut we choose to focus on ϕ(w) = 1 p+1 ∥w∥ p+1 which gives ∇ϕ(w) = w p (termwise exponentiation). This allows for simple run-time computation. For more intuition, in the p = ∞ limit the distributed averaging consensus problem [47] turns into an easier "distributed maximum problem," which is what allows for HyperPrism's tighter synchronization.The phrase "distributed mirror descent" can have multiple interpretations. In the existing literature, it typically refers to the process of taking local mirror descent steps and then linearly aggregating [23, 48, 6, 59] In the case of HyperPrism, mirror descent is transformed into the distributed algorithm in a different manner. It still takes local mirror descent steps but also uses the aggregation function in Equation ( 3). This approach proves to be more effective because, under the optimal model w * , distributed mirror descent with Equation (3) remains stable, while the linear aggregation functions do not. In other papers' analysis (e.g., [11] , [58] ), this inaccuracy is not obvious since the dominant terms are often related to communication costs, but in the perfect communication case, using a generalized mean is necessary for exact theoretical convergence.Hypernetworks (HNs) are deep neural networks used to generate weights of other target networks [12] . HNs can learn the mapping relationship between the embedding vector and the target network, and adaptively generate the target network based on the input. Shamsian et al. apply HNs in federated learning to generate personalized model parameters [50] . Ma et al. present pFedLA [38] using HNs to generate aggregation weights of each model layer in personalized federated learning. Previous works [53, 55] found the different layer parameters of the local model have different impacts on model aggregation. For example, the locally learned feature representations are prone to over-fitting and thus cannot generalize well when each device only has insufficient data. In HyperPrism, the mapping function chosen is a form of the Bregman divergence, specifically ∇ϕ(w) = w p , where p represents the degree of WPM. Our experiments indicate that different functional layers of the model respond differently to the degree of p. Increasing p has a more significant impact on accelerating convergence in linear or fully-connected layers. This observation inspires us to select the appropriate p for different functional layers adaptively. Through the chain rule [38] , we demonstrate that HN can effectively correlate the mirror mapping with the objective function, enabling it to discover the optimal p for each device's model parameters in each round. 4 The Design of HyperPrism Framework 
overview: 
 With the above preliminaries and motivation presented, we now give the design of HyperPrism framework. The framework assumes that each device i maintains its local model denoted by w i ; for simplicity of specification, we use a round-based specification, in which the model holds by device i after t communication rounds is denoted asw (t) i .The framework first chooses the aggregation weights a. Each a where eEQUATIONEQUATION(t) ij = 1 N (t) i +1when i, j are connected, otherwise e(t) ij = 0. N (t) iis the number of neighboring devices i at t-th round. a(t)ij generated in this way satisfy Assumption 5.1. Note that the connectivity e (t) ij is determined by underlying communication systems, which consider both feasibility and security requirements (e.g., low probability of detection). HyperPrism applies to generic DML, the parameter server and all-reduce frameworks can be seen as special cases (where all a (t) ij values equal 1 m ). In summary, the core of HyperPrism is that at every round t, device i computes its local model for the next round t + 1 as:EQUATION)where ∇ϕ is the mapping function with a selected degree of p for each round.To determine the p of various layers at each round, HyperPrism establishes a relationship between the optimization problem and the degree of p by utilizing hypernetworks, which adaptively choose the optimal p for different layers to achieve the best performance.Without loss of generality, HyperPrism also decouples the local model into the representation and decision parts. The representation part θ includes components like convolutional and embedding layers. The decision part φ includes components like fully-connected layers. Thus, for device i, we have w i = {w θ,i , w φ,i }. Each device i holds a local hypernetwork HN i and a randomly generated embedding vector v i . Every hypernetwork consists of several fully-connected layers and employs the softmax for the output layer. HN i takes v i and gradient of the local model as input, then output p i = {p θ,i , p φ,i } for w θ,i and w φ,i parts, respectively. Then the mapping function evolves to ∇ϕ(w i ) = {(w θ,i ) p θ,i , (w φ,i ) pφ,i }. The hypernetwork on device i can be defined asEQUATIONwhere ψ i denotes the parameters of HN i . Hence, a new objective function can be derived from the original problem asEQUATIONHyperPrism can transform the optimization problem for model parameters w i into the HN's v i and ψ i . HNs adaptively output p θ,i and p φ,i based on the input, and simultaneously update both v and ψ Algorithm 1 HyperPrism Framework. m }. 3: Initialize all devices' models w (0) , hypernetworks ψ (0) , and embedding vectors v (0) . 4: for t = 0 to T do 5:for each device i in parallel do 6:Local Update:w (t) i ← w (t) i -η∇f (t) i . 7: Compute p (t) θ,i and p (t) φ,i by HN i (v (t) i ; ψ (t) i ). 8:Compute ∇ϕ(w(t) i ) ← {(w (t) θ,i ) p (t)θ,i , (w(t) φ,i ) p (t)φ,i }.Send ∇ϕ(w Compute weights a (t) ij as Equation ( 4), (5) .(12:Aggregate received models using a (t) ij .13:Update w (t+1) i according to Equation ( 6).14:EQUATIONas Equation ( 11), (12) .15:end for 16: end for 17: return w(T ) i .by gradient descent at each round. Specifically, the gradient of v i and ψ i can be computed based on the chain rule [38] asEQUATIONEQUATIONThen, v i and ψ i can be represented asEQUATIONEQUATION)Algorithm 1 demonstrates the full procedure of HyperPrism. In our real deployment, we adopt the asynchronous protocol. 
analytical results: 
 We rigorously analyze the properties of HyperPrism. We fully introduce all assumptions, analyze the convergence behavior of HyperPrism, and finally compare HyperPrism to previous works.Assumption 5.1 (Connectivity). The network graph G(t) = (V, ε(A (t) )) and the connectivity weight matrix A (t) satisfy the following:• A (t) is doubly stochastic for all t ≥ 1; that is m j=1 a (t) ij = 1 and m i=1 a (t) ij = 1.• There exists a scale ζ > 0, such that a(t) ij ≥ ζ for all i and t ≥ 1, if {i, j} ∈ E t . • There exists an integer B ≥ 1 such that the graph (V, E kB+1 ∪ • • • ∪ E (k+1)B ) is stronglyconnected for all k ≥ 0. Definition 5.2 (Bregman Divergence). The Bregman Divergence of a function ϕ is defined asEQUATION)Note that for ϕ(x) = ∥x∥ 2 , D ϕ (x, y) = ∥x -y∥ 2 .Definition 5.3 (Uniform Convexity). Consider a differentiable convex function ϕ : R d → R, an exponent r ≥ 2, and a constant σ > 0. Then, ϕ is (σ, r)-uniformly convex with respect to a ∥EQUATIONNote that for r = 2, this is known as strong convexity. This assumption also implies that D ϕ (x, y) ≥ σ r ∥x -y∥ r . Assumption 5.4 (Smooth Gradient). Assume that the functions f i are convex with its gradients ∇f i (•) satisfying L-Lipschitz continuity [59] , namely:EQUATIONfor all x, y pairs. This final assumption is used in nearly every mathematical analysis of DML, including [54, 3, 52, 59, 43, 51, 49] . In our analysis, it is used to bound the distance between local models.In our experiments, we focused on the weighted power mean, generated by using ϕ(x) = 1 p+1 ∥x∥ p+1 which gives ∇ϕ(x) = x p . Such ϕ is uniformly convex as seen in Proposition 5.5, and the rest of the analysis will be generalized to all uniformly convex ϕ. Proposition 5.5 (Uniform Convexity of Power Functions). For p ≥ 1, the function φ p (x) = 1 p+1 ∥x∥ p+1 is uniformly convex with degree p + 1. This is because ∥∇φ p (x) -∇φ p (y)∥ = ∥x py p ∥ ≥ 1 2 p-1 ∥x -y∥, and as a corollary we haveEQUATIONGiven the function ϕ, HyperPrism instructs local models to take a WPM of received models. Using the Connectivity Assumption (Assumption 5.1), we can bound the distance between local models. Lemma 5.6 (Consensus). Under Assumption 5.1, for each device i, after t rounds:∥∇ϕ(wEQUATIONwhereEQUATIONi ) -f * and ζ is a constant related to the graph connectivity. Theorem 5.7 (Convergence Behavior). Consider a ( 1 2 p-1 , p+1) uniformly convex ϕ and the sequence w (t) i under Algorithm 1 with constant step size η. Then, under Assumptions 5.1, and 5.4, if x * is the value that minimizesEQUATIONIn this Theorem 5.7, the first term uses Lemma 5.6, and is caused by the differences between local models held at different devices. This is the effect of the diverging forces mentioned in earlier sections. Note that ϑ and κ are constants related to the graph connectivity, reflecting the influence of time-varying communication links. The second term represents the error caused by a non-zero learning rate in the Distributed Mirror Descent process, and the third term represents the lingering effects of the initialization. For more proof details, please refer to the Appendix. The first term becomes O(m • P √ mη) because since κ < 1, the κ t term goes to 0 as t gets large. The second term is O(m P √ η) is smaller than the first term, so we drop it. The final term is clearly O( m ηT ).Corollary 5.9. For η = O( p+1 T p m ), the upper bound becomes O(m p √ ηm + m ηT ) = O( p+1 m p+2 T ). P √ ηm + m ηT ) O( P +1 m P +2 T ) NA Bubeck [3] (m = 1, r = 2) O(η + 1 ηT ) O( 1 √ T ) Yes Srebro [52] (m = 1) O(η 1 r-1 + 1 ηT ) O( 1 r √ T ) Yes Nedic [44] (GD → r = 2) O(ηm 2 + m ηT ) O( m √ m √ T ) Yes Yuan [59] (r = 2 MD) O(ηm 2 + m ηT ) O( m √ m √ T )YesTable 1 summarizes the big O notation convergence rates of f (w (t) ) -f * . Previous works also assume bounded gradients. Our analysis recovers the same bounds as the state-of-the-art in DML and single-device mirror descent with generic uniform convexity assumptions. Nedic [44] is standard gradient descent, and thus has ϕ(x) = 1 2 ∥x∥ 2 , which is (1, 2)-uniformly convex, corresponding to p = 1. Yuan [59] considers distributed mirror descent under strong convexity (r = 2) equivalent to p = 1, and based on their analysis, the bound should be O(ηm 2 + m ηT ).6 Evaluation 
experimental setup: 
 The experimental platform consists of 8 Nvidia Tesla T4 GPUs, 4 Intel XEON CPUs, and 256GB of memory. All the models and training scripts are implemented in RAY [40] and PyTorch [24] .HyperNetworks Setup. We construct a hypernetwork model comprising three fully-connected layers and two additional output layers activated using softmax. The outputs of the fully-connected layers are fed into each of the two output layers to generate the degree of P for various parts.Time-varying Communication Links Setup. We employ the NS3 platform [13] to simulate realistic time-varying communication environments consisting of multiple distributed devices. Each device is configured with the WiFi 802.11a protocol and communication among themselves in Ad-Hoc mode.To quantify the degree of connectivity of the communication links, we define the topology density as the ratio of available tunnels to the total tunnels.Models and Dataset. We use MNIST [22] and CIFAR-10 [21] datasets distributed among devices in non-IID settings. We construct two models based on layer functionalities. The Logistic Regression (LR) [14] model consists solely of linear layers used for the MNIST dataset. On the other hand, the CNN model consists of both convolutional and fully-connected layers used for CIFAR-10.Non-IID Data Partitioning. To distribute datasets in a non-IID fashion, we employ Dirichlet distribution [35] to allocate all samples among devices. The Dirichlet Degree α is used to control the non-IID degree. The α = 0.1 represents the extreme scenario where each device possesses samples from only one class, while α = 10 equals the IID scenario. These distributions reflect a challenging and realistic training environment.Metrics. We consider two metrics to measure the performance of HyperPrism.• Average Accuracy. We evaluate the performance of the local model per device using a global test set that contains samples with all categories. The average Top-1 accuracy of all devices is calculated in each round to measure overall performance and convergence rate.• Convergence Speed. We track the loss of each round and the number of rounds to investigate the rounds needed to reach the convergence point for specific accuracy and loss.Baselines. We conduct a comparative analysis of HyperPrism with state-of-the-art methods for DML in non-IID data and time-varying communication links, including p = 1 methods, which is one of the most influential and widely applied works in DML studies. SwarmSGD [41] , Mudag [57] and ADOM [20] . To ensure a fair comparison, we made minor adjustments to each baseline method.• p=1. We define it as a class of methods using a linear aggregation function, including DPSGD [34] , and its variants [43, 45] . These methods can be seen as special cases of HyperPrism without the mirror mapping process.• SwarmSGD. We set the number of local SGD updates equal to 1, where the selected pair of devices performs only a single local SGD update before aggregation. It can also be viewed as a special case of HyperPrism where the topology density is very low.• ADOM and Mudag. We set the condition number k to 10, and the number of features d equals the number of classes. The gossip matrix W (t) at round t is chosen to be the Laplacian of time-varying communication links divided by its largest eigenvalue.Hyperparameters. For all experiments, the learning rate and batch size are both fixed at 0.01 and 128. We generate time-varying communication graphs with different sizes and densities and evaluate 100 rounds total. The graph changes every round. Impact of different degrees of p. We present the impact of varying degrees of p in HyperPrism on model performance in Figure 2 . In Figure 2 (a), it is evident that the model converges more swiftly and attains greater accuracy as p grows larger. Figure 2 (b) demonstrates that accuracy exhibits a significant fluctuation with different p. This underscores the substantial impact of the p value selection on the performance of HyperPrism. These findings highlight the importance of choosing an appropriate value for p in HyperPrism to achieve optimal performance.Comparison of time cost. In HyperPrism, each hypernetwork contains only 3 linear layers with 64 nodes per layer to ensure a minimal extra computational resource cost. We record the time cost in Figure 2 (c). Although HyperPrism does result in a higher time cost per iteration, it notably decreases the total number of rounds required for convergence, thereby reducing the overall time needed to achieve a specific accuracy.Performance of HyperPrism. To showcase the practicality of HyperPrism, we consider the basic configuration with Dirichlet = 0.1, density = 0.5, and m = 50 devices. The accuracy results for all benchmarks and HyperPrism are presented in Figure 3 . The convergence speed is summarized in (c) CNN, Dirichlet=10 a remarkable superior performance over state-of-the-art baselines with convergence accuracy and convergence speed improvements of up to 4.87% and 98.63%, respectively.Impact of non-IID. To investigate the impact of the non-IID Dirichlet degree on HyperPrism, we experiment with various α = 0.1, 1, 10. The corresponding results are presented in Figure 3 , and Table 2 . All methods exhibit poorer performance as the non-IID degree becomes more extreme, which aligns with common intuition. However, HyperPrism demonstrates enhanced stability and faster convergence speed, especially at highly non-IID degrees.We further evaluate the performance of HyperPrism with varying scales m ∈ {20, 50, 100}. The results are summarized in Table 4 . It can be noticed that most of the baseline methods exhibit deteriorating performance as the number of devices increases. The ADOM even barely converges under 100 devices. In contrast, HyperPrism is minimally affected by the scale and maintains superior acceleration and model performance.Communication Graph Densities. To further analyze the impact of connected densities on model performance, we present the performance and convergence speed of HyperPrism with various density ∈ {0.2, 0.5, 0.8} in Table 3 . In the extreme non-IID case, the performance of baselines deteriorates as the communication becomes denser. Particularly, ADOM exhibits significant fluctuations at a density of 0.8. However, HyperPrism maintains better performance across different densities. This can be attributed to the fact that as the communication becomes denser, the information exchanged between devices becomes more complex. Given HyperPrism's resilience to non-IID scenarios, it maintains good performance in such cases. 
conclusion: 
 In this work, we studied the important problem of divergence forces in decentralized DML, due to data heterogeneity (i.e., non-IID data) and time-varying communication links. We propose a non-linear class aggregation DML framework with adaptive Kolmogorov Means for aggregation to enhance the convergence speed and scalability. HyperPrism achieves superior performance in its dependence on the number of devices m, improving from m √ m to m p √ m, and achieving optimality in the limit p → ∞. We also conduct rigorous analysis and demonstrate that the loss bound of HyperPrism is O(( m P +2 T )1 P +1 ). In cases with few communication epochs (i.e., T ≤ m), employing a larger value of p yields improvements over traditional linear aggregation. Our theoretical results are consistent with state-of-the-art bounds in distributed gradient/mirror descent and single-device mirror descent, all under a general uniform convexity assumption. To verify the effectiveness of HyperPrism, we carry out extensive experiments that demonstrate that HyperPrism achieves a remarkable acceleration in convergence speed with improvements of up to 98.63%. Moreover, HyperPrism shows increased scalability.8 Appendix 
main notations: 
 We summarize the main notations in Table 5 . The learning rate G(t)The time-varying network topology graph consisting of V and ε VThe set of vertexes (i.e., edge devices)ε(A (t) )The set of directed time-changing edges between vertexes A tThe weight matrix of the topology aWe illustrate how HyperPrism leverages weighted power mean (WPM) to facilitate more efficient aggregation in Figure 4 . By examining the distribution of model parameters, we observe that the traditional linear averaging model (w/o WPM) has approximately 11.65% of parameters lying within the range of [-0.01, 0.01] after 80 rounds. In contrast, when utilizing WPM with degrees of p = 9 and p = 15, only 1.67% and 1.21% of parameters, respectively, fall within the same range. This indicates that WPM enables the model to effectively preserve a broader range of features. Consequently, HyperPrism can extract more information from the model parameters during aggregation, leading to enhanced performance. These results underscore the effectiveness of WPM in enabling HyperPrism to capture a wider range of features and facilitate a more informative aggregation. To aid analysis, introduce the sequence y i,j ∇ϕ(wEQUATION) = ∇ϕ(yEQUATIONAdditionally, we can defineEQUATION8.4 Proof of Lemma 5.6Proof. Define a transition matrix Φ(t, s) = A(t)A(t -1) • • • A(s + 1)A(s). Then, under Assumption 5.1, Corollary 1 in [43] states thatEQUATIONwhere ϑ and κ are defined in the lemma statement.We are able to write out a general formula for ∇ϕ(w(t+1) i ):∇ϕ(wEQUATIONas well as ∇ϕ(w (t) ):EQUATIONIt is known that for a L-Lipschitz continuous function f , if x * is the optimum of f , then:∥∇f (x) -∇f (x * )∥ ≤ 2L(f (x) -f * )Then, ∇ϕ(w(t) i ) -∇ϕ(w (t)) can be bounded by applying the Triangle Inequality and Equation (22) :EQUATIONThen, plugging in k = 0 givesEQUATIONFinally, shifting t down by 1 gives the desired bound.8.5 Proof of Theorem 5.7Firstly, note that under Algorithm 1 and Assumption 5.4, Weighted AM-GM gives:(r -1) • r-1 1 σr r-1 m • (mηG l ) r + m • σ r ∥y (t) -y (t+1) ∥ r ≥ r • ( r-1 1 σr r-1 m • (mηG l ) r ) r-1 • (m • σ r ∥y (t) -y (t+1) ∥ r ) 1 1 r = r • 1 r r • (mηG l ) r • ∥y (t) -y (t+1) ∥ r 1 r = (mηG l ) • ∥y (t) -y (t+1) ∥.We may follow the main line of reasoning that proves Theorem 5.7.Proof. We prove bounds for generic x. Note that w(t) = h -1 ( 1 m m i=1 h(w (t) i )). Then, we get:EQUATIONCauchy's inequality can bound the first term. The third term can be manipulated using Equation (20) . Then, combined with f i (w (t) ) ≤ f i (w(t) i ) + G l ∥w (t) -w (t) i ∥, we get η • (F (w (t) ) -F (x)) ≤ (2 • m i=1 ηG l • ∥w (t) i -w (t) ∥) + mηG l • ∥w (t) -w (t+1) ∥ + m i=1 ⟨∇ϕ(y (t) i ) -∇ϕ(w (t+1) i ), w (t+1) -x⟩.The factor of 2 comes from the 2nd term of Equation ( 27) added to the error from f i (w (t) ) ≤ f i (wEQUATIONi ∥. Note that the last term is equal to m • ⟨w (t) -w (t+1) , w (t+1) -x⟩. This is also equal to m(D ϕ (x, w (t) ) -D ϕ (x, w (t+1) ) -D ϕ (w (t+1) , w (t) )) by the Triangle Inequality for Bregman Divergences. We also substitute Claim 1 to replace the second term, so the value isEQUATIONBut, by uniform convexity, D ϕ (w (t+1) , w (t) ) ≥ σ∥w (t+1) -w (t) ∥ r , and thus this is alsoEQUATIONThen, taking the sum over T givesEQUATIONDividing through by ηT , substituting σ = 1 2 p-1 and r = p + 1, and substituting x = x * and Lemma 5.6 gives the desired result. EQUATIONThis allows us to see the weighted power mean as a way to decrease skew; this theorem is relevant in the setting of Lemma 5.6.EQUATIONThe1 p 1 m x p i 1 pterm is bounded, so the error is proportional to O( ϵ p ). 
 
 Appendix: The aggregation weights of device j in device i at t round D iThe private dataset of device i wThe set of all local models w iThe model held at device i f i (w) Device i's local loss functionF (w) F (w) = m i=1 f i (w) w (t) iThe aggregation step of device i:w(t) i = [∇ϕ] -1 ( m j=1 a (t) ij ∇ϕ(w (t) j )) ϕ(w) Define the function ϕ(w) = 1 p+1 ∥w∥ p+1 ∇ϕ(w)The mapping function ∇ϕ(w) = w p pThe adaptive degree of the WPM e(t) ij e (t) ij = 1 N (t) i +1when i, j are connected, otherwiseN (t) iThe number of devices neighboring device i at t-th round θ, φ ∇ϕ(w i ) = {(w θ,i ) pθ,i , (w φ,i ) pφ,i } HN i (v i ; ψ i )The Hypernetwork on device i ψ iThe parameter of HN i ζ A scale related to connectivity L Smoothness of gradient ∇f G l Upper bound on gradient derived from L D ϕ (x, y)The Bregman Divergence of a function ϕ σ, r (σ, r)-uniformly convex y (t) i ∇ϕ(y(t) i ) = m j=1 a (t) i,j ∇ϕ(w (t) j ) w (t) w (t) = [∇ϕ] -1 ( 1 m m i=1 ∇ϕ(w (t) i )) ϑ, κConstants related to graph connectivity, containing ζ Φ(t, s)A transition matrix Φ(t, s) = A(t)A(t -1) • • • A(s + 1)A(s) τ, k round numberAs a starter, we first prove a useful lemma on two real numbers with different signs. Note that all variables used in this subsection are generic ones not tied to the HyperPrism-based mechanism. Lemma 8.2 (Two Numbers with Different Signs). Given any two real numbers x > 0, y < 0, an odd integer p ≥ 1, and 0 ≤ α ≤ 1, we haveEQUATIONProof. Denote C = 1 2p and rewrite the right-hand side (RHS) of Equation (34) as C(αx + (1α)y) + (1 -C)y. We first note that we can scale both x, y by 1 |y| such that y = -1. As such, we can simplify our proof to only focus on the case of y = -1. We will use shorthand ∆, and rewrite asEQUATIONWe then consider two cases.Case 1: x ≤ 2p -1. Since -Cα(x + 1) ≥ -Cα(2p) = -α ≥ -1, by Bernoulli's inequality, we haveEQUATIONSince x+1 2 ≤ x p + 1, which can be verified by casework on x ≥ 1 or x ≤ 1, thenEQUATIONWe then finish this case by taking the power of 1/p on both sides.Case 2: x ≥ 2p -1. We will denote the expression on the left-hand side and right-hand side of Equation (34) with LHS and RHS, respectively. First note that both sides of the inequality in Equation ( 34) are -1 when α = 0 and y = -1. Next, we will show that the derivative with respect to α is always larger for the LHS of Equation ( 34) when x ≥ 2p -1. The LHS's derivative w.r.t α is clearly constant. Also, the d dα RHS clearly has local maxima at α = 0, 1. At α = 0, we clearly have that since x ≥ 2p -1, we have For α = 1, since x ≥ 1 2p (x + 1) -1 ≥ 0, we haveEQUATIONThus, for all 0 ≤ α ≤ 1, d dα LHS ≥ d dα RHS, and they are equal at α = 0, so we always have LHS ≥ RHS for 0 ≤ α ≤ 1 when x ≥ 2p -1, and we are done.Next, consider a list of real numbers x i , i = 1, . . . , m. We let M = min x i and U = max x i . We then prove the following lemma. Proof. The lower bound in Equation ( 42) is a direct result from the power mean inequality. As for the upper bound, we have where the last step uses the generalized AM-GM inequality.We now tie these two Lemmas together to prove Theorem 8.4, which to the best of our knowledge is a novel inequality on weighted power mean. Proof. First, note that if all x i s are positive, Equation (44) EQUATIONwhich completes our proof.In addition to a new lower bound of HyperPrism, we can also prove an upper bound of the HyperPrism at 1 2p m i=1 α i x i + 2p-1 2p U following similar steps. Note that these bounds are generic and can be applied outside of DML. These bounds are necessary for analyzing the benefits of using a HyperPrism aggregation function in DML because they guarantee that the HyperPrism takes at least some minimum consideration of each x i term. This consideration is crucial because each device's local objective function is unique, and they must all be considered to optimize the overall objective function. 


Decision: This paper proposes a novel way of non-linear aggregation in distributed machine learning, with the use of a hyper network. The advantage of this method is the reduction of communication rounds especially in the presence of non-IID data and time-varying communication links, thereby improving the convergence speed. The convergence of the algorithm has been analyzed theoretically and its performance has also been validated in experiments. It would be useful if the theoretical convergence results can be discussed further, such as by comparing the convergence rate with linear aggregation algorithms and explaining whether/how the convergence bound captures the effect of non-IID data and time-varying communication links. In addition, as mentioned by the reviewers, the presentation of this paper could be improved.

Reviews: review: paper_summary: paper_summary: This paper presents HyperPrism, a novel framework for decentralized machine learning (DML) that aims to address the challenges of non-IID data and time-varying communication links. The authors propose a non-linear aggregation method based on Kolmogorov Means and adaptive mapping functions, which they argue improves convergence speed and scalability compared to traditional linear aggregation methods. The paper includes theoretical analysis and experimental results to support their claims.
strength_weakness: strength_weakness: •	**Experimental Setup:** The experimental setup could be expanded to include more diverse datasets, models, and settings. The current experiments are limited to MNIST and CIFAR-10 with specific model architectures and the non-IID setting include only two extreme cases ($\alpha=0.1$ and $\alpha=10$). This would help to assess the generalizability of HyperPrism's performance improvements. 

•	**Clarity:** The paper could benefit from improved clarity in some sections. For example, in section 3 the authors discussed HyperPrism without first introducing it. The motivation for using mirror descent and Kolmogorov Means was unclear. The connection between these concepts and the challenges of non-IID data and time-varying communication could be made more explicit. 

•	**Hyperparameter Tuning:** The paper does not provide sufficient details on how the hyperparameters for the different methods were chosen. This makes it difficult to assess the fairness of the comparison. 

**Others minor issues:**

•	There seems to be a missing index $i$ on $w$ in equation 1. 

•	Heterogeneity in referencing in related work section. Some references seem to be typed manually (or at least they are not connected to any entry in bibliography).
questions: questions: See weaknesses.
limitations: limitations: Overall, the paper addresses an important problem in DML and proposes a novel solution with promising theoretical and empirical results. However, the experimental setup and clarity issues mentioned above make it just fall short of this venue’s bar.
paper_summary: paper_summary: This paper addresses challenges in distributed machine learning (DML) caused by non-IID data and 
unstable communication links. The proposed HyperPrism framework uses mirror descent and adaptive 
mapping to project models into a mirror space for better aggregation and gradient steps. 
It employs adaptive weighted power means (WPM) for efficient model aggregation, significantly 
improving convergence speed and scalability. The framework's analytical results show that HyperPrism 
effectively handles decentralized DML challenges, offering a robust solution for edge device data processing.
strength_weakness: strength_weakness: While the paper is strong overall, there are a few areas that could be improved. 

First, some sections are densely packed with technical details, which might be challenging for readers 
not deeply familiar with the subject. Simplifying these sections or providing additional explanations 
could enhance accessibility. 

In addition, the paper could provide more comparative analysis with existing methods to highlight the 
specific advantages and potential limitations of HyperPrism. However, due to the limited period of rebuttal, 
it is just optional and not necessary. 

Lastly, a more detailed discussion on the computational overhead and scalability of the proposed framework 
in extremely large-scale settings would be beneficial. Addressing these weaknesses would strengthen the paper 
and its contributions.
questions: questions: How does HyperPrism perform compared to other non-linear aggregation frameworks or adaptive learning methods in similar scenarios? 
Are there specific benchmarks or datasets where it excels or underperforms?
limitations: limitations: The potential limitation of the proposed method is not discussed in this manuscript. It will be of great importance if the authors can give an insight to the potential readers by comparing the pro/cons of the proposed method compared to the latest works in this topic.
paper_summary: paper_summary: Traditional DML methods are limited to (1) data heterogeneity and (2) time-varying communication links. In this work, they present a non-linear class aggregation framework HyperPrism that leverages Kolmogorov Means to conduct distributed mirror descent with the averaging occuring within the mirror descent dual space. The proposed method can improve the convergence speed up to 98.63% and scale well to more devices compared with the state-of-the-art, all with little additional computation overhead compared to traditional linear aggregation.
strength_weakness: strength_weakness: 1.	The authors claim that they use hypernetworks to predict p (since they use softmax, type might be int), the exponent of the mapping function. There is a little deviation from the common practice of hypernetworks, where they usually output more complex parameter weights, like the weight of classifiers (a 768 * 100 matrix, type is float). I would prefer the authors remove the part of hypernetworks, and simply introduce HN as a simple MLP.
2.	While HNs act as a crucial part of HyperPrism, the analyses of them seem limited. Since HN has softmax layers, I would be curious about (1) the value set of p and its influence, (2) the variance of p on each distributed machines when given different gradients, (3) the results of preset p (≠ 1) (like directly use the most outputted p of HNs rather than 1). I would suggest these ablation studies be added.
3.	While the performance of HyperPrism is excellent on the major experiments, the baselines used for comparison seems out of date (before 2021), it would be more reasonable if the authors could provide more recent baselines (in 2022 or 2023).
4.	The results of the experiments are a bit confusing, especially in terms of improvement. For example, in the last column of Table 1, Conv Rds of the proposed method is 13, while the best baseline is 14, how is 85.86% computed ? It is computed with the least performed baseline? It would be better if the authors could provide detailed explanations.
questions: questions: See weaknesses
limitations: limitations: N.A

Comments: title: Thank you very much for your recognition.
Thank you very much for your recognition!
comment: Thank you very much for your recognition.
Thank you very much for your valuable suggestions and acknowledgment. It is essential to highlight that our method is specifically designed for decentralized machine learning scenarios with strong data heterogeneity and time-varying communication networks. Excessively complex large datasets often hardly make achieving global convergence extremely challenging, even non-convergence. This is why existing studies in this field predominantly use the MNIST and CIFAER10 datasets for evaluation, e.g., [R1, R2]. We also actively explore implementing our method across a broader spectrum of scenarios.

The responses to the two above minor issues are as follows:

1) We confirm no additional index 'i' is needed in Equation (1). The 'w'  of Equation (1) denotes the aggregated local model from neighbors' devices, represented as $f_i(w)=\mathbb{E}_{\zeta_i\sim D_i} [\mathcal{F}(w_i;\zeta_i, G(t))]$ (Please refer to Line 109, Page 3). The notation in Equation (1) then represents selecting a single model which optimizes the function F(w) = sum of f_i(w).

2) We confirm that our manuscript is generated using the NeurIPS 2024 LaTeX template, and manual reference typing is not applicable in this context. There may be some formatting issues during the PDF conversion process that lead to some references not being linked to the bibliography.

[R1] Vogels T, He L, Koloskova A, et al. Relaysum for decentralized deep learning on heterogeneous data. Advances in Neural Information Processing Systems, 2021, 34: 28004-28015. 

[R2] Le Bars B, Bellet A, Tommasi M, et al. Refined convergence and topology learning for decentralized SGD with heterogeneous data. International Conference on Artificial Intelligence and Statistics. PMLR, 2023: 1672-1702.
Thank you for the detailed response. I acknowledge that I missed some of the additional settings presented in table 1. Therefore I am willing to increase my score to 5. I will not go beyond score of 5 because I still find the diversity of the experiments limited (e.g. datasets, models, lack of recent baselines). 

PS: could you also address the minor issues I mentioned in my review?
The responses seem reasonable, so I'll stick with my original rating.