PaperlayerView exact revision on arXiv

Source-linked AI summary

Local Epochs as Optimal Stopping: Theory and Evidence in Federated Learning

Ke Xiao, Qiyuan Wang, Christos Anagnostopoulos

arXiv:2610.21832v1cs.LGcs.DC

TL;DR

Fixed, hand-tuned local epoch counts leave open when clients should stop training as local gains diminish and drift accumulates. The paper formulates epoch selection as drift-priced optimal stopping and develops FedStop. FedStop limits mean-accuracy loss against tuned FedAvg while improving worst-decile accuracy under strong label skew.

  • Problem

    Fixed local epoch counts are tuned constants that can be fragile and do not explain how budgets should change during training or across clients.

    Source: §1 · p. 1

  • Method

    The paper models local epochs as finite-horizon optimal stopping, pricing squared drift against loss reduction, and derives a one-stage threshold rule for FedStop.

    Source: §1 · pp. 1–2

  • Results

    FedStop costs at most one point of mean accuracy against tuned FedAvg and raises worst-decile accuracy by up to eight points under strong label skew.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Training stage, rather than client identity, determines how much local work pays, and the stopping decision takes constant time per epoch with a few hundred scalars of memory per client.

    Source: §1 · p. 2, §10 · p. 18

    The optimality guarantee requires non-reversing stopping decisions, while local progress is more autocorrelated than the analyzed margin allows.

    Source: §9 · p. 17

Abstract

from arXiv · show

How long should each client train before communicating? Federated learning usually answers with a fixed, hand-tuned number of local epochs, yet an extra epoch helps only while the loss it removes outweighs the client drift it adds. We cast local training as finite-horizon optimal stopping: a client trains while the next epoch's expected loss reduction exceeds its priced squared drift. We prove that this one-stage look-ahead threshold is optimal whenever stopping decisions are nested, derive closed forms for geometric and power-law returns, and show that an estimated rule loses utility only through wrong-signed decisions. Epoch-level evidence supports the model's structure (geometric diminishing returns and additive drift) and reveals a simple regularity: the stage of training, not the client, sets the budget. The resulting rule, FedStop, is never more than one point below tuned FedAvg in mean accuracy. Where its budgets vary, it raises worst-decile accuracy by 4-8 points over tuned FedAvg and an adaptive-interval baseline, and every one of these gains remains significant after a Holm correction. Against the client-side stopping rules ALT and FLASH, it has the highest CIFAR-10 worst decile on every seed and is level elsewhere. It stays stable where per-client early stopping collapses, and each stopping decision takes constant time.

1 Introduction

Fixed local-epoch budgets are fragile and do not explain how training duration should change over time or across clients. The paper models epoch selection as drift-priced optimal stopping, tests its predictions against training logs, and introduces FedStop, which improves worst-decile accuracy while limiting mean-accuracy loss.

Problem

  • A tuned fixed epoch count is fragile: the best configuration is two epochs on CIFAR-10 but one on CIFAR-100, where two epochs diverge.A fixed budget also does not explain why it works or whether it should change during training or across clients.Source: §1 · p. 1

Model

  • The model treats each epoch as a trade-off between local-loss reduction and added drift, framing the stopping decision as a finite-horizon problem.FedStop stops when another epoch is no longer worth its drift.Source: §1 · pp. 1–2

Theory

  • The one-stage threshold is optimal when stopping decisions are nested, and an estimated rule loses utility only through wrong-signed decisions.The paper also derives closed forms for geometric and power-law returns, including declining drift.Source: §1 · p. 2

Evidence

  • Epoch-level logs support model predictions including diminishing returns and additive drift, and show that training stage, rather than client identity, sets the budget.The paper reports where the model holds and where it does not.Source: §1 · p. 2

Practice

  • FedStop costs at most one point of mean accuracy against tuned FedAvg and raises worst-decile accuracy by up to eight points under strong label skew.It has the highest CIFAR-10 worst decile among ALT and FLASH on every seed, is level with them on naturally partitioned data, and remains stable where per-client early stopping fails.Source: §1 · p. 2

2 Related Work

Prior work studies local-step effects, shared aggregation intervals, and client-side stopping, but the paper distinguishes its approach by modeling each client’s epoch budget as a drift-priced stopping decision.

Client drift

  • FedProx, SCAFFOLD, and FedNova address drift by changing local objectives or aggregation, making them complementary to training-duration selection; FedStop uses FedNova-style normalization.Source: §2 · p. 2

Local computation

  • Prior local-SGD analyses bound or tune one budget per round, while this paper models a client-side stopping decision and tests the model’s assumptions directly.Source: §2 · p. 2

Client-side and global stopping

  • ALT and FLASH use representation divergence and validation-loss signals, respectively, whereas FedStop derives a threshold from an explicit priced-drift model; unlike fairness-oriented methods, it reports but does not optimize worst-client accuracy (Shenaj et al., 2025; Panchal et al., 2023; Li et al., 2020b; Mohri et al., 2019).Prior server-side stopping work decides when to close a round or terminate training; FedStop sets each client’s duration within a round.Source: §2 · p. 2

3 Preliminaries

The preliminaries describe how clients train and aggregate local models in federated learning, then introduce finite-horizon stopping and the condition under which a one-stage look-ahead rule is optimal.

3.1 Federated learning

  • In FedAvg, sampled clients train locally from the broadcast model and the server averages their returned models; when clients take different numbers of steps, step normalization corrects the implicit upweighting of longer-trained clients.Client drift is the distance from the broadcast model to the client’s local model.Source: §3.1 · p. 3

3.2 Finite-horizon optimal stopping

  • Finite-horizon optimal stopping chooses a stopping time using only information available at each decision point, with the objective of maximizing expected reward.Source: §3.2 · p. 3
  • Backward induction compares stopping now with continuing optimally, while the one-stage look-ahead rule compares stopping now with stopping one step later.When the one-stage stopping sets are nested, the one-stage rule is optimal.Source: §3.2 · p. 3

4 Problem Formulation

The formulation prices accumulated squared client drift against local-loss reduction to define a client’s stopping utility. Its aggregation-bound motivation also clarifies which constants must be selected empirically and which timing costs are outside the model.

4.1 Local progress and drift

  • Squared drift accumulates as the sum of per-epoch increments, and a client cannot preview an epoch before paying its cost.The setup assumes synchronous rounds and excludes reporting delay from the client’s stopping utility.Source: §4.1 · p. 4
  • A client’s budget is chosen from its drift-priced utility and does not account for the time it takes to report back.A longer budget may delay the round, but deadline-aware or asynchronous stopping is outside the scope.Source: §4.1 · p. 4

4.2 Drift-priced utility and the stopping problem

  • The net utility increment subtracts a priced marginal squared-drift increment from the marginal local-loss decrease.The drift price is λ > 0.Source: §4.2 · p. 4
  • The stopping problem ranges over stopping times up to Emax, with the first local epoch compulsory so every participating client returns a trained model.Source: §4.2 · p. 4

4.3 Motivation from an aggregation bound

  • The drift-priced utility is motivated by a lower bound on aggregate descent under smoothness, local-objective, and gradient-discrepancy conditions.The bound applies to the plain aggregate and data-dependent stopping times.Source: §4.3 · p. 4
  • Setting λ = (L−ν)/2 makes each client’s utility the quadratic part of its contribution to a guaranteed aggregate-loss decrease, while a remaining term further penalizes drift.Source: §4.3 · p. 4
  • Because the curvature constants are unavailable for neural networks and step-normalized aggregation rescales the terms, λ is selected on validation data.Source: §4.3 · p. 4

5 FedStop: Optimal Stopping for Local Training

FedStop frames local epoch selection as finite-horizon optimal stopping, with a one-stage threshold that is optimal under nested stopping decisions. Its closed forms and deployment rule connect estimated progress and drift to client budgets, while the theory’s predictions specify when the threshold remains valid.

5.1 Threshold optimality

  • When stopping decisions are nested, the one-stage rule is optimal: continue while expected epoch benefit exceeds expected drift cost, stopping on ties.The theorem assumes integrable utility and deterministic, non-increasing net gains; its proof relies on nested one-stage stopping sets.Source: §5.1 · pp. 4–5

5.2 Closed forms for two progress profiles

  • Geometric returns yield an optimal budget that stops after one epoch exactly when κ ≥ ρ; power-law returns stop after one exactly when κ ≥ 2^-a.Under geometric returns, epoch counts vary logarithmically with the exchange ratio; under power-law returns they vary polynomially, and large exponents can make the second epoch expensive.Source: §5.2 · p. 5
  • With declining drift, the one-stage rule remains optimal when drift declines more slowly than progress; if it declines faster, stopping decisions need not be nested.For power-law returns, nestedness requires σ ≥ (1 − 1/Emax)^a; with a = 2.4 and Emax = 20, this requires σ ≥ 0.88.Source: §5.2 · p. 6

5.3 The cost of estimating the gains

  • Estimated stopping loses utility only when an estimate gives the wrong sign for whether an epoch’s net gain is positive.Each disagreement is charged the absolute net gain of the added or forgone epoch, so errors near the threshold cost little and decisions far from it are more robust.Source: §5.3 · pp. 6–7

5.4 Correlated progress

  • Trajectory-dependent correlated gains preserve optimality when the stopping sets remain nested, and the rule adjusts expected gain by ϕεk.The adjustment accounts for the latest progress deviation; when ϕ = 0, the rule reduces to the geometric and power-law corollaries.Source: §5.4 · p. 7

5.5 From theory to a deployable rule

  • FedStop estimates profile shape from recent participations and loss scale within the current round, then continues while its plug-in look-ahead statistic is positive.Each stopping decision takes O(1) time; the dominant deployment cost is evaluating training loss once per epoch.Source: §5.5 · p. 8

5.6 Predictions of the theory

  • The theory predicts geometric diminishing returns and additive, approximately stationary drift, with budgets driven by training stage rather than client identity.As global progress shrinks, κ grows and budgets fall together; the evidence plan also tests whether round-level quantities dominate client-specific budgets through random swaps.Source: §5.6 · p. 9

6 Experimental Setup

The experiments use five main federated-learning settings and compare FedStop with tuned FedAvg and several aggregation and stopping baselines. The protocol separates hyperparameter selection from reporting seeds and uses epoch-level logs and tail-accuracy metrics to assess behavior.

Settings

  • Experiments cover CIFAR-10 at two Dirichlet partitions, CIFAR-100, FEMNIST, and Shakespeare, using 100 clients with 10% uniform participation per round.Two additional CIFAR-10 partitions are used only for fitting progress profiles; settings use plain SGD and a cap of 20 local epochs.Source: §6 · pp. 9–10

Baselines and controls

  • FedStop is compared with tuned FedAvg, FedNova, an adaptive aggregation interval, FedProx, per-client early stopping, ALT, and FLASH.The client-side early-stopping and ALT/FLASH comparisons are run on CIFAR-10 at α = 0.1, FEMNIST, and Shakespeare.Source: §6 · p. 10
  • Random-swap and equal-split controls replay FedStop’s round budgets to test whether matching budgets to clients matters beyond preserving each round’s total training.The controls execute slightly more SGD steps than FedStop because clients differ in size; a round-wise validation oracle provides a reference unavailable to a real federation.Source: §6 · p. 10

Protocol

  • Hyperparameters are selected on seeds 1–5, while results use paired seeds 6–10; diverged runs remain in the evaluation at chance accuracy.The window is fixed at w = 4, and test data are never used for selection.Source: §6 · pp. 10–11

Epoch-level logs

  • Epoch-level logs record squared drift, drift increments, per-epoch displacement, and budgets; FedStop additionally logs local training loss after each epoch.Drift checks use geometric random-swap runs so observed epochs do not depend on a client’s own trajectory.Source: §6 · p. 11

Metrics

  • Evaluation reports mean test accuracy and worst-decile accuracy, with two tail definitions on FEMNIST and Shakespeare and Holm correction for headline tail comparisons.The cross-fitted tail ranks clients by validation accuracy, while the shared-index tail uses clients ranked lowest by tuned FedAvg on selection seeds.Source: §6 · p. 11

7 Testing the Theory Against Local Training

Controls suggest FedStop’s value lies more in its per-round schedule than in matching particular clients to budgets. Epoch-level evidence supports geometric diminishing returns and largely additive drift, while budgets track training stage more than stable client identity.

7.1 Diminishing returns (P1)

  • Geometric returns extrapolate better than power-law returns for 75%, 73%, 79%, and 63% of four CIFAR-10 partitions and 78% of FEMNIST.The geometric rule is active in 98% of CIFAR-10 decisions at α = 0.1 and 94% on FEMNIST, but never active on CIFAR-10 at α = 0.5, CIFAR-100, or Shakespeare.Source: §7.1 · p. 11

7.2 Drift cost (P2)

  • Per-epoch squared displacements sum to a median 0.96 (CIFAR-10) and 0.94 (FEMNIST) of final squared drift, supporting additive drift despite the first-epoch anomaly.The first epoch moves the model 3.6 (FEMNIST) to 4.9 (CIFAR-10) times further than the second, but it is compulsory and does not enter the optimal stopping decision.Source: §7.2 · p. 12

7.3 Determinants of client budgets (P3, P4)

  • Mean geometric budgets fall from 3.8 to 2.0 epochs across the first and last hundred CIFAR-10 rounds, and from 7.4 to 4.7 on FEMNIST.Round alone explains 44% of CIFAR-10 budget variance and 18% on FEMNIST; adding client identity raises explained variance by only about ten points on both.Source: §7.3 · p. 12
  • Randomly swapping client budgets and equal-split controls do not establish equivalence, though both point the same way; holding each client’s power-law budget at its run average, with the same epochs, lowers the CIFAR-10 tail by 3.9 points ([2.6, 5.3]).Source: §7.3 · pp. 12–13

7.4 Nested decisions (P5)

  • Although Corollary 5’s margin fails on every geometric client-round with at least four epochs, exact dynamic programming matches the one-stage rule using the same loss-scale estimator.The estimated reversal rates are 19% on CIFAR-10 and 27% on FEMNIST, but are indicative only because they use short, self-stopped trajectories and omit unobserved continuations.Source: §7.4 · p. 13

8 Performance in Federated Training

FedStop addresses fixed-budget fragility with progress-responsive local training, preserving mean accuracy while improving worst-decile accuracy where budgets vary. Its monitoring and client-side-rule comparisons show both practical trade-offs and settings where its tail advantage is strongest.

8.1 A fixed budget is fragile

  • The best fixed FedAvg configuration is two epochs at ηE = 0.1 on CIFAR-10 (63.8%) but one epoch at ηE = 0.1 on CIFAR-100 (31.4%), where two epochs diverge on every seed.At the main learning rate, the tuned CIFAR-100 budget is instead two epochs, so the fixed budget must be retuned across settings.Source: §8.1 · pp. 13–14, Figure 3 · p. 14

8.2 Accuracy and worst-client performance

  • FedStop never trails validation-tuned FedAvg by more than one mean-accuracy point; active-budget worst-decile gains over tuned FedAvg are 8.25 points on CIFAR-10 and 4.84 on FEMNIST.The four tail gains over tuned FedAvg and Wang et al. on CIFAR-10 and FEMNIST remain significant after Holm correction.Source: §8.2 · p. 14, Table 5 · p. 15
  • On CIFAR-10, FedStop reaches a 39.0% worst-decile tail with 24.9k local epochs, while no fixed budget matches it as cheaply and per-client early stopping ends at chance.FedStop's tail advantage over tuned FedAvg opens early and persists through training; its per-round schedule matters more than how budgets are divided among clients.Source: §8.2 · pp. 14–15

8.3 Controlling the cost of loss monitoring

  • Evaluating loss on the same fixed subset of min(128, n_i) examples bounds monitoring work by a constant per epoch, independent of local dataset size.The CIFAR-10 geometric-rule tail fell by 1.3 points in the primary run; re-selection chose the same drift price, and a repeat on the same seeds showed no tail difference from full monitoring.Source: §8.3 · pp. 16–17

8.4 Client-side stopping rules

  • On CIFAR-10, FedStop has the highest worst-decile accuracy against ALT and FLASH, ahead by 8.95 and 3.68 points respectively, but neither difference survives Holm correction.The three rules are level on FEMNIST and Shakespeare; the comparison ports stopping criteria into the same protocol rather than reproducing the complete systems.Source: §8.4 · p. 17, Table 7 · p. 17

9 Discussion

The discussion identifies boundaries on the optimality guarantee, the model's single-round scope, and validation-based drift-price selection. It also notes that the evidence is limited to small models and simulated federations, and outlines directions to address these gaps.

Discussion

  • The one-stage optimality guarantee requires non-reversing stopping decisions, while observed local progress is more autocorrelated than the stated margin condition allows.An exact trajectory-aware rule gains nothing over the one-stage rule, but the guarantee does not cover these trajectories.Source: §9 · p. 17
  • The single-round model predicts that budgets anneal but not how fast, and neural-network curvature constants are unavailable to fix the drift price, which is therefore chosen on validation data.Source: §9 · p. 17
  • The evidence covers small vision and language models in simulated federations of 100 clients; larger models and cross-device scale remain untested.Source: §9 · p. 18
  • Future work could address dependent progress, couple client stopping to the server trajectory across rounds, or reduce monitoring overhead by correcting recorded training losses.Source: §9 · p. 18

10 Conclusion

The conclusion recasts local-epoch choice as optimal stopping with priced drift, deriving a one-stage threshold and interpretable closed forms. Its evidence supports stage-dependent budgets and improved worst-off-client accuracy at limited mean-accuracy cost.

Conclusion

  • FedStop protects worst-off clients at a cost of at most one point of mean accuracy, while the stage of training—not client identity—determines how much local work pays.The decision computation and memory are negligible beside training, and the method degrades gracefully when its single parameter is set too high.Source: §10 · p. 18

Broader Impact Statement

The broader-impact statement frames FedStop as a way to improve accuracy for the least-served decile under strong data heterogeneity, while identifying deployment costs and privacy-relevant information disclosures.

Broader Impact Statement

  • Where budgets vary, clients train for more local epochs than under tuned FedAvg, consuming more energy and battery; fixed-subset monitoring reduces this overhead.The worst-off clients' gain can come with a slight loss for the best-served clients, and reported local-step counts reveal dataset size and training duration.Source: Broader Impact Statement · p. 18

Supplementary material

The supplementary material documents the reporting protocol and extends the comparisons to SCAFFOLD and additional baseline configurations. SCAFFOLD’s results vary sharply by dataset, while the extra CIFAR-10 baseline comparisons show where alternatives match or differ from FedStop.

Supplementary material

  • The appendices report prespecified study configurations not tabulated in the main text, along with diagnostics and implementation details, using the main text’s protocol.Coverage includes further baselines, ablations, budget controls, training dynamics, monitoring costs, and workload prediction.Source: Supplementary material · p. 21

A Protocol

  • All reported numbers use seeds 6–10, paired to FedStop by partition, initialization, and client sampling; confidence intervals are unadjusted for multiple comparisons.Collapsed runs are retained at chance accuracy, and paired contrasts are omitted for configurations with collapsed runs.Source: A Protocol · p. 21

B SCAFFOLD

  • SCAFFOLD collapsed on all five CIFAR-10 and four of five FEMNIST seeds, but was most accurate on Shakespeare, 4.77 points above geometric FedStop.Its Shakespeare tail was also 2.74 points higher; the comparison was untuned, and the authors note a tuned comparison would test whether the collapses are intrinsic.Source: B SCAFFOLD · p. 21

C Additional baseline configurations

  • FedNova matched geometric FedStop’s CIFAR-10 tail at E = 5, 10, and 20, using two to eight times its epochs and lower mean accuracy.The adaptive-interval option C was level with FedStop using its original ϕ; with option A’s ϕ it trained longer for lower mean accuracy.Source: C Additional baseline configurations · p. 21

D Ablations of the power-law rule

The supplementary analyses test power-law-rule variants and examine budget allocation, training dynamics, monitoring costs, and workload prediction. They report that more elaborate stopping variants do not improve the power-law rule’s tail with a confidence interval excluding zero, while also documenting settings and diagnostics that qualify the results.

D Ablations of the power-law rule

  • No power-law ablation improved its tail with an interval excluding zero; sequential and upper-confidence-bound rules trained much longer and collapsed on five and four of five seeds.The exact AR(1) rule matched the one-stage rule; against the one-stage rule with the same shrunk scale, its accuracy difference was +0.07 and tail difference +0.15, with both intervals including zero.Source: D Ablations of the power-law rule · pp. 21–22

E Budget controls

  • The budget-control decomposition estimated +1.61 tail points from the trend over rounds and +2.30 from per-client allocation of a constant budget, versus +3.91 total; all component intervals included zero.Equal split executed 4.8% more SGD steps on CIFAR-10 and 1.5% more on FEMNIST; its tail comparison established neither superiority nor equivalence.Source: E Budget controls · pp. 22–23

F Training dynamics

  • Per-client early stopping ended at chance accuracy on every CIFAR-10 seed, while tuned FedAvg’s FEMNIST worst decile stayed about four points below geometric FedStop’s.On CIFAR-10, the training-curve comparison reports two runs that never learned and three that fell back after peaking at 27–43%.Source: F Training dynamics · p. 23

G Variants on the naturally partitioned data

  • On FEMNIST and Shakespeare, forcing a second epoch cost accuracy on Shakespeare, where geometric FedStop chose one epoch; a CIFAR-10 drift price transferred to the other settings was level with the locally selected price.Figure S1’s caption identifies client-by-window epoch budgets for one reporting seed, with clients sorted by local data size and unsampled windows shown white.Source: G Variants on the naturally… · p. 23, Figure S1 · p. 24

H Shape estimator and pilot priors; I Client-side stopping rules

  • The shape estimator fits the current and prior four participations with profile-specific searches and pilot priors; ALT and FLASH thresholds are selected on seeds 1–5 using validation data and calibrated-work tie-breaking.The selected drift prices span 6.7×10^-4 to 1.27×10^4, while the power-law exponent rarely leaves its prior.Source: H Shape estimator and pilot priors · p. 24, I Client-side stopping rules · p. 24

J Computational cost and subset monitoring

  • Fixed-subset monitoring lowered work in every setting, but geometric FedStop still used 149% of two-epoch FedAvg’s work on CIFAR-10 and 203% on FEMNIST to reach the target.The subset changed some decisions; on CIFAR-10, geometric clients executed 88 more epochs per run, and no mean-accuracy difference excluded zero.Source: J Computational cost and subset… · p. 25, Figure S2 · p. 27

K Synthetic benchmark; L Predicting client workload

  • On the synthetic benchmark, FedStop improved over one-epoch FedAvg by 3.4 points in accuracy and 5.2 points in worst-decile accuracy, but was statistically level with every other method.Budget history predicted the next participation’s budget with mean absolute error 0.30 epochs on CIFAR-10 and 1.44 on FEMNIST, outperforming the previous cohort mean in both settings.Source: K Synthetic benchmark · pp. 25–26, L Predicting client workload · p. 26

M Diagnostics from the epoch logs

  • Epoch-log diagnostics indicate that the deployed drift estimate prices continuation above later realized drift, while short-trajectory estimates are explicitly described as indicative only.The geometric shape sat at its lower search bound in most CIFAR-10 decisions, and under the fitted autoregressive model a stopping decision would reverse in about a fifth to a quarter of cases.Source: M Diagnostics from the epoch logs · p. 26

Related papers

Loading 2610.21832v1…