Paper: Task Generalization in Decision-Focused Learning (ICLR-2024)

Reviewer comment (Theory):
Since \(c\) is assumed to be unknown, in which sense may we assume that they are known in the training data? Is there any real-world application that can support such a setting?

Author response:
**Can $c$ be assumed to be present in training data**

In real-world predict-then-optimize applications, $c$ is typically uncertain at solving time, but revealed later on. For instance, consider a production company that must schedule its production processes based on uncertain future energy prices (which take the role of $c$), which are to be predicted based on correlated contextual features. Although the energy prices are uncertain at decision-making time, they will become known in due time, and can thus become part of a set of historical training data. Similarly, consider a delivery company that must make its routing decisions based on future traffic conditions (which again take the role of $c$). Although the traffic conditions are not yet known when the routes are to be computed, in time, they will become known.

Reviewer comment (Novelty):
The proposed method is reasonable but not very significant. On the one hand, it is not clear in theory why samples from multiple tasks can improve the generalization performance. On the other hand, while being more complex, SPOExplGen does not demonstrate a significant improvement in generalization performance when compared to SPOImplGen.

Author response:
**Why the methods work**

While the merits of our work are mostly evidenced empirically, and while a theoretical analysis was outside of the scope of this paper, there is certainly a strong intuitive explanation of why training on multiple tasks improves generalization performance. When a predictive model is trained on a single task, the only aim in training is to achieve the best possible performance on that task. While this is commonly done in existing DFL literature, we show that such models perform badly on altered optimization tasks. When, instead, training implicitly on a distribution of tasks, it is the average performance across this distribution that matters, and thus the trained models tend to have a more robust generalization performance. However, because the predictive trade-offs required to perform well on one task may run counter to the trade-offs required for another task – and because these models do not know which task they are predicting for at any point in time – they generally are not able to excel on individual tasks. So, to be able to excel on individual tasks, the model must receive some knowledge about the current task, which is where task representations and explicit generalization come in.

A further theoretical analysis of the task generalization problem and the proposed approaches for tackling it would definitely strengthen this line of work. But given the 9-page limit of this conference paper, we considered this out of scope.


**The benefit of explicit generalization over implicit generalization**

There are indeed cases where explicit generalization does not offer a large benefit over implicit generalization. However, the claim that SPOExplGen does not demonstrate a significant improvement in generalization performance when compared to SPOImplGen is false. In many cases, the benefit is large: in knapsack weight generalization, the improvement in regret runs up to 64% (Table 1). In knapsack item subset generalization for large subsets, there is an improvement of 40% (Table 1). In facility location capacity generalization, the improvement goes up to 61% (Table 2), and for facility location demand generalization, it goes up to 51% (Table 2). These are all significant improvements, on top of the improvements that SPOImplGen already achieves with respect to regular SPO.

Reviewer comment (Experiments):
Since the paper considers samples from multiple tasks, it would be interesting to explore the generalization performance between different instance sizes.

Author response:
**Task generalization between different instance sizes**

Indeed, task generalization across different instances sizes is certainly an interesting setting to investigate. In the present work, we always assume the existence of a ‘base task’ on which the MSE and SPO models are trained. So changing the instance size with respect to the base task could mean two things: making it smaller or making it larger. In our knapsack item subset generalization experiment, we explored the effect of generalizing over smaller instances. Setting up an experiment around generalizing to larger instances is more tricky, however. This is the case because our experimental setup (which was inspired by related work such as (Elmachtoub & Grigas, 2022) and (Mandi et al., 2022, 2023)) involves a separate ground-truth predictive mapping between the features and each decision variable’s cost coefficient. In other words, the mapping from the features to one decision variable’s cost coefficient is different from the mapping from the features to another decision variable’s cost coefficient. The implication of this is that, when a new decision variable appears at test time, no informed prediction can be made for this decision variable’s cost coefficient. So, in order to investigate generalizing over larger tasks, another experimental setup is needed, in which the mapping from features to cost coefficients is shared between all decision variables.

Still, this would be interesting to explore in future work, since a setting in which the predictive mapping is shared across decision variables still reflects many realistic scenarios. Additionally, this kind of generalization is promising from a scalability point of view. As we discuss in our future work section, the ability to effectively train on small tasks, in order to later do inference on larger tasks, can make DFL training significantly faster.

Reviewer comment (Experiments):
The problem sizes are relatively small: item size 40 in the knapsack problem, and the 10 customers and 5 facilities in the facility location problem.

Author response:
**Small problem sizes in experiments**

While it is true that the optimization problems used in the experiments are relatively small, this is the case for most existing DFL work as well (Mandi & Guns, 2020; Mulamba et al., 2021; Elmachtoub & Grigas, 2022; Mandi et al. 2022, 2023; Tang & Khalil, 2022a, Pogancic et al., 2022). The reason for this is that training a model in a decision-focused manner is highly expensive, because for each instance in each epoch, a (potentially NP-hard) optimization problem must be solved. In research papers these problems are typically solved to optimality, restricting experimental settings to ones with relatively small optimization problems. In real-world applications with larger problems, approximate solving may be used instead.

Reviewer comment (Novelty):
The motivation for task generalization is not clear — why do we care about task generalization in DFL settings? If we have a new task, why don't we start training from scratch? Task generalization in prediction tasks makes more sense, e.g., when you have zero-shot or few-shot data, but here in DFL we don't have such issues.

Author response:
**Motivation for task generalization**

We, like the other reviewers, do not share the concern that the need for task generalization is insufficiently motivated in the paper, but we are open to hear what the reviewer specifically feels is missing. For reference, the central motivation given in the paper is the following:

> In this work, we identify and address an open problem that has thus far received much less exposition: the limited ability of decision-focused models to generalize across new optimization tasks. Existing DFL methods lead to models that perform well on the optimization problem that they have been trained on, but that generalize poorly to changing tasks, i.e., changing constraints. This is a serious limitation of the state of the art, as real-world optimization problems change all the time. The available resources in scheduling problems, the traversable edges in graph problems and the customer demands in vehicle routing problems, are all subject to constant change.

For the sake of clarity, we want to expand on this here with an example. Consider the following scenario. A delivery company must deliver goods to its customers using delivery trucks. Each day, the company needs to make routing decisions for its trucks, based on the traffic conditions of that day, which are unknown at decision-making time. In order to make routing decisions, a decision-focused predictive model is first used to predict the traffic conditions from known correlated features. In this scenario, it is of vital importance that the predictive model can generalize across decision-making tasks, because each day, there may be different customer demands (both in terms of quantity and location), a different set of inaccessible roads, and a different number of available trucks. Similarly, many other optimization problems that must be solved on a regular basis can be expected to change in unpredictable ways. If the model cannot account for this, and is overfitted on the task it has been trained on, this can lead to highly suboptimal decision-making. It is for this reason that task generalization is of such high importance.

Unfortunately, retraining a model each time the task changes is not a feasible solution to the task generalization problem. This is the case because training models using DFL is computationally highly expensive; much more so than training using traditional prediction-focused methods, as mentioned in the paper. The reason for this is that, in every training iteration, a (potentially NP-hard) optimization problem must be solved. This makes retraining frequently infeasible. Rather, a robust model that can adapt to changing tasks is required, which is what we investigated in this work.

We want to highlight the lack of existing investigation into this topic. Task generalization is a dimension that has thus far received no exposition in the literature on DFL, but that is of vital importance to the broad application of DFL to real-world decision-making.

Reviewer comment (Theory):
The two methods of task generalization are mostly heuristic, and there is no guarantee or in-depth analysis of whether/why they can work. Task representation using GNNs is not new, either.

Author response:
**Lack of theoretical guarantees**

In this work, we exposed the task generalization issue in DFL, and empirically evaluated two ways of dealing with this issue, with great experimental results. A theoretical analysis of the task generalization issue, and to which degree these methods can be expected to reduce it, would certainly be valuable, but is deliberately kept outside of the scope of this conference paper. To add a meaningful theoretical analysis on top of our empirical evaluation would be hard given the 9-page limit, but would be a great addition to a potential journal extension to this work. Do note, however, that the derivation of theoretical guarantees is rarely present in existing DFL work. A deeper theoretical understanding is certainly important, but presently absent from most related work as well.

Reviewer comment (Experiments):
It is unclear how many sampled tasks are needed to train. Beyond reproduction, this is also very important to understand if the comparison is fair, and to understand what computational overhead is required.

Author response:
**Missing descriptions of number of sampled tasks**

As Algorithm 1 shows, a new task is sampled for each training instance in each epoch, and as detailed in section 4.1, the training set contains 800 instances, on which each model is trained for 100 epochs. This means that each model is trained on a total of 800 * 100 = 80.000 tasks.

Note, also, that training on an arbitrary number of tasks is always possible, and does not require additional supervision from the user: the ground-truth solution $z^\star_t(c)$ of a task $t$ can always be computed, as long as ground-truth cost vector $c$ is known. The baseline SPO method and the generalization methods all use the same training sets of 800 $(x, c)$ pairs, and thus get exactly the same information to work with. They are also trained for the same amount of epochs. Thus, the comparison is fair, regardless of the amount of tasks the generalization methods sample at training time. Also note that the only added computational cost of the generalization methods is the computation of the ground-truth solution of a newly sampled task (line 5 of Algorithm 1).

Reviewer comment (Experiments):
For the base task, the generalized methods are even better than those trained on the specific base task, which is counter-intuitive. It may be because the generalization takes in more training data; if so, this does not appear to be a fair comparison.

Author response:
**Generalizing models performing better on base task**

The generalizing models generally do not perform better on the base task than the models trained specifically on that task. This can be seen in the ‘base regret’ columns of the tables of results, where the SPOImplGen and SPOExplGen models perform worse than the SPO model (aside from the first column in the knapsack capacity generalization experiment, where the SPOImplGen model does slightly better, which can be attributed to randomness). This is true in all experiments, except for the one on the facility location problem, where the model trained with explicit generalization does notably better on the base task than the SPO model trained specifically on that task. This is certainly not due to training on more tasks; after all, the implicit generalization model trains on the same amount of tasks as the explicit generalization model, and performs worse on the base task than the SPO model trained specifically on that task. Rather, it is due to the use of explicit task representations.

This surprising but highly promising result is discussed in the final paragraph of section 4.2. It highlights the advantage of explicit task representation in DFL: through an explicit representation of the optimization problem in the input to the predictive model, the model can more effectively learn the right predictive trade-offs required to achieve a low task loss on the problem at hand, especially for problems with strong dependencies between the decision variables, such as the facility location problem. This is an interesting finding and can serve as a great starting point for future work into the possible uses of explicit task representation in DFL, as we also highlight in our future work section.

Reviewer comment (Evaluation):
The advantage of the explicit generalization is very small compared to the implicit method in Tables 1 and 3. It is not clear how the proposed task representation method works.

Author response:
**The benefit of explicit generalization over implicit generalization**

As argued in the paper, explicit generalization does not always offer a large benefit over implicit generalization. There are indeed cases where the advantage is small or nonexistent, such as in knapsack capacity generalization (Table 1), knapsack subset generalization for small subsets (Table 1), or start and end node generalization in the shortest path (Table 3). There are, however, also many cases where the benefit is large: knapsack weight generalization (up to 64% improvement, Table 1), knapsack item subset generalization for large subsets (up to 40% improvement, Table 1), facility location capacity generalization (up to 61% improvement, Table 2) and facility location demand generalization (up to 51% improvement, Table 2).

Reviewer comment (Presentation):
The references are messy. For example, on page 3, Niepart et al. 2021 is mentioned twice under two different ways of handling integer variables.

Author response:
**Double citation**

Thank you for pointing out the double citation of Niepert et al. (2021). The second reference was supposed to be a reference to a different paper. We have fixed this in the revision.

Reviewer comment (Novelty):
The paper does not have a related work section, and many important DFL papers are not discussed, not even the two earlier and classic DFL papers (where the majority of the ideas discussed in the Introduction are from): Donti et al. 2017, "Task-based end-to-end model learning in stochastic optimization" and Amos et al. 2017, "Optnet: Differentiable optimization as a layer in neural networks".

Author response:
**Lack of a related work section and works that are not discussed**

This work indeed does not have an explicit related work section. However, it does discuss the relevant related work in section 2, which is dedicated to giving the needed background on DFL. We disagree with the statement that many important DFL papers are not discussed. The three major paradigms in solving the zero-gradient issue in DFL are identified and accompanied by several citations. For a more comprehensive overview, the reader is then referred to a survey paper on DFL. Finally, by far the most strongly related work, by Tang & Khalil (2022a), is discussed elaborately. We would be interested to hear which important papers the reviewer feels are missing.

The Amos et al. (2017) and Donti et al. (2017) papers are indeed not discussed. Our motivation for this had to do with the scope of what we considered related work: we limited ourselves to the many works that operate on the same kinds of problems as our work: LPs and ILPs. The Amos et al. (2017) and Donti et al. (2017) papers are instead concerned with continuous nonlinear problems, which come with a separate set of challenges, and which do not suffer from the zero-gradient issue central to most DFL research. For this reason, we did not discuss them. However, because of their seminal status in the literature regarding continuous problems, we added citations to them in the revision.

Reviewer comment (Reproducibility):
In the GNN training, it is unclear if it is trained together with the decision loss and the linear prediction layer; if so, it introduces a considerable amount of new parameters and may make the training even more challenging.

Author response:
**GNN training**

The GNN is indeed trained end-to-end alongside the linear prediction layer using the task loss, as described at the end of page 5, and illustrated in Figure 1. It is also correct that the use of a GNN introduces additional parameters to be trained. However, we did not observe this making training more challenging. The explicit generalization models (using GNNs) were able to converge in around the same amount of epochs as the other models (not using GNNs) and did not have noticeably less stable learning curves.

Reviewer comment (Novelty):
There has been recent advance in the area of parametric surrogate learning [1,2]. Authors in those papers also capitalize on generalization property of the learned surrogates. I'd appreciate authors' discussion and analysis on this matter, possibly including some comparison.

Author response:
**Relation to recent advancements in parametric surrogate learning**

Indeed, works [1] and [2] propose ways of learning parametric surrogate losses which can then be used in a DFL loop. The authors of these works indeed characterize their methods as generalizable, but this refers to a very different kind of generalizability than the one we refer to in the present work. [1] and [2] are generalizable in the sense that they do not involve the use of hand-crafted surrogates specific to the optimization problem at hand. Instead, they only require access to a black-box solver, and thus do not exploit the structural properties of the optimization problem. In other words, they are called generalizable because they can be applied to all kinds of optimization problems. For example, they can be used to learn a surrogate for a shortest path problem, or for a knapsack problem, or for a portfolio optimization problem, etc. The techniques are generalizable in the sense that they are broadly applicable. But these works are not about how well a learned surrogate generalizes to altered optimization tasks, which is a very different kind of generalizability, and which is exactly what we refer to with ‘task generalization’ in the present work.

Reviewer comment (Experiments):
What kind of specific data augmentation was performed for the first mechanism? Is it just randomly drawing A,b or some specific data augmentation techniques were applied? If yes, what was their effect?

Author response:
**The kind of data augmentation**

Both mechanisms (implicit and explicit generalization) sample tasks from a task distribution. The task distributions differ between the experiments. For example, in the weight generalization experiment for the knapsack problem, the different tasks have different item weights, sampled from a normal distribution (which corresponds to changes in the A matrix of the ILP formulation). The other 5 experiments all involve other task alterations. But each of these indeed corresponds to a random sampling of alterations to either A or b from some distribution, where the shape and width of this distribution differ between the experiments and between the column groups in the tables of results. For the exact specification of the task alterations for each of the 6 experiments, please see section 4.2 in the paper.

Reviewer comment (Novelty):
The experimental setup in section 3 is a straightforward generalization of existing techniques, such as Tang & Khalil (2022a). The beginning of section 3 states that they differ from Tang & Khalil (2022a) because they consider a distribution over tasks rather than a fixed set of tasks, but this difference just makes equation 5 an expectation over tasks rather than a sum.

Author response:
**Relation to Tang & Khalil (2022a)**

The Tang & Khalil (2022a) paper addresses multi-task DFL where the set of tasks is fully known and kept the same in the training and test set. They investigate two approaches. The first is a single-cost approach, where the model outputs a single cost vector that is subsequently used in all of the tasks. This model is trained using an aggregate of the task losses for the different tasks. The second is a multi-cost approach, where the model has some shared layers, followed by a multi-headed layer, with a separate head for each task, each of which outputs a cost vector specific to the associated task. The shared layers are trained on an aggregate of the task losses for all tasks; the heads are trained specifically on the task loss for their associated task.

There is a similarity between the implicit generalization training scheme and the single-cost methodology used in Tang & Khalil (2022a). Both are simple schemes methodologically, and involve training the predictive model on multiple tasks, rather than on only one. So, we argue that the single-cost contribution in Tang & Khalil (2022a) and our contribution with implicit generalization are not really of a technical nature. The contributions instead lie mostly in the insights gained, which are clearly different in the two works.

The approach in Tang & Khalil (2022a) is aimed at and evaluated in terms of multi-task performance, meaning that the model is trained and evaluated on a fixed and fully-known set of tasks. That is, the evaluation is done entirely in-sample: the tasks used at test time are the same as the ones used at training time. Tang & Khalil’s main finding was that when predictions must be used in several related tasks that are known in advance, it often suffices to train one predictive model on all of these tasks together, rather than training a separate predictive model for each task. 

In our paper, the focus is task generalization, i.e., out-of-sample performance, on tasks that have not been seen before during training. So, rather than working with a fixed set of tasks, implicit generalization involves sampling tasks from the task distribution at training time. We found that models trained implicitly in a decision-focused way on a distribution of tasks perform much better across that distribution than models trained on individual tasks from that distribution. Additionally, we found something that runs somewhat counter to the finding in Tang & Khalil (2022a): while the models trained with implicit generalization perform well across the task distribution, they perform worse on any individual task than a model trained specifically for that task. This was the motivation for explicit task representations, which is where our work does offer a contribution that is of a more technical nature.

It is somewhat unclear to us what it would mean to compare the way Tang & Khalil (2022a) train their multi-task model with the way we train our implicit generalization model. In a sense, implicit generalization is the application of their single-cost approach to the task generalization setting.

Regarding the multi-cost approach of Tang & Khalil (2022a): this approach is not applicable to our experimental setting. Each head of the model is trained on a specific task, so when an out-of-sample task occurs at test time, it is unclear from which head to take the predicted cost vector. On the other hand, our training scheme with explicit representations could be applied to the experimental setting from Tang & Khalil (2022a). We chose not to do so because this lies outside of the scope of our work: our research questions relate to task generalization (i.e., out-of-sample performance), whereas the experimental setting from Tang & Khalil (2022a) is only suited to answer questions about in-sample performance. Additionally, we do not expect the explicit generalization model to generally outperform the multi-cost approach from Tang & Khalil (2022a) in their in-sample setting, since the latter has a dedicated head for each task, which is unlikely to be outperformed by a model that is trained with respect to a whole distribution of tasks (although our experiments on the facility location problem do suggest that the opposite is possible, which is a finding that must be further explored in future work).

Reviewer comment (Experiments):
Because the methodology is similar to Tang & Khalil (2022a), it would have been insightful to compare directly to a) their method on the experimental settings here, and/or b) the SPOImplGen and SPOExplGen methods on the experimental settings there. If possible, this could add an insightful bridge between the existing works, and if not possible it could be insightful to discuss more why not. (I acknowledge comparisons to their models/settings may not make sense if they assume more contextual information is available.)

Author response:
**Relation to Tang & Khalil (2022a)**

The Tang & Khalil (2022a) paper addresses multi-task DFL where the set of tasks is fully known and kept the same in the training and test set. They investigate two approaches. The first is a single-cost approach, where the model outputs a single cost vector that is subsequently used in all of the tasks. This model is trained using an aggregate of the task losses for the different tasks. The second is a multi-cost approach, where the model has some shared layers, followed by a multi-headed layer, with a separate head for each task, each of which outputs a cost vector specific to the associated task. The shared layers are trained on an aggregate of the task losses for all tasks; the heads are trained specifically on the task loss for their associated task.

There is a similarity between the implicit generalization training scheme and the single-cost methodology used in Tang & Khalil (2022a). Both are simple schemes methodologically, and involve training the predictive model on multiple tasks, rather than on only one. So, we argue that the single-cost contribution in Tang & Khalil (2022a) and our contribution with implicit generalization are not really of a technical nature. The contributions instead lie mostly in the insights gained, which are clearly different in the two works.

The approach in Tang & Khalil (2022a) is aimed at and evaluated in terms of multi-task performance, meaning that the model is trained and evaluated on a fixed and fully-known set of tasks. That is, the evaluation is done entirely in-sample: the tasks used at test time are the same as the ones used at training time. Tang & Khalil’s main finding was that when predictions must be used in several related tasks that are known in advance, it often suffices to train one predictive model on all of these tasks together, rather than training a separate predictive model for each task. 

In our paper, the focus is task generalization, i.e., out-of-sample performance, on tasks that have not been seen before during training. So, rather than working with a fixed set of tasks, implicit generalization involves sampling tasks from the task distribution at training time. We found that models trained implicitly in a decision-focused way on a distribution of tasks perform much better across that distribution than models trained on individual tasks from that distribution. Additionally, we found something that runs somewhat counter to the finding in Tang & Khalil (2022a): while the models trained with implicit generalization perform well across the task distribution, they perform worse on any individual task than a model trained specifically for that task. This was the motivation for explicit task representations, which is where our work does offer a contribution that is of a more technical nature.

It is somewhat unclear to us what it would mean to compare the way Tang & Khalil (2022a) train their multi-task model with the way we train our implicit generalization model. In a sense, implicit generalization is the application of their single-cost approach to the task generalization setting.

Regarding the multi-cost approach of Tang & Khalil (2022a): this approach is not applicable to our experimental setting. Each head of the model is trained on a specific task, so when an out-of-sample task occurs at test time, it is unclear from which head to take the predicted cost vector. On the other hand, our training scheme with explicit representations could be applied to the experimental setting from Tang & Khalil (2022a). We chose not to do so because this lies outside of the scope of our work: our research questions relate to task generalization (i.e., out-of-sample performance), whereas the experimental setting from Tang & Khalil (2022a) is only suited to answer questions about in-sample performance. Additionally, we do not expect the explicit generalization model to generally outperform the multi-cost approach from Tang & Khalil (2022a) in their in-sample setting, since the latter has a dedicated head for each task, which is unlikely to be outperformed by a model that is trained with respect to a whole distribution of tasks (although our experiments on the facility location problem do suggest that the opposite is possible, which is a finding that must be further explored in future work).

Reviewer comment (Novelty):
While the experimental results through section 4 are scientifically well-executed and documented, I do not find them especially insightful or surprising. The problems are relatively standard for DFL and I would not expect methods to generalize between tasks. Between models, it is also unsurprising the SPOExplGen model conditional on the task information performs the best as it has the most information.

Author response:
**Summary of insights gained**

We would summarize the most important new insights worth spreading to the community as follows.

1) While the application of DFL techniques generally leads to predictions with better performance on the downstream optimization task than predictions that are made to be as accurate as possible, these techniques strongly overfit the model to the optimization task at hand. This means that when the downstream task is altered at test time, the downstream task loss quickly worsens, sometimes even to the point where the decision-focused model becomes ‘counterproductive’, and where the use of a much simpler prediction-focused model leads to better downstream decisions.

2) To combat this, one may instead implicitly train on a distribution of tasks. This greatly improves the out-of-sample generalization performance, making the model more robust to task alterations that can occur after training. However, models trained implicitly on a distribution of tasks perform worse on any particular individual task from that distribution than a model trained specifically for that task, and this performance gap grows as the task distribution widens.

3) To improve on this, especially when the task alterations occur on the left-hand side of constraints, one can use our proposed explicit generalization methodology.

4) Surprisingly, on our most complex benchmark, the model trained with explicit generalization (with explicit task representations) even performed better on individual tasks from the task distribution than models trained specifically on those tasks (without explicit task representations). This means that the merit of task representations is not merely to limit the individual task performance loss incurred by training on many tasks, but that they even have the potential to improve performance on individual tasks. This finding is further discussed in the last paragraph of section 4.3, and seems to us like a great starting point for future work.