Distributed Machine Learning with Sparse Heterogeneous Data

1 Introduction . The wide adoption of mobile phones , wearable and smart devices , has created a influx of data which requires processing and storage . Due to the size of these datasets and communication limitations , it is then often not feasible to send all the data to a single computer cluster for storage and processing . This has motivated the adoption of decentralised methods , such as Federated Learning [ 33 , 43 ] , which involves both storing and processing data locally on edge devices . This increase in data sources has led to applications that are increasingly high-dimensional . To be both statistically and computationally efficient in this setting , it is then important to develop approaches that can exploit the structure within the data . A natural assumption in this case is that the data is sparse in some sense . For instance , a subset of features is assumed to be responsible for determining the outcome of interest or , in the case of compressed sensing [ 13 ] , the data is assumed to be generated from a sparse signal . Fitting statistical models on data collected and stored across a variety of devices presents a number of computational and statistical challenges . Specifically , the distributed data sets can be heterogeneous 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . owing to data coming from different population distributions e.g . each device can have different geographic locations , specifications and users . Meanwhile , from a computational perspective , it is often unfeasible , due to network limitations and robustness , to have a single central “ master ” device collect and disseminate information . This creates a demand for statistical methodologies which are both : flexible enough to model potential statistical differences in the underlying data ; and can be fit in a decentralised manner i.e . without the need for a centralised device to collate information . In this work , we set to investigate the statistical efficiency of a method for jointly fitting a collection of sparse models across a collection of heterogeneous datasets . More precisely , models are associated to nodes within a graph , with edges then joining two models if their difference is also assumed to be sparse . The approach of penalising the differences between models ( in an appropriate norm ) has , for instance , been applied within both Federated Learning [ 28 ] and Hyperspectral Denoising [ 23 , 14 ] to encode heterogeneous data . In our case , we consider linear models and obtain precise insights into when jointly fitting models across heterogeneous datasets yield gains in statistical efficiency over other methods , such as , the group Lasso and Dirty Model for multi-task learning [ 24 ] . In addition to being solvable in a centralised manner with standard optimisation packages , the framework we consider is then directly applicable to decentralised applications , as information only needs to be communicated across nodes/models/devices which share an edge within the graph . 1.1 Our Contribution . We consider a total variation scheme that penalises differences between models that share an edge in the graph . This scheme encodes the intuition that if the signal differences are sufficiently sparse then , to recover all signals in the graph , it is more statistically efficient to first recover a single signal associated to a particular reference node ( root ) and then recover the signal differences associated to edges . Following the celebrated Basis Pursuit algorithm [ 5 ] , we consider the solution that minimises the ` 1 norm of the model associated to a root node of the tree and the differences between models that share an edge . We refer to this method as Total Variation Basis Pursuit ( TVBP ) . A noisy variant similar to Basis Pursuit Denoising [ 9 ] is then considered , where the linear constraint is substituted for a bound on the ` 2 norm of the residuals . We call this method as Total Variation Basis Pursuit Denoising ( TVBPD ) . Note that variants of TVBPD have been successfully applied within the context of hyperspectral data [ 23 , 14 ] and thus , this work now provides theoretical justification in this case . Given this framework and assuming sub-Gaussian design matrices , we show that statistical savings can be achieved by TVBP and TVBPD as opposed to solving them either independently or with methods that consider the union of supports ( see Table 1 ) . In the noiseless case , TVBP requires a total sample complexity of O ( s+ n2s′ ) where s is the root sparsity , s′ is the maximum sparsity of the difference along edges and n is the number of nodes . In contrast , methods like the group Lasso and Dirty Model [ 24 ] have an O ( ns ) total sample complexity , which matches the case when nodes do not communicate with each other . Moreover , note that the TVBP approach does not need to know the true underlying graph G , whereas the optimal stepwise approach does . If the true graph G is known , TVBP can incorporate this information to yield improved sampled complexity ( see Table 2 in Section 2.2 ) . In the noisy setting , we demonstrate that TVBPD has an ` 1 estimation error bounded as O ( √ s + √ ns′ ) , where as the stepwise approach scales as O ( s + n √ s′ ) , and thus , achieves an O ( √ ns′ ) saving in this case . While the step wise approach achieves optimal total sample complexity in the noiseless case , in the noisy setting its estimation scales sub-optimally compared to TVBPD . Alongside total sample complexity savings over other joint recovery methods , we also show that TVBP is amenable to decentralised machine learning . Specifically , the objective can be reformulated with constraints that reflect the graph topology allowing a Distributed Alternating Direction Methods of Multipliers ( ADMM ) algorithm [ 3 ] to be applied . Theoretical findings are also supported by experiments ( Section 3.2 ) which demonstrate both : TVBP can outperform group Lasso methods [ 24 ] when model differences are sparse ; and TVBPD yields qualitative improvements in hyperspectral unmixing with the real-world AVIRIS Cuprite mine data set . A key reason for TVBP achieving sample complexity savings over variants of the group Lasso and Dirty Model [ 24 ] , is that the matrices associated to each tasks do not need to satisfy an incoherence condition with respect to their support sets . The incoherence condition is a common assumption within the analysis of sparse methods like the Lasso ( see for instance [ 49 ] ) and , in short , it requires the empirical covariance of the design matrix to be invertiable when restricted to co-ordinates in the signal support set . This immediately implies that the sample size at each task is larger than the support set of that given task . In contrast , for TVBP an alternative analysis is conducted by reformulating the problem into a standard basis pursuit objective with an augmented matrix and support set . In this case , the problem structure can be leveraged to show a Restricted Null Space Property holds when the sample size at each task scales with the sparsity along the graph edges . This highlights that the Total Variation penalty encodes a different structure when compared to methods like the group Lasso and Dirty Model [ 24 ] . As shown in the noisy case , this approach can then be generalised through the Robust Null Space Property [ 8 , 17 ] . 1.2 Related Literature . Learning from a collection of heterogeneous datasets can be framed as an instance of multi-task learning [ 6 ] , with applications in distributed contexts gaining increased attention recently . We highlight the most relevant to our setting . The works [ 52 , 48 ] have considered models penalised in an ` 2 sense according to the network topology to encode prior information . The ` 2 penalty is not appropriate for the sparse setting of our work . A number of distributed algorithms have been developed for the sparse setting , for a full review we refer to [ 1 ] . The works [ 24 , 44 , 50 , 29 , 51 , 38 ] have developed distributed algorithms following the group Lasso setting , in that , the signals are assumed to be composed of a common shared component plus an individual component . Within [ 24 , 44 , 50 ] this requires each node to satisfy an incoherence condition , while the setting in [ 29 , 38 ] is a specific case of a star topology within our work . For details on how the incoherence condition influences the sample complexity see discussion in Section 2.1 . The work [ 40 ] develops a manifold lifting algorithm to jointly recover signals in the absence of an incoherence assumption , although no theoretical guarantees are given . Federated machine learning [ 27 , 43 , 33 , 28 ] is a particular instance of distributed multi-task learning where a central node ( root ) holds a global model and other devices collect data and update their model with the root . Data heterogeneity can negatively impact the performance of methods which assume homogeneous data , see [ 33 ] for the case of Federated Averaging . This motivates modelling the data heterogeneity , as recently done within [ 28 ] where the difference in model parameters is penalised at each step . Our work follows this approach in the case of learning sparse models , and thus , it provides insights into when an improvement in statistical performance can be achieved over other baseline methods . Simultaneously recovering a collection of sparse vectors can also be framed into the multiple measurement vectors framework [ 15 ] , which has been precisely investigated for ` 1/ ` q regularisation for q > 1 . Specifically , ` 1/ ` ∞ was investigated within [ 55 , 36 , 47 ] and ` 1/ ` 2 in [ 31 , 37 ] . Other variants include the Dirty Model of [ 24 ] , multi-level Lasso of [ 32 ] and tree-guided graph Lasso of [ 25 ] . In the same context , a number of works have investigated variants of greedy pursuit style algorithms [ 16 , 11 , 12 , 46 ] . These methods assume a large overlap between the signals , with their analysis often assuming each task satisfies an incoherence condition [ 36 , 24 , 37 ] ( see Section 2.1 ) . The total variation penalty is linked with the fused Lasso [ 53 , 21 , 45 , 7 , 42 ] and has been widely applied to images due to it promoting piece-wise continuous signals which avoids blurring . As far as we are aware , the only work theoretically investigating the total variation penalty as a tool to link a collection of sparse linear recovery problems has been [ 10 ] . This work considers the penalised noisy setting and gives both asymptotic statistical guarantees and an optimisation algorithm targeting a smoothed objective . In contrast , we give finite sample guarantees as well as settings where statistical savings are achieved . The application of hyperspectral unmixing [ 22 , 23 , 14 ] has successfully integrated the total variation penalty within their analysis . Here , each pixel in an image can be associated to its own sparse recovery problem , for instance , the presence of minerals [ 23 ] or the ground class e.g . trees , meadows etc . [ 14 ] .

Summary: This paper studies distributed (linear model) learning on heterogeneous data, e.g. federated learning, with the assumption of sparsity. Given a siilarity graph, it minimizes the l1 norm of the sparse model at a root node as well as the total variation along the graph paths. This work provides statistical analyses on the sample complexity by reformulating the problem as a basis pursuit denoising problem, and show that sample complexity is improved compared to previous works. A distributed ADMM algorithm is derived, and synthetic data experiments and application to hyperspectral unmixing are demonstrated.