Mean-Regularized Multi-Task Learning

• The paper introduces a framework that regularizes individual task models toward the average model, balancing personalization and global trends.
• It presents a federated optimization algorithm using clipping and Gaussian noise to enforce joint differential privacy while retaining model performance.
• Empirical evaluations on benchmarks like FEMNIST and CelebA show that this approach outperforms baseline methods in stringent privacy settings.

Mean-regularized multi-task learning (MTL) is a framework in which multiple related machine learning tasks are solved simultaneously through an explicit regularization strategy that encourages each task’s model to remain close to the average of all task models. Originally captured by empirical risk minimization with mean-based regularization penalties, this approach is central to privacy-sensitive collaborative settings such as federated learning, where data resides with separate clients and strong privacy guarantees are required for individual contributions. Recent advances formally incorporate joint differential privacy (JDP) into mean-regularized MTL, yielding algorithms that achieve both privacy protection and competitive empirical performance (Hu et al., 2021).

1. Mean-Regularized MTL Objective

In mean-regularized MTL, each of $m$ tasks (or clients) $k$ has data $D_k = \{x_{k,i}, y_{k,i}\}_i$ and an associated parameter vector $w_k \in \mathbb{R}^d$. The per-task empirical loss is defined as

$$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$

Aggregating all task parameters as $W = [w_1; \ldots; w_m]$, the average model is

$\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$

The mean-regularized MTL objective takes the form

$$F(W) = \frac{1}{m} \sum_{k=1}^m \left[ \mathcal{L}_k(w_k) + \frac{\lambda}{2} \|w_k - \bar w\|_2^2 \right],$$

where $\lambda > 0$ is a regularization parameter. This formulation penalizes divergence from the average model, capturing a spectrum between purely personalized and fully global modeling.

An equivalent iterative interpretation is to introduce a global anchor $\tilde w$ and solve, for each round and task $k$0,

$k$1

using $k$2 steps of stochastic gradient descent (SGD), followed by an update of $k$3 as the mean of all $k$4.

2. Federated Optimization under Joint Differential Privacy

The mean-regularized MTL framework is amenable to communication-efficient learning in federated settings and can be made private via adaptations of the Federated Averaging (FedAvg) protocol. The server’s public state at iteration $k$5 is denoted $k$6. At each round:

• The server uniformly samples a subset $k$7 of $k$8 clients.
• $k$9 is broadcast to $D_k = \{x_{k,i}, y_{k,i}\}_i$0.
• Each $D_k = \{x_{k,i}, y_{k,i}\}_i$1 performs $D_k = \{x_{k,i}, y_{k,i}\}_i$2 steps of local SGD with respect to $D_k = \{x_{k,i}, y_{k,i}\}_i$3:

$D_k = \{x_{k,i}, y_{k,i}\}_i$4

• The update $D_k = \{x_{k,i}, y_{k,i}\}_i$5 is clipped as $D_k = \{x_{k,i}, y_{k,i}\}_i$6, enforcing $D_k = \{x_{k,i}, y_{k,i}\}_i$7.
• The server aggregates and perturbs the mean update:

$D_k = \{x_{k,i}, y_{k,i}\}_i$8

After $D_k = \{x_{k,i}, y_{k,i}\}_i$9 rounds, each client retains its own $w_k \in \mathbb{R}^d$0 as a private, personalized model. The privacy parameters are calibrated as follows: sampling rate $w_k \in \mathbb{R}^d$1; clipping norm $w_k \in \mathbb{R}^d$2; and noise $w_k \in \mathbb{R}^d$3 for a global $w_k \in \mathbb{R}^d$4-differential privacy guarantee.

3. Formal Privacy Guarantees: Differential and Joint Differential Privacy

Joint differential privacy is a relaxation suited for distributed settings and mechanism design. For mechanism $w_k \in \mathbb{R}^d$5:

$w_k \in \mathbb{R}^d$6

for every client $w_k \in \mathbb{R}^d$7, every pair of datasets $w_k \in \mathbb{R}^d$8 differing only in $w_k \in \mathbb{R}^d$9's data, and every event $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$0 over all clients except $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$1. This formalism ensures that no coalition of $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$2 clients gains significant information about any excluded individual.

The algorithm described achieves $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$3-DP for the broadcast sequence $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$4, provided $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$5 is set as above and $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$6 (for constant $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$7). By post-processing and the specific design, the collection of personalized outputs $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$8 satisfies $$\mathcal{L}_k(w_k) = \sum_{i=1}^{n_k} \ell_k(x_{k,i}; w_k).$$9-joint differential privacy (Hu et al., 2021).

4. Convergence and Utility Analysis

For non-convex objectives, assuming each $W = [w_1; \ldots; w_m]$0 is $W = [w_1; \ldots; w_m]$1-smooth and gradients are bounded by sufficient clipping (with large enough $W = [w_1; \ldots; w_m]$2), the averaged squared gradient norm over $W = [w_1; \ldots; w_m]$3 rounds is

$W = [w_1; \ldots; w_m]$4

where $W = [w_1; \ldots; w_m]$5 bounds the instantaneous loss and $W = [w_1; \ldots; w_m]$6 is the number of local steps. With $W = [w_1; \ldots; w_m]$7 and $W = [w_1; \ldots; w_m]$8, the asymptotic error terms are

$W = [w_1; \ldots; w_m]$9

indicating that the error due to noise injection scales as $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$0. Thus, increasing the number of clients amortizes privacy-induced error ("amplification by multi-task" effect).

For $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$1-strongly convex $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$2, exponential convergence is obtained to a joint solution neighborhood of radius $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$3:

$\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$4

5. Empirical Evaluation and Comparative Performance

Empirical results span benchmarks such as FEMNIST (205 clients; 4-layer CNN; 62-way classification), StackOverflow (400 clients; convex logistic regression; 500-way tag prediction), and CelebA (515 clients; 4-layer CNN; binary classification). The following baselines are evaluated:

• Private FedAvg and FedProx (global models, Gaussian mechanism noise)
• Non-private MTL
• Pure local training
• Private MTL with local finetuning (mean-regularization, KL-divergence, EWC)
• PP-SGD (personalized DP via billboard mechanism, convex tasks)

Metrics include test accuracy versus privacy budget $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$5, training loss, evolving $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$6 over rounds, and privacy-utility trade-off curves.

Key findings are:

• Private MTL (PMTL) outperforms private FedAvg and FedProx for all $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$7 and datasets, especially at small $\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$8 ($\bar w = \frac{1}{m} \sum_{k=1}^m w_k.$9).
• MTL’s utility gains are largely retained even with Gaussian noise: personalization continues to confer benefits.
• Local finetuning after MTL further improves accuracy (notably in non-convex CNN tasks) at negligible additional privacy cost.
• PP-SGD underperforms PMTL on convex tasks for stringent privacy ($$F(W) = \frac{1}{m} \sum_{k=1}^m \left[ \mathcal{L}_k(w_k) + \frac{\lambda}{2} \|w_k - \bar w\|_2^2 \right],$$0) but converges in performance under weaker privacy.
• Empirically, as the number of tasks $$F(W) = \frac{1}{m} \sum_{k=1}^m \left[ \mathcal{L}_k(w_k) + \frac{\lambda}{2} \|w_k - \bar w\|_2^2 \right],$$1 increases, the required noise standard deviation $$F(W) = \frac{1}{m} \sum_{k=1}^m \left[ \mathcal{L}_k(w_k) + \frac{\lambda}{2} \|w_k - \bar w\|_2^2 \right],$$2 decays as $$F(W) = \frac{1}{m} \sum_{k=1}^m \left[ \mathcal{L}_k(w_k) + \frac{\lambda}{2} \|w_k - \bar w\|_2^2 \right],$$3, matching theoretical predictions.

6. Significance and Applications in Privacy-Sensitive Collaborative Learning

Mean-regularized MTL, with differentially private algorithms, has broad applicability in privacy-sensitive domains such as healthcare, finance, and IoT. The established approach—using a federated-averaging–style loop, clipping, and Gaussian noise addition—supports joint differential privacy while enabling personalized model training per client. By combining personalization and formal privacy, mean-regularized MTL supersedes both global and local-only methods along privacy-utility frontiers and scales favorably with increasing clients due to amortized privacy costs (Hu et al., 2021).