Transformation-Invariant Learning and Theoretical Guarantees for OOD Generalization

Summary

The authors study a new framework for learning under distribution shift caused by a family of `invariance maps’ $\mathscr{T}: X \to X$. Informally, one may think of this as generalizing e.g. settings such as learning under invariance to rotation, where the maps may be rotation matrices and the learner should perform well at test time regardless of whether the raw data has been rotated. More formally, the learner is given a hypothesis class $H$, a family of maps $\mathscr{T}$, and sample access to a distribution $D$ over $X \times {0,1}$, and wishes to train a classifier that minimizes the test error over all possible shifts $T(D)$ for $T \in \mathscr{T}$, either globally in the sense of outputting $h$ satisfying $$sup_T \ \{ err(h,T(D))\} \leq OPT + \varepsilon$$ Where $OPT=min_{h_* \in H} sup_T err(h^*,T(D))$, or where error is measured with respect to the best possible over each individual T (the authors call this regret minimizing): $$sup_T \ (err(h,T(D))-OPT_T) \leq min_{h^* \in H} \ (sup_T \ (err(h^*,T(D))-OPT_T)) + \varepsilon$$ This model captures a number of settings considered in the literature. In particular, under reasonable assumptions on distributions $D,D’$ over $X$, there is always a map $T$ such that $D=T(D’)$ (Breiman’s theorem), so the setup does not lose substantial generality over classical study of covariate shift. It may also model transformation invariances as mentioned above, as well as certain types of adversarial attacks (where the adversarial shift is chosen before the adversary sees the test example). The main results concern the sample complexity of learning in this framework in the case were 1) H is known to the learner 2) H is unknown and the learner only has ERM access to H. In the first setting, the authors show the sample complexity of (proper) agnostic PAC learning in this model is controlled by the VC dimension of the composed class VC(H \circ T). In other words, it takes at least $VC(H \circ T)$ samples to properly learn the class, and for any $k$ there exists a VC 1 class H and maps T with $VC(H \circ T)=k$ which requires $\Omega(k)$ samples to learn. To complement their result, the authors also investigate $VC(H \circ T)$ for a few special cases, including when $H$ is a composition $f \circ W$ of an arbitrary linear map W and any fixed $f: X \to \{0,1\}$ and $T$ is a family of linear maps (in which case $VC(H \circ T) \leq VC(H)$), some basic non-linear $T$ coming from neural net architectures, and transformations on the hypercube. 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 $\mathscr{T}$ which ensures low error across a probabilistic mixture of hypotheses from H via an MWU type strategy. Finally the authors also consider a related setting in which the learner is in some sense allowed to ignore transforms with bad error, and the learner tries to return a hypothesis with low error across as many of the transforms $T \in \mathscr{T}$ as possible. At a technical level, most of the results are the consequence of the basic duality of transforming input data and transforming the hypothesis class, allowing the authors to move between $T(D)$ and $H \circ T$.

Strengths

The authors introduce a clean theoretical framework for various types of distribution shift and transformation invariance studied in the literature, capturing several interesting problems in the literature. The proposed framework should be of interest to the theory audience at NeurIPS, and is particularly nice in that it allows easy application of classical methods in VC theory to the problem. The authors give a (weak) characterization of its statistical complexity in terms VC(H \circ T), a natural combinatorial parameter of the concept and transformation class. The authors discuss several natural scenarios in which their upper bound gives new sample complexity bounds on transformation invariant learning, including settings occurring in modern neural net architectures (namely the composition of linear transformations with activation functions). The authors study several natural variants of the framework (global risk, regret, maximizing #transformations with low error) with similar corresponding upper bounds.

Weaknesses

On the other hand, beyond the main upper bound and examples given by the authors, many of the further results feel somewhat incomplete. For instance: 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. 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. Furthermore, the agnostic guarantees in this setting only hold for finite $T$. POST REBUTTAL: The authors have included some discussion re: the first point that helps ameliorate the issue (e.g. using boosting one can lower the batch size to ERM at the cost of log(T) calls). Combined with the fact that the authors main techniques are fairly standard (note in some cases, as mentioned in the strengths, it is an *advantage* of the framework one can use standard techniques! But this is less the case when these techniques do not seem to give the desired results), these issues prevent me from a stronger recommendation of the paper.

Questions

See above

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes

Summary

This paper studies a theoretical model of out-of-distribution (OOD) generalization, where the possible distribution shifts are encoded by a collection $\mathcal{T}$ of transformations on the domain. The goal in this setting is given a distribution $D$ to learn a hypothesis $h$ from some hypothesis class $\mathcal{H}$, such that $h$ achieves the best possible risk on the distribution $T(D)$ under the worst-case shift $T\in\mathcal{T}$. The paper gives algorithms achieving this goal in a variety of settings including when $\mathcal{H}$ and $\mathcal{T}$ are both known, and when given an ERM oracle for $\mathcal{H}$ and $\mathcal{T}$ is finite. The bounds are given in terms of the VC-dimension of the composition $\mathcal{H}\circ\mathcal{T},$ consisting of all function of the form $h\circ T$.

Strengths

1. The model introduced of a class of transformations corresponding to distribution shifts gives very mathematically clean results for sample complexity in terms of the VC-dimension of the composed class $\mathcal{H}\circ\mathcal{T}$. 2. The explanations of how the model can be applied to different practical settings involving distribution shift are clear and show broad applicability.

Weaknesses

The primary weakness is in distinguishing these results from prior work, particularly that on multi-distribution learning. In particular, when the class of transformations $\mathcal{T}$ is finite, then every problem instance from this paper's model is also an instance of multi-distribution learning with $D_1,...,D_k$ given by $T_1(D),...,T_k(D)$. Furthermore, as pointed out on line 41, if the domain of the learning problem is $\mathbb{R}^d$ then Brenier's Theorem implies that for any instance of a multi-distribution learning problem $D_1,...,D_k$ with sufficiently "nice" distributions, there is a distribution $D$ and set of $k$ transformation $T_1,\dots,T_k$ such that $D_i = T_i(D)$. Thus, in the case of finite $\mathcal{T}$ on domain $\mathbb{R}^d$ these two models are equivalent so long as only "nice" distributions are considered. The setting of finite $\mathcal{T}$ 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 $\mathcal{H}$). The above discussion leads to what may be an issue with Theorem 3.1 as stated. As far as I can tell, Theorem 3.1 combined with Remark 3.2 would imply a $O(d + \log k)$ upper bound on sample complexity in the case that the VC-dimension of $\mathcal{H}$ is $d$ and $|\mathcal{T}|=k$. However, prior work on multi-distribution learning [1] suggests that in the setting of $k$ distributions $D_1,...,D_k$ and a hypothesis class of VC-dimension $d$ there is a lower-bound on sample complexity of $\Omega(d + k\log k)$. Perhaps I have missed something here, but I don't see any obstruction to instantiating the lower-bound family of instances from [1] as a family of instances in the model of this paper. So there seems to be some discrepancy between the known lower bounds and Theorem 3.1. [1] Haghtalab, Nika, Michael Jordan, and Eric Zhao. "On-demand sampling: Learning optimally from multiple distributions." Advances in Neural Information Processing Systems 35 (2022): 406-419.

Questions

1. 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? 2. 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]?

Rating

4

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes.

Summary

The authors investigate transform-invariant binary classification problems in which the learner observes transformed features during the prediction phase and aims to construct a classifier robust against test-time transformations. Their first contribution is the derivation of a generalization error bound on the worst-case risk of transform-invariant classification, characterized by the VC dimension of the composite class. Subsequently, they propose an algorithm that achieves this generalization bound without direct access to the hypothesis class, relying solely on the empirical risk minimization (ERM) oracle associated with the hypothesis class. The second contribution is the establishment of a generalization error bound on the number of high-risk transformations, which is also characterized by the VC dimension of the composite class. The third result presents a generalization error bound on the worst-case regret, similarly characterized by the composite class's VC dimension. Empirical evaluations were conducted to demonstrate the behavior of the transform-invariant learning algorithm.

Strengths

The paper is well-written and easy to follow. It addresses a crucial topic for machine learning: robustness against transformations. This focus on developing classifiers that remain effective under various input transformations is highly relevant to real-world applications and aligns well with the conference themes. The authors provide generalization error bounds for various accuracy measures, each corresponding to potential practical scenarios. This comprehensive approach enhances the applicability of their work. The methodology to minimize worst-case accuracy solely relying on the ERM oracle is innovative. While it adapts the Multiplicative Weights method to the transform-invariant classification problem, the application of this method to this specific domain represents a novel contribution.

Weaknesses

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. 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.

Questions

- 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?

Rating

4

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors adequately address the limitations and potential impacts.

Summary

The work focuses on theoretical aspects of distribution shift, considering various frameworks that expand the standard PAC model of learning. The paper considers a hypothesis class $\mathcal{H}$ and a class of transformations of the domain $\mathcal{T}$. The learning algorithm is given some data drawn i.i.d. from some distribution $D$ and the overall goal is to produce a classifier that does well not only on fresh samples from the source distribution $D$ but also when a transformation $T$ from $\mathcal{T}$ is applied to $D$. Overall, the makes assumptions on the VC dimension of the class of functions obtained from composition of $\mathcal{H}$ and $\mathcal{T}$. In some parts of the paper it is further assumed that $\mathcal{T}$ is finite. Various parts of the paper consider different loss functions that depend on the data distribution and the class of transformations $\mathcal{T}$. The goal is to find a hypothesis $\hat{h}$ for which the value of a loss function is not more by $\epsilon$ greater than the best (i.e. smallest) value of the loss function among all functions in the hypothesis class $\mathcal{H}$. (Note that this optimum value also depends on the choice of the loss function.) The paper considers a number of specific goals in this broad setup: - In Section 2 the loss of a classifier $\hat{h}$ equals to the worst-case population risk that $\hat{h}$ will have as a result of applying a transformation $T$ from $\mathcal{T}$ to the data distribution. Note that this loss function quantifies the worst-case impact that a transformation from $\mathcal{T}$ can have on the population loss. Section 2 considers the classifier $\hat{h}$ that has the smallest introduces and analyzes the classifier obtained by minimizing the empirical worst-case risk that a hypothesis will have as a result of applying a transformation $T$ from $\mathcal{T}$ to the training dataset. - In Section 3 the class of transformations $\mathcal{T}$ is assumed to be finite and the algorithm has access to an ERM oracle for the hypothesis class $\mathcal{H}$. First, the realizable setting is considered (i.e. the setting in which there is some function in $\mathcal{H}$ that perfectly desribes the data even if it is transformed by an arbitrary transformation in $\mathcal{T}$). The paper shows that the classifier shown to exist in Section(2) can be obtained by (i) augmenting the training dataset by applying various transformations in $\mathcal{T}$ and (ii) applying the ERM oracle to the resulting dataset. - In Section 3 the paper also considers the agnostic setting. In this setting there can potentially be a transformation in $\mathcal{T}$ that causes the empirical risk to be non-zero. The paper gives an efficient algorithm that combines calls to the ERM oracle with the multiplicative weight update method to obtain a classifier shown to exist in section 2. Note that the run-time of the algorithm does not include the time to actually minimize the empirical risk when the ERM oracle is called, which can potentially be slow. - Section 4 considers the case in which there is potentially no classifier whose error stays small when **any** transformation in $\mathcal{T}$ is applied. Instead, Section 4 considers a loss function equal to the **number** of transformations in $\mathcal{T}$ that make the population loss of the classifier greater then epsilon. The paper gives a statistical generalization bound with respect to this loss. - Section 5 again considers a different loss function. The loss of a hypothesis $\mathcal{h}$ is defined as the worst-case difference (over transformation $T$ in $\mathcal{T}$) between (i) the population risk of $\mathcal{h}$ after transformation $T$ is applied (ii) the lowest population risk among all hypotheses in $\mathcal{H}$. In other words, if there is some transformation $T$ makes the error of $\mathcal{h}$ large, we do not penalize the algorithm if the worst-case error in the hypothesis class $\mathcal{H}$ is likewise large. With respect to this loss function as well, the paper again gives generalization bounds (similar to Sections 2 and 4) and an algorithmic reduction (Similar to Section 3). Additionally, an experiment is presented in which a Boolean function is learned from synthetic data, and the function is invariant under a class of permutations.

Strengths

- The question of mitigating distribution shift is central for modern machine learning. - The paper is written clearly and numerous examples are given.

Weaknesses

- The technical results follow by a fairly direct use of standard generalization bounds, together with standard tools from convex optimization such as multiplicative weights. - 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. However, the previous work does not enjoy the worst-case error guarantees across transformations.

Questions

No questions.

Rating

6

Confidence

4

Soundness

4

Presentation

4

Contribution

2

Limitations

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.

Summary

This paper proposed a transformation-based invariant learning framework for OOD generalization. The authors initiated a theoretical study for this framework, investigating learning scenarios where the target class of transformations is either known or unknown. Upper bounds on the sample complexity in terms of the VC dimension of the class composing predictors with transformations were obtained.

Strengths

Nice theoretical results, with various scenarios considered.

Weaknesses

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). 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? 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? What's the connection between this work and the importance weighting for distributional shift (such as [1])? [1] T. Fang, etc. Rethinking Importance Weighting for Deep Learning under Distribution Shift. NeurIPS 2020.

Questions

Please see the questions mentioned above.

Rating

4

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

N.A.

Rebuttal Response

We thank the authors' for their clarifications. The discussion of ERM complexity is enlightening, and I think is worth including in the paper (even if just in discussion). I think I was confused on the quantifier in Theorem 3.1, and the authors are correct; sorry! I will update my review to reflect this.