FedRSPO+: A Heterogeneity-aware Algorithm for Decision-focused Federated Learning

AI summary

Overview

• Research area: Decision-focused learning (DFL), federated learning (FL), and contextual optimization — specifically the intersection where predictive models are trained for downstream optimization under client heterogeneity.
• Technical level: Advanced. The paper relies on convex analysis tools (support functions, Bregman divergences, Moreau envelopes, projection stability) and derives regret, heterogeneity, and convergence bounds.
• Scope: One-sentence scope: the paper introduces a smoothed predict-then-optimize surrogate (RSPO+) and a modular federated training pipeline (FedRSPO+) that stabilizes collaborative decision-focused learning when clients have different objectives and different feasible sets, and validates it on a synthetic knapsack task, a PJM energy pricing case study, and a Warcraft shortest-path benchmark.

What This Paper Is About

Decision-focused learning trains predictors so that the downstream optimization decision is good, but existing methods assume the data is centralized. Moving to cross-silo federated learning is natural for privacy and regulatory reasons, yet standard federated pipelines optimize prediction quality rather than decision quality, and they ignore the fact that each client may have its own objective and its own feasible region.

The core difficulty is geometric: when clients have different linear objectives over polyhedral feasible sets, the optimal decision is a vertex, and arbitrarily small perturbations can switch the optimizer to a completely different extreme point — producing a jump on the order of O(diam(S_i)). This makes client updates and aggregation unstable. The paper's goal is a federated algorithm whose decision map is smooth enough to aggregate stably, while still remaining provably tied to the original downstream regret.

Key Contributions

1. RSPO+ surrogate. The authors define a regularized surrogate, RSPO+, by replacing the support function ξ_S of the feasible set with its λ-Moreau envelope, M_{λξ_S}(u) = (1/2λ)||u||² − (λ/2)d_S²(u/λ). This yields a closed-form loss in terms of projections onto S, ℓ_RSPO+(c, ĉ; λ) = (λ/2)(||u − Π_S(v)||² − ||u − Π_S(u)||²) with u = (c − 2ĉ)/λ and v = −c/λ, and the authors show the surrogate upper bounds downstream decision error and regret.
2. FedRSPO+ federated framework. A modular four-stage pipeline — calibration, federated training, personalization, and freezing — compatible with standard FL personalization and aggregation methods (FedAvg, FedProx, Ditto are tested), with per-client calibration of λ_j and η_j and co-annealed λ/η scheduling.
3. Heterogeneity-aware theory. The paper formalizes why naive federation of SPO+ fails, then derives point-wise cross-client loss bounds for RSPO+ that scale with objective heterogeneity c_d and feasible-set heterogeneity δ, vanish at homogeneity, and require no strong convexity — in contrast to SPO+ bounds that carry an irreducible diameter term D_min and need strong convexity for basic stability.
4. Empirical evaluation across three settings. Controlled synthetic fractional knapsack instances isolating objective, constraint, and data heterogeneity; a PJM electricity pricing case study on real data; and the 12×12 Warcraft shortest-path benchmark intended to probe structural mismatch.

Main Findings

• RSPO+ bounds the downstream regret. Theorem 1 gives ℓ_RSPO+ ≥ R(c, ŵ) + λ||ŵ − w*||², hence ||ŵ − w*|| ≤ sqrt(ℓ_RSPO+/λ). Corollary 2 gives an unregularized regret bound R_unreg(c, ŵ) ≤ ℓ_RSPO+(c, ĉ; λ) + (λ/2)R², so the surrogate controls original LP regret up to an O(λR²) bias.
• The surrogate is well behaved. Proposition 3 states ℓ_RSPO+ ≥ 0 as a divergence, is globally convex in ĉ, is 2D-Lipschitz continuous under compactness of S, and is 4/λ-smooth (hence λ/4-cocoercive). Its gradient is ∇_ĉ ℓ_RSPO+ = −2(Π_S(u) − Π_S(v)).
• Heterogeneity bounds improve on SPO+. For SPO+ the rate is O(c_d + δ_N + D_min); for RSPO+ it is O(c_d/λ + sqrt(δ/λ) + δ). The SPO+ bound does not vanish at homogeneity (D_min persists) and requires strong convexity; the RSPO+ bound vanishes at homogeneity and does not require strong convexity.
• Projection gives sqrt(δ) stability without strong convexity. For compact convex sets with Hausdorff distance δ, ||Π_{S1}(p) − Π_{S2}(p)|| ≤ sqrt(2δ(||p|| + R_max)), which holds for all convex sets including polyhedra.
• An explicit bias–stability tradeoff. The combined error scales roughly as λR² + c_d/λ + sqrt(δ/λ), which motivates annealing λ only until the regularized objective aligns with the original LP objective rather than driving λ to zero.
• Theorem 3 decomposes cross-client loss discrepancy into an envelope shift T(ĉ) = D_i c_d + δ(||c_j − 2ĉ|| + ||c_j|| + 2λR_max) and a decision shift B(ĉ) = (2||ĉ|| + 2λR_max)(c_d/λ) + sqrt(2δ(c_min/λ + R_max)).
• Knapsack: RSPO+ has the lowest relative regret across all heterogeneity configurations. In the homogeneous setting (0,0,0), federated RSPO+ reaches 2.3 ± 0.4 percent versus 9.7 ± 2.1 for SPO+, 8.5 ± 1.5 for PFYL, 16.3 ± 2.7 for DPO, and 15.9 ± 3.3 for MSE; after personalization RSPO+ is 1.9 ± 0.3. At the highest objective heterogeneity (1,0,0) RSPO+ is 3.0 ± 0.3 federated and 2.0 ± 0.3 personalized; at (0,1,0) it is 4.6 ± 0.5 and 2.6 ± 0.4; at (0,0,1) it is 9.3 ± 1.2 and 3.4 ± 0.2.
• Data heterogeneity is the most damaging axis. Across all methods, data heterogeneity hurt performance most, constraint heterogeneity affected regret moderately, and objective heterogeneity affected it minimally.
• Comparison with Interp-DFFL and OptNet (knapsack). FedRSPO+ is better in the homogeneous (1.86 ± 0.25 vs 2.42 ± 0.30), objective-heterogeneous (2.02 ± 0.26 vs 2.42 ± 0.33), and constraint-heterogeneous (2.58 ± 0.38 vs 3.51 ± 0.77) settings; under pure data heterogeneity (0,0,1), Interp-DFFL is stronger (2.47 ± 0.39 vs 3.37 ± 0.22), which the authors attribute to its explicit local–federated interpolation. The OptNet-style baseline is substantially weaker (9.4 to 11.3 percent), suggesting gains come from the RSPO+ surrogate rather than quadratic regularization alone.
• PJM ablation shows each component helps. Mean absolute regret falls from 175.453 ± 7.260 (MSE) to 81.676 ± 1.768 (SPO+) to 72.890 ± 4.231 (RSPO+), then 68.136 ± 1.402 with personalization, 63.006 ± 1.059 with per-client λ/η, and 60.654 ± 0.796 for the full pipeline — a 65 percent reduction over MSE and 26 percent over naive SPO+ federation. PFYL (67.8 ± 1.2) and DPO (74.7 ± 2.2) also performed well here.
• Modularity results. FedProx matches the full pipeline closely at 60.602 ± 0.623 with reduced variance, while Ditto performs comparably to RSPO+ with personalization only (69.482 ± 2.932), which the authors attribute to Ditto's proximal regularization conflicting with RSPO+'s gradual scheduling.
• Sweeps over communication, client count, and α. Improvement over SPO+ ranges from 1.1 percent (α = 10⁻³) to 5.4 percent (α = 0.1); from 0.9 percent (5 clients) to 12.5 percent (50 clients); and from 5.7 percent to 6.1 percent as the fraction of active clients rises from 10 percent to 100 percent. Increasing the number of clients produced the largest gain, 12.5 percent at 50 clients.
• Convergence. Theorem 2 gives an O(1/Ση_t + Ση_t²/Ση_t + Ση_t Δ_t / Ση_t) bound; with fixed λ and η_t ≍ 1/t this reduces to O(1/log T), and after freezing it becomes O(1/(Tη_f) + η_f).
• Warcraft shortest path. The provided paper content is truncated mid-sentence in this section ("adding the quadratic term dest"), so the Warcraft results are not reported in the available text.

Methodology in Plain English

The authors start from SPO+, the most widely used convex surrogate in decision-focused learning. SPO+ can be written as a Bregman divergence built on the support function of the feasible set. Because that support function encodes sharp vertex solutions, the authors replace it with its smoothed (Moreau envelope) version. The practical consequence is that the decision becomes a Euclidean projection of the scaled cost vector onto the feasible set, which moves continuously when the input changes — instead of jumping from vertex to vertex.

They then prove that this smoothed loss still upper bounds the downstream regret of the regularized decision, that the extra bias shrinks as the smoothing parameter λ shrinks, and that the loss difference between two heterogeneous clients depends on their objective gap and the Hausdorff distance between their feasible sets — with a square-root dependence on that distance and no need for strong convexity.

For the federated pipeline, each client first calibrates its own λ_j = ||c̄_j||/R_j and step size η_j = α||c̄_j||/(4R_j) using its average cost magnitude and feasible-set radius. Federated training then runs SGD on the regularized objective and aggregates with FedAvg weighted by client sample counts. A personalization phase runs additional local SGD while annealing λ downward with a schedule linear in 1/λ_t (λ_t = λ_0/(1 + κt)), so that the loss and gradient drift are uniformly controlled; the step size is co-annealed proportionally. Annealing stops ("freezing") when a data-dependent normalized regularization bias Ũ_t computed on a client validation subset falls below a tolerance τ, so the surrogate does not drift all the way back to the unstable unregularized problem.

Evaluation uses the original downstream LP decision, not the regularized one; the paper notes the theoretical bounds apply directly to the regularized decision and act as a surrogate-level guarantee for the evaluated LP decision, with the distinction disappearing below a problem-dependent exact-regularization threshold.

Why This Matters

This work is distinctive because it directly targets a failure mode that prediction-focused federated learning does not see: heterogeneity in the downstream optimization problem itself (different objectives, different feasible sets), not just heterogeneity in data distributions. By showing that a projection-based smoothing removes the diameter penalty that persists in SPO+-based federation even as heterogeneity goes to zero, it offers a principled alternative to interpolation or proximal heuristics, and it does so with bounds that hold without strong convexity. It also demonstrates that the surrogate is a loss rather than an aggregation rule, so it composes with existing FL personalization methods.

Real-world applications implied by the paper's settings:

• Electricity procurement and demand response: the PJM case study models aggregators and large consumers forecasting day-ahead prices to choose the cheapest k hours meeting a fixed daily demand.
• Logistics and routing: the Warcraft shortest-path benchmark stands in for path selection under heterogeneous origins, destinations, and terrain costs.
• Resource allocation under capacity constraints: the fractional knapsack formulation represents budget- or capacity-limited allocation decisions.
• Cross-silo analytics in regulated or competitive industries: the paper's motivating scenario is organizations that cannot centralize data for regulatory, privacy, or competitive reasons but still need collaborative decision models.

Industry relevance: the paper provides a communication-budget-aware training procedure with a stopping rule that operators can tune, reports relative regret under varying fractions of active clients and varying client counts, and shows compatibility with FedProx and Ditto — all of which matter when deciding whether to deploy collaborative decision models across organizational boundaries.

Future Directions

• Coordinating regularization-based FL methods with the schedule. The authors observe that Ditto's proximal regularization toward the global model conflicts with RSPO+'s gradual transition from regularized to unregularized objectives, and call this an interesting direction for future work.
• Closing the residual gap from the freeze rule. The freeze rule stops annealing before λ reaches the exact-regularization threshold, and the paper states that this residual gap is quantified in Section C.6 of the Appendix rather than eliminated.
• Combining RSPO+ with interpolation-style aggregation. Since RSPO+ is a loss and Interp-DFFL is an aggregation rule, the authors note they are not mutually exclusive and that Interp-DFFL is stronger under pure data heterogeneity, leaving open how to combine the two.
• Behavior under structural mismatch. The paper frames the Warcraft shortest-path experiment as a way to evaluate how RSPO+'s regularization affects performance when the optimization task relies on the structure of the unregularized objective — the question it was designed to probe, though the results are not present in the available content.

Target Audience

Researchers and practitioners working at the intersection of decision-focused learning and federated learning, particularly those with a background in convex optimization and contextual stochastic optimization. It will be most useful to readers comfortable with support functions, Bregman divergences, Moreau envelopes, and regret analysis, as well as to applied researchers in energy markets, operations, and privacy-constrained multi-organization analytics who want a concrete, modular algorithm with an ablation-style evaluation on real data and an open-source implementation (https://github.com/kziliask/RSPO_Plus).

Authors’ abstract

Decision-focused learning (DFL) trains predictive models for downstream optimization, but existing methods largely assume centralized data. In cross-silo settings, federated learning offers a natural alternative, yet standard federated methods optimize prediction over decision quality and do not address heterogeneity in downstream objectives or feasible sets. This heterogeneity is especially challenging for DFL because small perturbations in polyhedral problems can cause discontinuous changes in optimal decisions, destabilizing client updates and aggregation. We propose FedRSPO+, a heterogeneity-aware framework for decision-focused federated learning, built on RSPO+, a regularized predict-then-optimize surrogate that smooths the decision map through projection. We show that RSPO+ upper bounds decision error and regret for the regularized decision and, under exact regularization and consistent LP solution selection, for the original LP decision. We further derive cross-client heterogeneity bounds that depend on both objective and feasible-set heterogeneity, vanish at homogeneity, and require no strong convexity. FedRSPO+ uses an annealed, modular training procedure compatible with standard federated personalization and aggregation methods. Experiments on synthetic knapsack, shortest-path, and real-world energy pricing tasks compare against prediction-only federated learning and DFL baselines under varying heterogeneity and communication budgets. Results suggest that smoothing is a useful ingredient for stable collaborative decision learning and provide a heterogeneity-aware foundation for federated DFL.