Awesome

Official code for "Rich Feature Construction for the Optimization-Generalization Dilemma"

Overview

In Machine Learning, defining a generalized goal (e.g. the invariant goal in out-of-distribution generalization) and finding a path to the goal (e.g. the many optimization tricks) are two key problems. Usually, there is a dilemma between the two. i.e. either the generalization goal is weak/poor or the optimization process is hard. This optimization-generalization dilemma is especially obvious in the out-of-distribution area. This work tries to solve this dilemma by creating a RICH and SIMPLE representation, such that the optimization process becomes easier with the representation. As a result, we can pursue a stronger generalization goal.

A short Story

Two common questions, in many areas, are: "where is our goal?" and "how to reach the goal from our current position?". A successful project needs to answer both questions. The two questions, however, are contradicted in difficulty. When the goal is ambiguous, normally the path to the goal is blurry. When the path is clear and confident, normally the goal is plain. For the goal of "making a cup of espresso?", for instance, most people can have a clear precise path immediately. On the other hand, "Building a spacecraft to Jupyter" is an ambiguous goal. But most people have no idea about how to achieve it.

Can we build a spacecraft by purely thinking about the "spacecraft"? No. The spacecraft is built based on the development of diverse areas, such as material, computer, engine.

The story above revises the path to hard problems, that is "Search/develop diverse areas (directions). Then a clear path may appear upon them. Otherwise, continuing search more."

The rule above is also the key idea of the proposed Rich Feature Construction (RFC) method.

required environments

• wilds==2.0.0
• einops=0.4.1
• python=3.6.13
• pytorch=1.10.2
• torch-geometric=2.0.3
• torch-scatter=2.0.9
• torch-sparse=0.6.12
• torchvision=0.11.3
• tqdm=4.62.3
• transformers=4.17.0

Optimization difficulties of OOD methods (ColoredMNIST)

OOD methods are sensitive to the network initialization. We test nine OOD methods, IRMv1, VREx, FISH, SD, IGA, LfF, RSC, CLOvE, fishr, on the ColoredMNIST benchmark. Fig1 shows the OOD performance with different ERM pretrain-epochs. None of the nine OOD methods can work with a random initialzation.

To reproduce the results, run:

bash script/coloredmnist/coloredmnist_anneal.sh

generalization difficulties of OOD methods (ColoredMNIST)

Starting from a 'perfect' initialization where the model only uses the robust feature (OOD performance is maximized), what is going on if we continue training these OOD methods? Will they maintain the robustness? or decay to a spurious/singular solution? Fig2 (top) gives the latter answer.

To reproduce the results (top), run:

bash script/coloredmnist/coloredmnist_perfect_initialization_longtrain.sh

The proposed RFC on ColoredMNIST

The proposed RFC method creates a rich & simple representation to solve the optimization-generalization dilemma above. Tab1 shows the comparison of Random initialization (Rand), ERM pretrained initialization (ERM), RFC pretrained initialization (RFC / RFC(cf)). The proposed RFC consistantly boost OOD methods.

To reproduce the results, run:

bash script/coloredmnist/coloredmnist_rfc.sh

Aiming for the second easiest-to-find feature is not OOD generalization

A line of works seek OOD generalization by discovering the second easiest-to-find features, such as PI. Here we claim that the second easiest-to-find feature is not the robust solution in general. To showcase the idea, we create a 'InverseColoredMNIST' dataset where the robust feature (digits) is more predictive than the spurious feature (color).

To reproduce the results, run:

bash script/coloredmnist/inversecoloredmnist_rfc.sh

Camelyon17 experiments

Citation

If you find our code useful, please consider citing our work using the bibtex:

@inproceedings{zhang2022rich,
  title={Rich feature construction for the optimization-generalization dilemma},
  author={Zhang, Jianyu and Lopez-Paz, David and Bottou, L{\'e}on},
  booktitle={International Conference on Machine Learning},
  pages={26397--26411},
  year={2022},
  organization={PMLR}
}