Prediction-Focused Learning

• Prediction-Focused Learning is a design principle that prioritizes predictive performance by allocating model capacity to task-relevant features.
• It leverages structured mechanisms in latent-variable models, predict-then-optimize pipelines, and personalized federated learning to separate predictive signals from noise.
• Empirical studies reveal that while PFL ensures smoother training, it also presents trade-offs in decision quality, feasibility, and model interpretability.

Prediction-Focused Learning (PFL) designates a family of learning strategies in which model construction, representation selection, or personalization is organized around predictive performance on the task of interest rather than around undifferentiated data likelihood or a single consensus model. In the literature considered here, the term appears in supervised latent-variable modeling, in predict-then-optimize pipelines, and, with an overloaded acronym, in federated learning where “PFL” often denotes Personalized Federated Learning and the prediction-focused principle is instantiated as per-client predictive optimization (Ren et al., 2019, Jeon et al., 10 Sep 2025, Chen et al., 2022).

1. Terminological scope and core formulations

Across these lines of work, PFL is not a single algorithm but a design principle. In latent-variable models, it means allocating representational capacity to the parts of the input that improve prediction, while routing irrelevant structure into auxiliary channels (Ren et al., 2019, Sharma et al., 2021). In predict-then-optimize settings, it denotes the standard two-stage regime in which a predictor is trained only with a prediction loss such as MSE and the downstream optimization problem is solved only after prediction, without affecting training (Jeon et al., 10 Sep 2025, Capitaine et al., 19 May 2025). In federated learning, the acronym frequently refers to Personalized Federated Learning, where the operative notion of “prediction focus” is that each client’s predictive performance is primary, rather than the performance of a single global model (Zhao et al., 31 Dec 2025, Chen et al., 2022).

Usage of PFL	Central training signal	Representative works
Latent-variable and interpretable modeling	Task-relevant feature or dimension selection inside a generative model	(Ren et al., 2019, Sharma et al., 2021)
Predict-then-optimize	Prediction loss on uncertain parameters, followed by separate optimization	(Jeon et al., 10 Sep 2025, Capitaine et al., 19 May 2025, Mandi et al., 6 Oct 2025)
Personalized federated learning	Per-client or per-cluster predictive quality under non-IID data	(Chen et al., 2022, Zhao et al., 31 Dec 2025, Xing et al., 23 Oct 2025)

A common thread is that PFL privileges predictive sufficiency over exhaustive input modeling. In the topic-modeling formulation, the representation should remain interpretable and task-relevant (Ren et al., 2019). In mixture and HMM formulations, a limited number of components should be spent on dimensions that matter for predicting $Y$, not on dominant but irrelevant structure in $X$ (Sharma et al., 2021). In federated learning, prediction focus means that each client, cluster, or prompt-conditioned predictor should optimize its own prediction loss rather than conform to a single shared model (Zhao et al., 31 Dec 2025, Jimenez-Gutierrez et al., 23 Dec 2025).

2. Prediction-focused latent-variable models

A canonical formulation appears in "Prediction Focused Topic Models via Feature Selection" (Ren et al., 2019). That work starts from supervised topic modeling and identifies three failure modes of standard sLDA-style training: cardinality mismatch between many words and one label, model misspecification, and pollution from irrelevant vocabulary. Its remedy is a dual-channel generative model in which words are partitioned into a relevant channel, modeled by predictive topics $\{\beta_k\}$, and an irrelevant channel, modeled by a background topic $\pi$. A Bernoulli switch $\xi_{dn}$ determines whether token $n$ in document $d$ contributes to the predictive topic mixture $\theta_d$, and only the relevant channel influences the label model. The hard support constraint $\beta_k^\top \pi = 0$ enforces a vocabulary partition, while the word inclusion prior $p$ becomes the interpretable control over how much of the vocabulary may affect prediction.

The model’s variational selector $X$0 is indexed by vocabulary type rather than document, so a word is globally treated as predictive or irrelevant. This embeds feature selection inside the generative model rather than applying it as a preprocessing step. Empirically, the method preserves or improves prediction while sharply improving topic coherence. On Pang & Lee, pf-sLDA reports coherence $X$1 versus $X$2 for pc-sLDA and $X$3 for sLDA, with RMSE $X$4 versus $X$5 for pc-sLDA. On Yelp, it reports coherence $X$6 with RMSE $X$7. On the ASD EHR task, it improves AUC to $X$8 while also improving coherence (Ren et al., 2019).

"On Learning Prediction-Focused Mixtures" (Sharma et al., 2021) generalizes the same principle to Gaussian mixtures and HMMs. There, each input dimension has a Bernoulli switch $X$9. If $\{\beta_k\}$0, the dimension is modeled through cluster- or state-dependent emissions $\{\beta_k\}$1; if $\{\beta_k\}$2, it is assigned to a simple cluster-independent noise model $\{\beta_k\}$3. The posterior relevance variable $\{\beta_k\}$4 therefore becomes an interpretable relevance probability for dimension $\{\beta_k\}$5. The result is a generative model whose maximum-likelihood solution, under limited capacity, preferentially allocates mixture components or HMM states to predictive dimensions. The paper shows that this mechanism outperforms non-prediction-focused mixtures in AUROC under constrained capacity and is easy to optimize with variational EM (Sharma et al., 2021).

Taken together, these models establish an important point: prediction focus need not imply abandoning generative semantics. The model may remain fully generative, with maximum-likelihood training intact, while the architecture itself creates an asymmetry between predictive signal and nuisance structure. This suggests that PFL can be realized by structural priors and latent switches, not only by discriminative losses.

3. Predict-then-optimize, decision quality, and the limits of pure prediction loss

In the predict-then-optimize literature, PFL has a narrower and more operational meaning. A predictive model $\{\beta_k\}$6 estimates uncertain parameters $\{\beta_k\}$7 from features $\{\beta_k\}$8, then an optimizer computes a decision $\{\beta_k\}$9 by solving a downstream problem such as

$\pi$0

Training is performed only on a prediction loss $\pi$1, while the optimization stage is absent from gradient computation (Jeon et al., 10 Sep 2025).

"Prediction Loss Guided Decision-Focused Learning" (Jeon et al., 10 Sep 2025) gives a particularly explicit characterization of this regime. It treats PFL as stable because $\pi$2 is usually smooth and well-conditioned, but also as potentially misaligned with the downstream objective. The paper’s empirical comparison across knapsack, budget allocation, and portfolio optimization shows that PFL can have substantially higher regret in tasks where small prediction errors induce large combinatorial or ranking errors, while remaining comparatively competitive in portfolio optimization, where prediction quality is more directly related to decision quality (Jeon et al., 10 Sep 2025).

"Feasibility-Aware Decision-Focused Learning for Predicting Parameters in the Constraints" (Mandi et al., 6 Oct 2025) sharpens this critique for constrained optimization with uncertain constraint parameters. In that setting, prediction errors can change not only objective value but the feasible set itself. The paper argues that PFL, instantiated as MSE on predicted constraint parameters, does not distinguish between errors that preserve feasibility and errors that induce infeasible or severely suboptimal decisions. It introduces two losses: an Infeasibility-Preserving Loss, which penalizes predicted solutions that are infeasible under the true constraints, and an Optimality-Preserving Loss, which penalizes predicted constraints that make the true optimal solution infeasible. The combined loss

$\pi$3

exposes a single tunable parameter $\pi$4 governing infeasibility aversion (Mandi et al., 6 Oct 2025).

The conceptual implication is that pure PFL is often an incomplete surrogate. Better prediction error does not necessarily imply better decisions, and when parameters enter the constraints, it may not even imply feasible decisions. This is the main point of divergence between PFL and decision-focused learning: PFL optimizes parameter fidelity, whereas DFL optimizes the operational consequences of those parameters.

4. Hybrid, feasibility-aware, and online extensions

Later work does not simply reject PFL; it often repurposes its geometry or stability. In "Prediction Loss Guided Decision-Focused Learning" (Jeon et al., 10 Sep 2025), the prediction gradient and decision gradient are combined through a normalized direction

$\pi$5

where $\pi$6 and $\pi$7 are unit gradient directions and $\pi$8 is the geometric mean of gradient norms. The decaying coefficient

$\pi$9

lets prediction loss guide training early and shifts emphasis toward decision loss later. For $\xi_{dn}$0, the paper proves convergence to a Pareto stationary point under differentiability, Lipschitz smoothness, bounded geometric-mean gradient magnitude, and standard diminishing step-size assumptions, with rate $\xi_{dn}$1 (Jeon et al., 10 Sep 2025).

"Online Decision-Focused Learning" (Capitaine et al., 19 May 2025) extends the distinction between PFL and DFL to dynamic, non-stationary environments. It formalizes PFL as an online learner that minimizes a prediction risk such as MSE and then plugs predictions into a linear optimization over a simplex or a bounded convex polytope. The obstacle for DFL is that the true decision loss is piecewise constant in the predictor parameters because the optimizer’s argmin changes discontinuously. The proposed DF-OGD therefore regularizes the inner problem,

$\xi_{dn}$2

and then uses an optimistic update based on an approximate offline oracle. The paper establishes expected dynamic-regret bounds for both simplex and general polytope settings and shows empirically on an online knapsack task that DF-OGD achieves lower average cumulative decision cost than a prediction-focused PF-OGD baseline, even though PF-OGD attains lower MSE (Capitaine et al., 19 May 2025).

These works suggest a more nuanced role for PFL in decision problems. Prediction-focused training remains attractive because of its smoothness and curvature, but when downstream decision boundaries are sharp, it functions better as a stabilizing prior or auxiliary signal than as the sole training target.

5. Prediction focus in personalized federated learning

In federated learning, the acronym PFL usually denotes Personalized Federated Learning, but the underlying principle is still prediction focus: the target is each client’s predictive performance under non-IID data. "pFL-Bench: A Comprehensive Benchmark for Personalized Federated Learning" (Chen et al., 2022) formalizes this perspective through a generalized objective involving performance of a global model on participating clients, performance of local personalized models, generalization to new clients, and regularization between global and local models. It organizes methods into fine-tuning, model mixture, model interpolation, and model regularization, and reports that no single method dominates across all datasets. The benchmark also shows that fine-tuning often improves both participating-client accuracy and non-participating-client generalization, while many personalized methods trade better local accuracy for worse transfer to unseen clients (Chen et al., 2022).

Traffic forecasting offers a concrete instance of this interpretation. "AutoFed: Manual-Free Federated Traffic Prediction via Personalized Prompt" (Zhao et al., 31 Dec 2025) defines PFL for traffic prediction as per-client optimization of MAE, RMSE, MSE, and MAPE under strongly non-IID spatial-temporal data. Its architecture combines a local AGCRN-based personalized predictor with a federated representor that generates a prompt matrix $\xi_{dn}$3 for the decoder. The total client loss

$\xi_{dn}$4

makes the representation-learning term adaptive rather than manually tuned. The paper reports communication cost of about $\xi_{dn}$5 parameters per round, compared with about $\xi_{dn}$6 for FedAvg and $\xi_{dn}$7 for FedGCN, while achieving best or second-best performance across travel demand and traffic flow prediction scenarios (Zhao et al., 31 Dec 2025).

Several recent federated methods make the prediction-focused principle explicit in how they weight clients, partition parameters, or define clusters. "CO-PFL: Contribution-Oriented Personalized Federated Learning for Heterogeneous Networks" (Xing et al., 23 Oct 2025) replaces heuristic client weighting with dual-subspace contribution scores based on gradient direction discrepancy and leave-one-out prediction deviation, couples this with parameter-wise personalization masks, and reports gains on CIFAR10, CIFAR10C, CINIC10, and Mini-ImageNet. "RCC-PFL: Robust Client Clustering under Noisy Labels in Personalized Federated Learning" (Ali et al., 25 Mar 2025) performs one-shot, label-agnostic clustering by feature covariance statistics before training and matches a genie-aided optimum-clustering baseline under both class-independent and class-dependent noisy-label settings. "Clust-PSI-PFL: A Population Stability Index Approach for Clustered Non-IID Personalized Federated Learning" (Jimenez-Gutierrez et al., 23 Dec 2025) computes per-client label-skew PSI features, selects the number of clusters by silhouette score, and reports up to $\xi_{dn}$8 higher global accuracy than its baselines together with a relative fairness improvement of $\xi_{dn}$9 under severe non-IID label skew. "HEART-PFL" (Kim et al., 25 Mar 2026) combines hierarchical prototype alignment with adversarial knowledge transfer, trains only $n$0 adapter parameters, and reports personalized accuracies of $n$1 on CIFAR-100, $n$2 on Flowers-102, and $n$3 on Caltech-101 under Dirichlet non-IID partitions (Kim et al., 25 Mar 2026).

This federated strand broadens the meaning of PFL. The central object is no longer only a predictive representation or a two-stage predictor, but an allocation of shared and private computation that preserves client-specific predictive quality. A plausible implication is that prediction focus, in this literature, serves as the operational criterion for deciding what should be shared, what should remain local, and how much influence each client should exert on collective updates.

6. Evaluation criteria, misconceptions, and open problems

The literature evaluates PFL with markedly different metrics depending on the problem class. In topic models, the relevant pair is predictive performance together with interpretability, measured by RMSE or AUC alongside normalized PMI coherence (Ren et al., 2019). In mixtures and HMMs, AUROC is paired with held-out $n$4 and $n$5, reflecting the tension between predictive utility and generative fit (Sharma et al., 2021). In decision problems, normalized regret, infeasibility rate, and dynamic regret become primary, and these are often only weakly correlated with MSE (Capitaine et al., 19 May 2025, Mandi et al., 6 Oct 2025). In personalized federated learning, evaluation expands further to include weighted and unweighted client accuracy, generalization to new clients, fairness, convergence speed, computation, communication, and robustness to attacks (Chen et al., 2022, Qin et al., 2023).

Several misconceptions recur across the literature. One is that PFL and DFL are interchangeable. The decision-focused papers instead treat them as distinct training paradigms: PFL optimizes parameter prediction, DFL optimizes decision quality, and the two may disagree sharply when the downstream map from parameters to decisions is discontinuous or feasibility-sensitive (Jeon et al., 10 Sep 2025, Mandi et al., 6 Oct 2025). Another is that prediction focus requires a purely discriminative model. The topic-model and mixture-model papers show the opposite: fully generative models can be made prediction-focused through structured partitions between predictive and non-predictive channels (Ren et al., 2019, Sharma et al., 2021). A third is that personalization in federated learning must mean one fully local model per client. Clustered models, prompt-conditioned decoders, local heads on shared backbones, and parameter-wise masks all implement prediction-focused personalization without collapsing into isolated training (Zhao et al., 31 Dec 2025, Jimenez-Gutierrez et al., 23 Dec 2025).

Open problems also differ by setting. In latent-variable PFL, a single noise topic or a binary relevant/irrelevant partition may be too crude when nuisance structure is itself rich or context-dependent (Sharma et al., 2021, Ren et al., 2019). In decision-focused settings, solver calls remain expensive, exact gradients are often unavailable, and theoretical analyses usually require smoothness, oracle access, or bounded non-stationarity (Capitaine et al., 19 May 2025, Mandi et al., 6 Oct 2025). In federated settings, strong results often come with extra clustering, prompt-generation, masking, or distillation overhead, while theory lags behind practice on convergence, robustness, and privacy (Xing et al., 23 Oct 2025, Zhao et al., 31 Dec 2025, Kim et al., 25 Mar 2026). Robustness is a particularly salient issue: "Revisiting Personalized Federated Learning: Robustness Against Backdoor Attacks" (Qin et al., 2023) shows that partial model-sharing methods such as FedBN and FedRep can significantly boost robustness against backdoor attacks, whereas full model-sharing personalized methods do not show comparable robustness; the same study proposes a lightweight post hoc defense, Simple-Tuning, which retrains only the last linear layer locally and substantially reduces attack success rates (Qin et al., 2023).

Taken as a whole, the literature presents PFL as a unifying principle rather than a fixed recipe. Whether the object is a topic model, a mixture, a portfolio optimizer, or a federated predictor, the central question is the same: which aspects of the model, the data, or the collaboration structure should be preserved because they improve prediction on the target task, and which aspects should be excluded, regularized, or localized because they do not.