Transformation-Invariant Learning and Theoretical Guarantees for OOD Generalization

We thank you for your detailed review. We address your comments below.

When $H$ is unknown, the authors give a similar result in the realizable case using data augmentation, and a weaker variant for the agnostic setting for finite which ensures low error non-uniformly (that is, for each fixed map T, the strategy will do well, but this may not be the case for all T simultaneously as above) across a probabilistic mixture of hypotheses from H via an MWU type strategy.

We want to clarify and emphasize that in Theorem 3.1, we learn a mixture strategy (a randomized predictor) with a low error guarantee that holds simultaneously across all T (Exactly the same guarantee in Theorem 2.1).

The lower bound based on VC(H \circ T) is weaker than what one typically hopes for in a combinatorial characterization. Namely while the authors show there exist classes for which this bound is tight, a strong lower bound (as holds for standard learning) would show this is tight for every class. Is there a parameter characterizing the model in this stronger sense? The lower bound also only holds for proper learners.

This is a great point, and it is an interesting open question for future work to have a full characterization of what is learnable in our model in terms of $\mathcal{H}$ and $\mathcal{T}$.

As mentioned in Lines 117–124, we think that improper learners could achieve better sample complexity but this remains an open question. One of the main reasons for including the lower bound (Theorem 2.2), is to argue that the sample complexity of the simple proper learner in Equation 2 is tight in general, and we can not hope to improve the sample complexity of this proper learner further. While there could be complex improper learners with better sample complexity, we believe that having a proper learner that has a simple description (similar to the ERM principle) is valuable even when it does not achieve the optimal sample complexity.

In the setting where H is unknown, the use of data augmentation is a bit troubling in the sense that the size of the augmented sample could be extremely large (indeed infinite for infinite T). This isn’t captured by the model the authors introduce which simply assumes the ERM oracle can handle a call of any size, but this does not really seem to capture the computational cost of learning even at an oracle level.

This is a great point. We argue that for the purpose of achieving the guarantee in Equation (1) in the realizable case, it is necessary to call ERM on the full augmented sample. Consider an example of a distribution supported on a single point $(x,-)$ on the Real line where $x = 5$, consider transformations of $x$: $z_i = x+i$ for all $i\geq 1$ induced by some class $\mathcal{T}$. Now, calling ERM on a finite subset of these transformations will return a predictor that labels $x, z_1, …, z_{k}$ with a label $-$, and labels $z_{k+1}, \dots$ with +, which fails to satisfy our objective.

Perhaps we should clarify that we are considering finite transformations (similar to the agnostic case). In which case, the computational cost of a single ERM call on $m*|T|$ samples, grows by a factor of $|T|$. But, with a fairly standard boosting argument, we can also run a classical boosting algorithm that makes multiple calls to ERM. Specifically, we run boosting starting with a uniform distribution on the augmented sample. In this case, each ERM call is on samples of size $O({\rm vc}(\mathcal{H}))$, and we make at most $O(\log(m*|T|))$ many ERM calls. So, we have a tradeoff between the size of a dataset given to ERM on a single call, and the total number of calls to ERM.

Furthermore, the agnostic guarantees in this setting only hold for finite $T$, and hold in the seemingly much weaker non-uniform way mentioned above. It is fairly unclear if either of these constraints are actually necessary.

It is an interesting direction to extend the result in Theorem 3.1 to classes with infinite $\mathcal{T}$, perhaps with a certain parametric structure.

We thank you for your detailed review. We clarify below how our work is different from multi-distribution learning, and how the lower bounds on multi-distribution learning do not contradict our upper bound result in Theorem 3.1 / Remark 3.2.

Could you explain the main differences between this model and that of multi-distribution learning? In particular, are there some natural instances of multi-distribution learning which do not fall into this model?

We summarize the main differences below:

In our model, the learning algorithm has access to samples from a source distribution $D$ and access to transformations $T_1, …, T_k$. Thus, by drawing a single sample $(x,y)$ from $D$, the algorithm can generate $k$ samples: $(T_1(x), y), …, (T_k(x), y)$ from $T_1(D), …, T_k(D)$. And, we only count the number of samples drawn from $D$.

In contrast, in multi-distribution learning, the learning algorithm has access to samples from $k$ distributions $D_1, …, D_k$ and it does not know the relationship between $D_1, …, D_k$. And, what is counted is the number of samples drawn in total from all $k$ distributions. While, as you noted and based on line 41, it is possible in principle to find a distribution $D$ and $k$ transformations $T_1, …, T_k$ such that $D_i = T_i(D)$. But, that would require the algorithm to learn the transformations $T_1, …, T_k$ based on samples from $D_1, …, D_k$ which is going to be more costly than solving the original multi-distribution problem directly.

In our model, we only consider a covariate shift setting, i.e., the labels do not change under transformations (see Line 29). While multi-distribution learning is more general and allows different labeling functions across the $k$ distributions.

Thus, given the above, we can not in general reduce a multi-distribution learning problem instance to a problem instance in our model.

As described above under weaknesses, why do the lower bounds for multi-distribution learning not apply to the results of Theorem 3.1? Have I misunderstood something about the model in this paper, or perhaps the lower bound instances from [1]?

The lower bounds from [1] do not apply to Theorem 3.1, because in the setting of this theorem, the learning algorithm is provided access to transformations $T_1, …, T_k$. And, given a dataset $S$ of $m$ samples from $D$ “source”, the learner can generate $S_1 = T_1(S), …, S_k = T_k(S)$ a total of $mk$ samples by applying the transformations $T_1, …, T_k$. We only count the number of samples drawn from $D$ which is $m$.

While the lower bound instance from [1] (e.g. Theorem 4.2) involves $k$ distributions $D_1, …, D_k$ with disjoint supports and the learner does not know the transformations $T_1, …, T_k$ nor a source $D$ satisfying $T_i(D) = D_i$. The cost of learning such transformations $T_1, …, T_k$ based solely on samples from $D_1, …, D_k$ will be at least $k$ (i.e., drawing at least one sample from each distribution). In general, learning transformations is more costly (and likely an overkill) to achieve the goal of multi-distribution learning.

Furthermore, as pointed out on line 41, if the domain of the learning problem is then Brenier's Theorem implies that for any instance of a multi-distribution learning problem with sufficiently "nice" distributions, there is a distribution and set of transformation such that . Thus, in the case of finite on domain these two models are equivalent so long as only "nice" distributions are considered. The setting of finite is particularly important, as it is the setting for which the main results in this paper give oracle-efficient algorithms (given an ERM oracle for ).

As alluded to above, another challenge here is: because our model only considers covariate shifts (i.e., the labeling doesn't change across transformations), in order to reduce multi-distribution learning to our model, the transformations that need to be learned have to be "label preserving". This is another challenge that will likely make it even more difficult to reduce a problem instance of multi-distribution learning to our model.

We hope that this clarifies the differences between multi-distribution learning and our model. We are happy to answer any further questions, and would appreciate an updated response from you.

Thank you for the follow-up response. We address your comments below.

Thank you for your response, it definitely clarifies the relationship to multi-distribution learning, and confirms that everything is fine with Theorem 3.2. To make sure I understand correctly: the main difference is that the transformations $T_i$ are given to the learning algorithm in advance, and the learner only pays for one sample from $D$, from which it is able to generate $k$ transformed samples from $T_1(D), \dots, T_k(D)$. In contrast, in multi-distribution learning the learner pays for $k$ samples in order to see one sample from each of $D_1,\dots, D_k$.

Yes, we want to confirm that your understanding above is correct. We totally agree that this distinction is crucial and we will revise the comparison on line 313 to clarify that the reduced sample complexity is due to the additional structure of knowing the $k$ transformations, rather than an inherent improvement over multi-distribution learning. Specifically, we plan to include the following revision:

Multi-Distribution Learning. This line of work focuses on the setting where there are $k$ arbitrary distributions $D_1, \dots, D_k$ to be learned uniformly well, where sample access to each distribution $D_i$ is provided [see e.g., 26]. In contrast, our setting involves access to a single distribution $D$ and transformations $\{T_1, \dots, T_k\}$, that together describe the target distributions: $\{T_1(D), \dots, T_k(D)\}$. From a sample complexity standpoint, multi-distribution learning requires sample complexity scaling linearly in $k$ while in our case it is possible to learn with sample complexity scaling logarithmically in $k$ (see Theorem 3.1 and Remark 3.2). The lower sample complexity in our approach is primarily due to the assumption that the transformations $\{T_1, \dots, T_k\}$ are known in advance, allowing the learner to generate $k$ samples from a single draw of $D$. In contrast, in multi-distribution learning, the learner pays for $k$ samples in order to see one sample from each of $D_1,\dots, D_k$. Therefore, while the sample complexity is lower in our setting, this advantage arises from the additional information/structure provided rather than an inherent improvement over the more general setting of multi-distribution learning. From an algorithmic standpoint, our reduction algorithms employ similar techniques based on regret minimization and solving zero-sum games [24].

Thank you again for your insightful feedback, which has significantly improved the clarity and accuracy of our paper.

We thank you for your detailed review. We address your comments below.

We emphasize that our contributions are conceptual and theoretical. At the conceptual-level, to the best of our knowledge, this is the first work to model the OOD/distribution shift problem through a transformation function class, which is relevant to many real-world scenarios as discussed in the introduction. As is well known, studying OOD/distribution shift problems requires specific modeling of the relationship between source and target distributions (e.g., bounded density ratios between source and target distributions, which often don't hold in real-world scenarios). Instead, we propose a new way of modeling the problem through the lens of transformations. This allows us to obtain new learning guarantees against classes of distribution shifts, and while these may not require new sophisticated tools, they nonetheless provide bounds that were not covered by prior work.

The theoretical results appear to be direct applications of existing uniform bounds on empirical processes and VC dimension theory. The introduction of transform-invariance does not significantly impede the application of these uniform bounds, as they represent worst-case bounds for functions to be optimized. Specifically, Theorem 2.1 is derived by applying the original VC theory (by Vapnik and Chervonenkis) to the composite class. Theorem 4.1 is a straightforward consequence of the optimistic rate by Cortes et al. for the composite class. Theorem 5.1 is also obtained by directly applying the original VC theory to the composite class. Consequently, the technical novelty underlying the theoretical results is not clearly evident.

We believe that it is an asset to be able to leverage the structure of the composition of hypotheses and transformations, and to obtain sample complexity bounds by utilizing existing VC theory (as also highlighted by reviewers RtrA and fSmB), to obtain new insights for this relevant OOD setting. We further demonstrate this by providing in Section 2.1 several examples highlighting how a simple bound can imply several interesting results and guarantees that are not known before.

The authors do not sufficiently elucidate the implications of their experimental results. The significance and benefits of the empirical evaluation are not adequately explained, leaving the reader unclear about the practical impact of the proposed approach.

We will add this elaboration to the paper. Please see below for further clarification.

Can you provide more insight into the practical implications of your empirical evaluations? What specific benefits does your transform-invariant learning algorithm offer in real-world scenarios?

We only provide a simple illustrative experiment showing how the use of transformations can help learn with fewer samples / training data. Specifically, in Figure 1, both learning tasks can be learned better (i.e., higher accuracy) with the use of transformations compared with the baseline of standard training (with no transformations).

Investigating implications of our theoretical results in real-world scenarios is of course an interesting and important future direction, but it lies beyond the scope of this work.

Thank you for your detailed review. We address one your comments below.

One could argue that one of the most important use cases is when $\mathcal{T}$ form a group of symmetry transformations. As discussed on page 9, in this setting the technique of data augmentation was studied previously.

As mentioned in page 9 (Lines 333–338), we emphasize that the technique of data augmentation does not enjoy the worst-case error guarantees across transformations as we prove in this work (using a different learning rule).

Limitations are discussed adequately. However, I believe it would be helpful to additionally have a dedicated short section titled limitations that briefly summarizes the limitations.

As suggested, we will reiterate the limitations together in a single short section and discuss future open questions.

Thank you for your response. I updated the text of my review accordingly and I will finalize my rating once the rebuttal discussion with other reviewers concludes. For now, my rating remains 6: Weak Accept.

Thank you for your detailed review. We address your comments below.

When proposing a new learning rule, it is better to demonstrate its utility by comparing it to existing and similar ones, such as ERM (at least) and DRO (similar).

ERM can be viewed as concerned only with the identity transformation map $T(x) = x$. Therefore, when interested in a collection of transformations e.g. rotations, then ERM can be provably shown to do worse than our proposed learning rule in Equation 2. Here is an Example/Scenario that further illustrates this difference: when the training set consists of images of upright cats, ERM might learn the feature of "upright," leading the model to fail when the test set is composed of upside-down cat images. Our method, however, would still work in this scenario.

We discuss the relationship with the DRO line of work in (Lines 306–312).

As the authors mentioned in Related Work, IRM is another popular learning framework for OOD generalization, but some works concluded that ERM may outperform IRM in some cases. Some other works showed that ERM may outperform DRO. Will the learning framework proposed in this paper be underperformed by ERM under some scenarios?

From a theoretical standpoint, ERM and our proposed learning rule (e.g., Equation 2) enjoy different guarantees. When we are only interested in performing well on a single source distribution, then ERM suffices. But, when interested in a collection of distributions parametrized by transformations $\mathcal{T}$, our learning rule outperforms ERM.

Consider a more realistic definition of $\mathcal{T}$: $T(x)=x+R(x)$, as in the ResNet, because Q and P should not be too far away. In such a case, what will the learning rule, the theoretical guarantee, and the algorithm be like?

In this case, the learning rule/algorithm corresponds to solving an alternating minimization-maximization problem. The minimization is with respect to parameters of the model responsible for classification (e.g. an MLP on top of the ResNet architecture) and the maximization is with respect to transformations (e.g. the parameters of the $R$ function you describe). The theoretical guarantee is that with enough samples (scaling with VC dimension of ResNet), the learned model will minimize worst-case error (across all $R$ functions).

Also, please see Lines 144–155 for discussion of a similar example concerned with feed-forward neural networks.

What's the connection between this work and the importance weighting for distributional shift (such as [1])?

Importance weighting addresses the problem of covariate shift / distribution shift from a different angle relying on bounded density ratios between source and target distributions. Our work instead considers transformations, i.e., when $P$ and $Q$ are related by a transport map such that $Q = T(P)$. This allows us to address classes of distribution shifts that may be difficult to address with importance weighting, specifically when the density ratio is not bounded. I.e., when distributions $P$ (source) and $Q$ (target) have disjoint supports, the density ratio is not bounded, and therefore we can not rely on importance weighting. However, it may still be possible to learn using the language of transformations that we contribute in this work.