Decision-Focused Federated Learning Under Heterogeneous Objectives and Constraints

2026-04-21

AI Summary

This work addresses the challenge of heterogeneous objective functions and constraints across clients in federated learning when solving downstream optimization tasks. To tackle this issue, the authors introduce the predict-then-optimize paradigm into federated learning and propose a joint training framework based on Smart Predict-then-Optimize plus (SPO+). The key innovation lies in decomposing the SPO+ loss bound via support functions to disentangle shifts in objectives from shifts in feasible sets, and in establishing local and federated decision rules grounded in excess risk analysis. The method integrates SPO+ surrogate losses, FedAvg aggregation, and stability analysis under strongly convex or polyhedral feasible regions. Experiments demonstrate that the framework exhibits robust performance in strongly convex settings, while under polyhedral constraints, its efficacy is primarily influenced by constraint heterogeneity—particularly pronounced for clients with larger sample sizes.

Technology Category

• Machine Learning: Distributed Machine Learning & Federated Learning
• Search and Optimization: Non-convex Optimization
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Systems and Infrastructure for Web, Mobile and WoT: Federated Web and WoT systems, including distributed, federated and edge-based data processing

Abstract

We consider what we refer to as {Decision-Focused Federated Learning (DFFL)} framework, i.e., a predict-then-optimize approach employed by a collection of agents, where each agent's predictive model is an input to a downstream linear optimization problem, and no direct exchange of raw data is allowed. Importantly, clients can differ both in objective functions and in feasibility constraints. We build on the well-known SPO+ approach and develop heterogeneity bounds for the SPO+ surrogate loss in this case. This is accomplished by employing a support function representation of the feasible region, separating (i) objective shift via norm distances between the cost vectors and (ii) feasible-set shift via shape distances between the constraint sets. In the case of strongly convex feasible regions, sharper bounds are derived due to the optimizer stability. Building on these results, we define a heuristic local-versus-federated excess risk decision rule which, under SPO+ risk, gives a condition for when federation can be expected to improve decision quality: the heterogeneity penalty must be smaller than the statistical advantage of pooling data. We implement a FedAvg-style DFFL set of experiments on both polyhedral and strongly convex problems and show that federation is broadly robust in the strongly convex setting, while performance in the polyhedral setting degrades primarily with constraint heterogeneity, especially for clients with many samples. In other words, especially for the strongly convex case, an approach following a direct implementation of FedAvg and SPO+ can still yield promising performance even when the downstream optimization problems are noticeably different.

Problem

• Federated Learning
• Heterogeneous Objectives
• Feasibility Constraints
• Predict-then-Optimize
• Decision-Focused Learning

Innovation

• Decision-Focused Learning
• Federated Learning
• Heterogeneous Constraints
• SPO+
• Optimization-Aware Prediction