PaperlayerView exact revision on arXiv

Source-linked AI summary

Federated Multi-Task Learning

Virginia Smith, Chao-Kai Chiang, Maziar Sanjabi, Ameet Talwalkar

arXiv:1705.10467v2cs.LGstat.ML

TL;DR

Federated learning must cope with heterogeneous, non-IID data and systems constraints across distributed devices. This paper frames the problem as federated multi-task learning and introduces MOCHA, a systems-aware optimizer with convergence analysis. MOCHA addresses communication cost, stragglers, and fault tolerance, while experiments demonstrate empirical performance on federated datasets.

  • Problem

    Federated learning seeks models over data generated on distributed nodes with distinct, non-IID distributions and uneven dataset sizes.

  • Method

    MOCHA is a systems-aware optimization framework that learns separate but related node models through federated multi-task learning.

  • Results

    MOCHA’s theory addresses high communication cost, stragglers, and fault tolerance, and its empirical evaluation uses federated datasets.

  • Takeaways & Limitations

    Federated multi-task learning provides a modeling choice for statistical heterogeneity while MOCHA extends distributed optimization to practical federated systems constraints.

  • Takeaways & Limitations

    MOCHA does not currently apply to non-convex deep learning models, and asynchronous-method connections remain future work.

Abstract

from arXiv · show

Federated learning poses new statistical and systems challenges in training machine learning models over distributed networks of devices. In this work, we show that multi-task learning is naturally suited to handle the statistical challenges of this setting, and propose a novel systems-aware optimization method, MOCHA, that is robust to practical systems issues. Our method and theory for the first time consider issues of high communication cost, stragglers, and fault tolerance for distributed multi-task learning. The resulting method achieves significant speedups compared to alternatives in the federated setting, as we demonstrate through simulations on real-world federated datasets.

1 Introduction

Federated learning must address both non-IID, unevenly distributed data and practical systems constraints such as communication bottlenecks and heterogeneous devices. The paper uses multi-task learning and introduces MOCHA to handle these challenges with convergence guarantees and empirical evaluation.

  • Federated nodes generate non-IID data from distinct distributions, with substantially varying numbers of data points.
  • Communication is a major bottleneck in networks containing many nodes with heterogeneous storage, computation, and connection capacities.
  • Multi-task learning fits separate but related models for nodes, addressing statistical heterogeneity more directly than a single global model.
  • MOCHA generalizes COCOA for federated multi-task learning and addresses network size, node heterogeneity, stragglers, and fault tolerance.
  • The paper provides convergence guarantees that account for federated systems challenges and evaluates MOCHA on federated datasets.

2 Related Work

Prior federated learning methods commonly learn one global model, limiting their handling of non-IID data and node structure. Existing distributed multi-task methods also leave important systems concerns, convergence guarantees, or practical federated assumptions unresolved.

  • Learning Beyond the Data Center: Existing federated learning approaches aim to learn a single global model, limiting their ability to handle non-IID data and structure among nodes.
  • Multi-Task Learning: Multi-task learning can capture relationships among non-IID and unbalanced data, but existing distributed methods do not adequately address federated systems challenges.
  • Distributed Multi-Task Learning: Several distributed multi-task methods lack flexibility between communication and computation, limiting their handling of fault tolerance and stragglers.
  • Distributed Multi-Task Learning: Asynchronous methods mitigate stragglers but may lack fault tolerance, convergence guarantees, or rely on impractical bounded-delay assumptions.
  • Learning Beyond the Data Center: The paper uses “federated learning” exclusively because “on-device learning” ambiguously refers to both model training and model serving.
  • Learning Beyond the Data Center: Privacy is important in federated learning, and the paper notes that privacy benefits associated with global federated learning also apply to its approach.

3 Federated Multi-Task Learning

The paper formulates federated learning as multi-task learning, fitting related local models while addressing distributed optimization challenges such as communication, stragglers, and dropped nodes. MOCHA extends primal-dual optimization with data-local updates, centralized relationship updates, and approximate synchronous computation.

  • 3 Federated Multi-Task Learning: Federated multi-task learning fits separate but related models for nodes whose data distributions differ across the network.The framework represents local models as columns of W and models task relationships through Ω and the regularizer R.
  • 3.1 General Multi-Task Learning Setup: The general MTL formulation uses arbitrary convex local losses and a regularizer that promotes structure among task models.Ω may be known beforehand or estimated jointly, while R takes Ω as input.
  • 3.2 MOCHA: A Framework for Federated Multi-Task Learning: MOCHA alternates between distributed updates of W and centralized updates of Ω.With Ω fixed, W depends on distributed data and centralized task structure; with W fixed, Ω depends only on W.
  • 3.3 Federated Update of W: MOCHA decomposes the federated W update into data-local quadratic subproblems that use only each node’s local data.The method extends distributed primal-dual optimization and communicates the block v_t = X_tα between nodes at each iteration.
  • 3.4 Practical Considerations: MOCHA mitigates stragglers by allowing each node to approximately solve its subproblem with its own per-iteration quality parameter θ_t^h.A node with θ_t^h = 1 makes no progress and may be treated as dropped; the method remains synchronously updated and tolerates a small fraction of periodically dropping nodes.
  • 3.4 Practical Considerations: The paper leaves asynchronous updating schemes for future work because bounded-delay assumptions in many such methods limit fault tolerance.This marks a scope boundary between approximation-based synchronous updates and asynchronous alternatives.

4 Convergence Analysis

MOCHA's federated W-update uses flexible per-node, per-iteration approximations to accommodate stragglers and dropped nodes, with convergence guarantees under assumptions on returned updates. The analysis covers smooth and Lipschitz losses and extends to asymptotic convergence under milder straggler assumptions.

  • Federated Update: MOCHA allows each node's approximation quality to vary between 0 and 1 at every iteration, supporting heterogeneous computation and unreliable participation.The parameter θ^h_t captures the quality of node t's local result at iteration h.
  • Assumptions: Assumption 2 requires nonzero probability of receiving a node result and an expected improvement over the previous iterate.The assumption permits periodic drops but excludes nodes that never return updates.
  • Convergence Guarantees: Under smooth losses and Assumptions 1–2, Theorem 1 establishes finite-horizon convergence for any target ε_D with a suitable number of iterations.The convergence characterization is stated for the federated MOCHA update.
  • Convergence Guarantees: The analysis also permits asymptotic convergence as H → ∞ under milder assumptions on stragglers.This result is provided through an appendix corollary.
  • Convergence Guarantees: For L-Lipschitz losses, Theorem 2 provides a sub-linear convergence guarantee.The theorem applies to non-smooth losses such as the hinge loss used by SVM models.
  • Relation to Prior Work: The theorems guarantee convergence in the federated setting under mild assumptions and generalize prior COCOA results to multi-task learning.The guarantees are stated for the dual and analogously for the duality gap.

5 Simulations

The simulations evaluate federated multi-task learning on real-world datasets and compare multi-task, local, global, and distributed optimization methods. Multi-task learning achieves the lowest average error across datasets, while MOCHA remains robust to communication cost, statistical and systems heterogeneity, and periodic node drops.

  • Datasets: The benchmark uses real-world federated datasets including Google Glass, Human Activity Recognition, and Vehicle Sensor data.Participants or sensors are modeled as separate tasks in each dataset.
  • Multi-Task Learning for the Federated Setting: The evaluation compares multi-task models with fully local and fully global models using repeated train-test splits and cross-validation.Models are trained on 75% of the data, tested on 25%, and averaged over 10 trials.
  • Multi-Task Learning for the Federated Setting: Multi-task learning significantly outperforms local and global models on every dataset in average error across tasks.The global model performs worst for most datasets, including under highly skewed data, while MTL continues to outperform both alternatives.
  • Straggler Avoidance: MOCHA and COCOA remain robust to high communication cost, whereas mini-batch methods degrade as communication increases.MOCHA also remains robust to statistical heterogeneity, while COCOA is affected by stragglers caused by difficult subproblems.
  • Straggler Avoidance: MOCHA handles systems heterogeneity by varying local work, while the experiments report that additional systems challenges would further reduce COCOA's performance.The simulated variability ranges from 10%–100% of local data points in high-variability environments to 90%–100% in low-variability environments.
  • Tolerance to Dropped Nodes: MOCHA remains robust when nodes periodically drop, both during a single W update and across the full method.The method does not converge to the correct solution if a node never sends updates, validating Assumption 2.

6 Discussion

The paper presents MOCHA as a systems-aware optimization framework for federated multi-task learning that addresses communication, stragglers, and fault tolerance. Its current scope excludes non-convex deep learning models, though the authors note possible connections to convexified approaches.

  • Discussion: MOCHA is a systems-aware optimization framework for federated multi-task learning addressing high communication cost, stragglers, and fault tolerance.The paper describes these as issues considered by both the method and its theory.
  • Scope: MOCHA does not currently apply to non-convex deep learning models.The authors note possible connections with convexified deep learning models in kernelized federated multi-task learning.

B Multi-Task Learning

The paper's multi-task learning framework covers several ways to model relationships among tasks, including clustered, graphical, and probabilistic structures. MOCHA can solve these formulations and can be adapted when tasks are shared across nodes.

  • MTL Formulations: The framework addresses multiple MTL formulations by separating the W update from the Ω update.The W update is handled in Section 3, while this section describes Ω updates for representative formulations.
  • Clustered MTL: Clustered MTL assumes task weights are close within clusters, with formulations for known or learned cluster structure.Convex relaxations can be used when a perfect clustering structure would make the problem non-convex.
  • Probabilistic Priors: Probabilistic-prior MTL models dependence among task weights through covariance structure and can represent both positive and negative task relationships.The formulation is made jointly convex by omitting log |Ω| and constraining tr(Ω).
  • Graphical Models: Graphical-model MTL represents conditional independence through zeros in a task precision matrix and updates Ω using graphical lasso methods.The corresponding optimization is a sparse precision-matrix estimation problem.
  • Shared Tasks: MOCHA can accommodate shared tasks across nodes without changing the local solvers.The central node aggregates results for shared tasks, reducing Ω's size and simplifying its update.

C Convergence Analysis

This section develops convergence analysis for MOCHA under aggregation, smoothness, and Lipschitz assumptions. The lemmas progressively relate local subproblem improvements and dual objectives to global convergence.

  • Convergence setup: MOCHA uses an aggregation parameter γ ∈(0, 1] to scale federated updates, although Algorithm 1 presents the γ = 1 case.The updates are α_t ← α_t + γ∆α_t and v_t ← v_t + γ∆v_t.
  • Convergence setup: The analysis defines per-task quantities and uses Assumption 2 to bound average task performance, which is needed for global convergence guarantees.The corresponding lemma bounds the average performance of θ^h_t.
  • Dual progress: Lemma 4 bounds the dual objective of an iterate using the previous dual objective and local subproblem objectives.Its proof uses local subproblem definitions, smoothness of R* and the choice of σ′.
  • Dual progress: Lemma 5 bounds per-iteration improvement in the dual objective and remains valid for non-smooth functions when µ = 0.For smooth losses, Lemma 6 further relates one-step dual improvement to the duality gap.
  • Non-smooth losses: For L-Lipschitz losses, additional bounds on the J term are required because Lemma 5 alone does not provide sufficient decrease when µ = 0.Lemma 7 supplies the needed bound under the L-Lipschitz assumption.

C.1.1 Proof of Theorem 1

The proof derives convergence by recursively applying sufficient-decrease inequalities and bounding the duality gap. It also extends the analysis to smooth losses and discusses how aggregation affects update aggressiveness.

  • Proof of Theorem 1: Recursive application of the sufficient-decrease inequality and expectation bounds yields a bound on the dual distance to optimality.The initial duality gap is bounded to obtain the final result.
  • Proof of Theorem 1: sγ(1 −¯Θ) E[G(α(H))] ≤ E[D(α(H)) − D(α(H+1))] ≤ E[D(α(H)) − D(α⋆)] ≤ ϵD.The inequalities use the sufficient-decrease relation, optimality of α⋆, and the bound on dual distance.
  • Smooth-loss extension: For µ-smooth losses, Corollary 8 establishes E[D(α(H))−D(α⋆)] → 0 as H →∞ under either of its stated conditions.This is presented as an asymptotic extension of Theorem 1 under Assumption 1.
  • Lipschitz-loss extension: For L-Lipschitz losses, the proof replaces the smooth-loss decrease argument with Lemma 5, Lemma 7, recursive inequalities, and a careful choice of s.The analysis cannot bound dual-objective decrease using equation (23) in this case.
  • Aggregation trade-off: When γ < 1, σ′ is smaller and local subproblems are less restrictive, producing more aggressive ∆α updates that require more conservative aggregation.Despite this trade-off, γ = 1 has the best empirical performance in most practical scenarios.

E.1 Datasets

The empirical study uses publicly available federated datasets characterized by task counts, feature sizes, per-task data sizes, and data skew.

  • Dataset characteristics: Table 2 reports the number of tasks m, feature size d, and per-task data size n_t for each federated dataset.The standard deviation n_σ measures skew in training-data sizes across tasks.
  • Dataset characteristics: The datasets used in the empirical study are publicly available.

E.2 Multi-Task Learning with Highly Skewed Data

The study evaluates multi-task, local, and global models on datasets whose task sizes differ by at least two orders of magnitude. Multi-task learning still significantly outperforms the alternatives despite this extreme skew.

  • Skewed-data construction: Highly skewed datasets are created by sampling original training data so task dataset sizes differ by at least two orders of magnitude.The resulting dataset sizes are reported in Table 3.
  • Results: The global model performs slightly better in the highly skewed setting, particularly for Human Activity.This comparison is reported for the results in Table 4.
  • Results: Multi-task learning still significantly outperforms all models on the highly skewed datasets.Table 4 reports average prediction errors as means and standard errors over 10 random shuffles.

E.3 Implementation Details

The implementation evaluates federated optimization methods under statistical and systems heterogeneity, communication costs, and node dropouts. It reports that MOCHA remains robust and outperforms alternatives across variability settings, while complete node loss violates its convergence condition.

  • Estimated Time: The estimated federated runtime combines local floating-point operations with communication size and frequency using mobile-network clock-rate and bandwidth/latency measurements.Communication cost includes both bandwidth and latency.
  • Statistical Heterogeneity: MOCHA and mini-batch methods adapt local computation or batch size to statistical heterogeneity, unlike COCOA's fixed approximation parameter.COCOA can become slower across nodes when subproblem difficulty and data skew vary.
  • Systems Heterogeneity: MOCHA uses a global clock cycle, with each worker determining local computation according to system and statistical conditions.The simulations model constrained nodes by assigning variable numbers of local updates.
  • Systems Heterogeneity: MOCHA significantly outperforms all other methods in both low- and high-variability settings and remains robust to systems-related heterogeneity.COCOA receives no additional simulated variability, producing overly optimistic results for that baseline.
  • Fault Tolerance: MOCHA tolerates periodic node dropouts, but it does not converge to the correct solution when a node drops out permanently.The convergence assumption requires each node's per-round dropout probability to remain below one.
Loading 1705.10467v2…