Ditto: Fair and Robust Federated Learning Through Personalization

In plain words

Federated learning trains one shared model on data that stays on many separate devices. Each device's data can look quite different from the others, so one shared model may serve some devices well and others badly. A few devices that send bad updates can also spoil the shared model.

Ditto keeps the shared model and gives each device its own personal copy as well. The personal copy learns from that device's data, but it is pulled back toward the shared model. A single setting controls the pull. With no pull, each device just trains alone. With a very strong pull, every device effectively uses the shared model. The authors argue that a setting in between makes the results more accurate, more even across devices, and more robust to attacks than methods built only for fairness or only for robustness.

TL;DR

Ditto is a lightweight personalization layer for federated learning. Each device trains a personal model with an L2 pull toward the current global model, and the global model is trained by an ordinary federated solver such as FedAvg. The authors report that this matches or beats recent personalization methods on accuracy, and beats robust and fair baselines on robustness and fairness. A closed-form analysis covers only linear problems.

Problem

Federated learning must satisfy accuracy, fairness and robustness at once. Fairness here means that accuracy is even across devices, measured as the standard deviation of test accuracy across devices §2. Robustness means that the mean test accuracy on benign devices stays high under training-time attacks on data or updates §2. Robust aggregators, gradient clipping and fair reweighting methods each address one of these, and the authors show that they can work against each other. Fair methods that upweight rare or hard devices can overfit to corrupted devices, and robust methods can discard informative updates from unusual devices §1; §4.3; fig. 4; fig. 5. The paper considers attacks A1 (label poisoning), A2 (random Gaussian updates) and A3 (model replacement with scaled updates) §2.

Method

Objective. Ditto adds one personal model v_k per device k. It minimises the device's local loss plus a pull toward the optimal global model w^*:

\min_{v_k}\; F_k(v_k) + \frac{\lambda}{2}\left\|v_k - w^*\right\|^2 \quad \text{s.t.} \quad w^* \in \arg\min_w G\big(F_1(w),\dots,F_K(w)\big) §3.1

Here G is the global objective, and FedAvg uses a weighted average of local losses for it §3.1. When \lambda = 0 the personal models are purely local. As \lambda grows, the objective approaches the global one §3.1.

Solver. The paper optimises the global and personal models jointly and alternately, not one after the other §3.2. In each round:

1. The server samples a subset of devices and sends them the current global model w^t alg. 1.
2. Each selected device updates the global model with the solver for G, starting from w^t. With FedAvg this is r local SGD steps on F_k app. B, alg. 2.
3. Each device updates its personal model for s local steps with v_k \leftarrow v_k - \eta\,(\nabla F_k(v_k) + \lambda(v_k - w^t)) §3.2; alg. 1; app. B. The pull is toward w^t, the global model at the start of the round, not the device's own global update.
4. The device sends only its global update back. Personal models never leave the device, so the communication matches that of the global solver §3.2.
5. The server aggregates the global updates into w^{t+1} alg. 1.

The method for the personal step, in the FedAvg case used for most experiments:

# Ditto round with FedAvg as the global solver  [app. B, alg. 2]
# w: global model w^t; personal[k]: personal model v_k of device k
# grad_F(model, k): mini-batch gradient of F_k  [§3.2]

def ditto_round(w, personal, devices, lam, lr, r, s):
    # Server selects the devices S_t and sends them w^t  [alg. 1]
    updates = []
    for k in selected(devices):
        # Global part: r local SGD steps on F_k from w^t  [app. B, alg. 2]
        w_k = w
        for _ in range(r):
            w_k = w_k - lr * grad_F(w_k, k)
        # Personal part: s steps on F_k + lam/2 * ||v - w^t||^2  [§3.2]
        for _ in range(s):
            pull = lam * (personal[k] - w)
            personal[k] = personal[k] - lr * (grad_F(personal[k], k) + pull)
        # Only the global update is sent back  [§3.2; alg. 1]
        updates.append(w_k - w)
    # Server averages the global updates  [app. B, alg. 2]
    w = w + mean(updates)
    return w, personal

Choosing λ. The paper treats λ as a per-device hyperparameter. Each device selects λ from a small candidate set using its own validation data §4 Setup; app. D.2. The candidate grids depend on whether the attack is treated as strong (defined in app. D.2). Section 3.3 gives closed-form values of λ* for linear regression: λ* = σ²/(nτ²) without adversaries, and an adversarial version that decreases as the number of adversaries and the attack strength grow §3.3; app. A. Here τ measures how related the devices are and σ the noise. A device does not know τ, so the closed form cannot be used directly, and the paper does not give an estimator for it.

Convergence. If the global model converges at rate g(t), the personal models converge at a rate C\,g(t) §3.2; app. B. With FedAvg as the global solver this gives O(1/t) for the personal models app. B. The analysis assumes strong convexity, bounded stochastic gradients and bounded heterogeneity M app. B.

Robust aggregation. Robust aggregators can replace the averaging step on the server, for example clipping or multi-Krum, and this can improve Ditto further §4.4; app. D.3.

Hyperparameters

Results

Author-reported. Setup: federated benchmarks Fashion MNIST and FEMNIST, plus CelebA, Vehicle and StackOverflow in the appendix; attacks A1, A2 and A3 with several adversary ratios; metrics are mean accuracy on benign devices and its standard deviation §4 Setup; app. C.1; app. D.4.

Headline claims (author-reported):

• Robustness: about 6% absolute higher mean benign accuracy than the strongest robust baseline, averaged over datasets and attacks §1; §4.1; app. D.4.
• Fairness: about 10% lower variance across devices than the fair baseline TERM, while accuracy on clean data is about 5% higher §1; §4.2.
• Personalization: similar or better accuracy than other personalized methods, with slightly lower standard deviation §4.4; tab. 2.

Table 1: average (standard deviation) of test accuracy, Fashion MNIST and FEMNIST tab. 1. The Ditto row matches the dynamic-λ row of app. D.2, but the caption does not say how λ was chosen for this table.

* app. D.4 reports .364 (.13) for this cell. See Pitfalls.

Table 2: other personalization methods, FEMNIST and CelebA tab. 2. Ditto values here differ slightly from Table 1 (see Pitfalls).

Table 3: Ditto with robust baselines, FEMNIST tab. 3. Ditto with clipping is best in five of six columns. Plain Ditto is best at A1 80%.

Figures (described; no images are used on this card):

• fig. 1: toy point-estimation problem. The λ given by the theory gives the lowest error and standard deviation among Ditto's solutions.
• fig. 2: as devices become more related, global beats local; Ditto with λ* beats both.
• fig. 3: mean benign accuracy under A1 on Fashion MNIST and FEMNIST. Ditto is the most robust under almost all attacks, especially at high adversary ratios.
• fig. 4: fair methods with larger t lose more accuracy under attack on FEMNIST and CelebA.
• fig. 5: robust baselines are either robust but unfair, or not robust. Ditto sits in the robust and fair corner.
• fig. 6: under A3 at 20% adversaries, joint optimisation beats finetuning from w^*.

Paper vs code

Official repository: github.com/litian96/ditto, MIT licence, commit 26d9b29a62f3872dcdf071be08e897607c283bb4 LICENSE. The paper links it in §4 Setup. Discrepancies, each checked against both sides:

• λ selection differs. The paper selects λ per device on validation data, using the grids in app. D.2. The code selects λ at every local step: it tries each candidate in {0.1, 1, 2}, takes one lookahead step, and keeps the candidate with the lowest validation loss ditto.py. The same grid is used for every attack. The strong-attack grid and the small-validation rule of app. D.2 are not implemented. The released scripts set dynamic_lam=0 and λ=1 run_fashion_clean_ditto_lam1.sh, so reproducing the dynamic-λ numbers needs --dynamic_lam=1.
• Device sampling differs. The paper says devices are selected randomly alg. 1. The code samples without replacement with probability proportional to the number of training samples fedbase.py, and the scripts pass --sampling=2 run_fashion_clean_ditto_lam1.sh.
• Local steps are shared. The paper uses r global steps and s personal steps app. B, alg. 2. The code has one local_iters for both. Each iteration takes a personal gradient step ditto.py and then a global step on the same mini-batch ditto.py.
• Step size is constant. The convergence lemma uses a decaying step size 2/((t+1)(\mu+\lambda)p_k) app. B. The code uses the constant learning_rate for the personal update ditto.py, and the model's TensorFlow training op for the global update cnn.py.
• Vehicle learning rate differs. The text says 0.1 for all datasets other than StackOverflow and Fashion MNIST app. D.4. The Vehicle script uses 0.01 run_vehicle_clean_ditto_lam1.sh.
• Boosting factor is set by the code. Model replacement multiplies the malicious update by clients_per_round ditto.py. With the scripts' value of 10 this is a factor of 10. The paper gives no factor in the main text, only ">10×" as the strong-attack threshold app. D.2.
• Random-update attack. Malicious updates are Gaussian with a standard deviation computed from the honest update ditto.py. This matches §2 (A2).
• Label attacks match the paper. Corrupted labels are flipped for CelebA and Vehicle and drawn uniformly at random for FEMNIST and Fashion MNIST ditto.py. This matches the threat model in §2 (A1).
• Corrupted devices are fixed by seed. The corrupted set comes from np.random.seed(1234567 + seed) ditto.py.

Only the Ditto trainer and the run scripts were read in detail. The baseline trainers were not reviewed.

Pitfalls

• λ grid gap. The rules cover devices with fewer than 4 and more than 5 validation points app. D.2. A device with exactly 4 or 5 is covered by neither rule. The paper does not say what those devices should use.
• Per-step λ selection costs more. Each local iteration runs two mini-batch gradient computations, one for the personal model and one for the global model ditto.py, plus one loss call ditto.py. In dynamic mode it also runs three validation-loss evaluations ditto.py. The paper reports no wall-clock or FLOP cost.
• Stateful devices. Each device must keep its personal model between the rounds in which it is selected §3.2. A device that is not selected for a long time uses a stale personal model. In simulation, the code keeps all personal models on the server in self.local_models ditto.py.
• Closed-form λ is not usable in practice.* The formulas need τ, σ, the adversary count and their spread §3.3. The analysis assumes Gaussian data, X_i^\top X_i = \beta I and a uniform prior app. A. The paper does not give an estimator, and the experiments tune λ on validation data instead.
• Convergence does not transfer to the experiments. The theory assumes strong convexity and a decaying step size app. B. The experiments use non-convex CNNs with a constant step, and the released code does the same.
• Numbers differ between tables. Ditto's clean FEMNIST accuracy is .834 (.09) in Table 1 and .836 (.10) in Table 2; its A1 50% value is .762 (.11) in Table 1 and .767 (.10) in Table 2 tab. 1; tab. 2. Table 1 gives global A3 20% FEMNIST as .314 (.13), while app. D.4 gives .364 (.13). Table 3 shows Ditto A1 80% as .669 and app. D.4 as .672 tab. 3. The λ used in Table 2 is not stated. Use the appendix tables for implementation targets, and expect small differences.
• The "about 10%" variance reduction is ambiguous. The paper does not say whether it is relative or absolute §1; §4.2.
• Vehicle attack ratio is approximate. The script's "50%" run corrupts 11 of 23 devices, which is about 48% run_vehicle_a1_50_ditto_lam1.sh.
• Joint optimisation matters. Finetuning from w^* is worse under attack, and the paper notes that stopping early requires a choice that is hard to make without knowing which devices are benign §3.2; app. D.1; fig. 6.
• Clipping threshold. Clipping uses the median update norm of each round app. D.4. This was not verified in the code, which calls l2_clip at ditto.py.

Lineage

• Builds on: FedAvg as the default global solver §3.1; mean-regularised multi-task learning, which Ditto differs from by regularising toward the global model rather than the average personal model §2; FedProx's L2 term, which resembles Ditto's but yields a single global model §3.1, remark.
• Compared against: fair methods q-FFL and TERM §4.2; robust aggregators (median, Krum, multi-Krum, clipping, k-norm, k-loss) and FedMGDA+ §4.1; personalization methods L2SGD, MOCHA, EWC, SKL, Per-FedAvg, mapper and APFL §4.4; app. C.2.
• Superseded by: not known from content in this wiki.