Paper: Connect Later: Improving Fine-Tuning for Robustness with Targeted Augmentations (ICLR-2024)

Reviewer comment (Theory):
The proposed framework hinges on the manual design of the transformation distribution and requires knowledge about the distribution shift between source and target domains to create effective targeted augmentations. This property limits the applicability of the proposed method in scenarios where the shift is unknown or diverse.

Author response:
**We describe a general procedure for designing targeted augmentations from the source and target data distributions in Section 4 of the paper**, as well as an example of designing the redshifting augmentation used for AstroClassification within this framework. While these augmentations do require some manual design, **a relatively small amount of manual design leads to strong empirical results (SOTA on AstroClassification, iWildCam-WILDS, Camelyon17-WILDS + best performance on new dataset Redshifts).**

Reviewer comment (Evaluation):
On iWildCam, Connect Later only marginally outperforms ERM+targeted augmentations, which indicates that if there is diverse distribution shifts, the proposed method is not quite effective.

Author response:
We ran additional seeds for iWildCam to clarify these results, for a total of 15 seeds. We report 95% confidence intervals on the mean estimate in parentheses and demonstrate that **the mean F1 results of ERM+targeted augs are outside of the 95% confidence intervals of Connect Later.**

|                      | ID Test Macro F1 | OOD Test Macro F1 |
|---|---:|---:|
| ERM                  | 46.4 (0.5)       | 30.4 (0.6)        |
| Standard fine-tuning | 46.4 (0.8)       | 31.2 (0.6)        |
| ERM + targeted augs  | 51.4 (0.6)       | 36.1 (0.7)        |
| DANN                 | 48.5 (3.2)       | 31.9 (1.6)        |
| CORAL                | 40.5 (1.6)       | 27.9 (0.5)        |
| Noisy Student        | 47.5 (1.0)       | 32.1 (0.8)        |
| Connect Later        | 51.7 (0.8)       | 36.9 (0.7)        |

**To further strengthen our results, we tested on another benchmark dataset, Camelyon17-WILDS, and show that Connect Later sets a new state-of-the-art on this benchmark for DenseNet121.** Connect Later applied with the stain color jitter targeted augmentation achieved 94.9% +/- 0.4% accuracy OOD compared to the current state-of-the-art (93.8 +/- 0.2%, cite ICON). **Overall, Connect Later achieves the state-of-the-art OOD performance on iWildCam-WILDS (for ResNet-50), Camelyon17-WILDS (for DenseNet-121), and AstroClassification.**


|                      | ID Val Avg Acc | OOD Test Avg Acc |
|---|---:|---:|
| ERM                  | 89.3 (0.9)     | 65.2 (1.1)       |
| Standard fine-tuning | 92.3 (0.2)     | 91.4 (0.9)       |
| ERM + targeted augs  | 96.7 (0.0)     | 90.5 (0.4)       |
| DANN                 | 86.1 (1.3)     | 64.5 (1.2)       |
| CORAL                | 92.3 (0.7)     | 62.3 (1.9)       |
| ICON (SOTA)          | 90.1 (0.3)     | 93.8 (0.2)       |
| Connect Later        | **98.5 (0.0)**    | **94.9 (0.4)**       |

Reviewer comment (Experiments):
Could the authors test the sensitivity of the performance of Connect Later with respect to the design choice of transformation distribution?

Author response:
- **During the rebuttal period, we verified that Connect Later (with targeted augmentations) outperforms Connect Later with generic augmentations at fine-tuning time in OOD performance, tested on AstroClassification (79.9% vs. 72.7%)**. The difference in ID accuracy between Connect Later and standard fine-tuning+generic augmentations is much smaller (80.54% vs 79.9%) in comparison. This suggests that the choice of augmentation is important, especially for OOD performance.
- For variations on targeted augmentations, Gao et al. 2023 tested a variant (Copy-Paste All Backgrounds) of the targeted augmentation we use for iWildCam (Copy-Paste Same Y). Copy-Paste segments out the animal from the image and pastes it onto a different background. Copy-Paste Same Y requires that the background is from a camera trap that has observed this animal species before, preventing, e.g., camels being pasted in water. Copy-Paste All Backgrounds does not leverage this requirement. **Gao et al. found a 2% degradation in OOD performance with Copy-Paste All Backgrounds compared to Copy-Paste Same Y, indicating that the realism of the targeted augmentation is important.** 
- For pretraining augmentations, we performed an ablation on the masking percentage for the AstroClassification task (Fig 3) and demonstrated that both ID and OOD performance are relatively robust to the chosen masking percentage between 20-90%.

Reviewer comment (Novelty):
If target augmentations involve knowledge of the distribution shift, why pre-training is needed? If we are able to generate the target data, where other options include using powerful generative models [1], why do we need pre-training? It seems to be more a question of when is pre-training necessary.

Author response:
> If target augmentations involve knowledge of the distribution shift, why pre-training is needed? If we are able to generate the target data, where other options include using powerful generative models [1], why do we need pre-training? It seems to be more a question of when is pre-training necessary.

**The ablations in the paper show that pretraining is necessary**: specifically, ERM+targeted augs removes the pretraining step from Connect Later and results in worse OOD accuracy on all datasets compared to Connect Later (67.5% vs. 79.9% OOD accuracy on AstroClassification, 36.3% vs. 37.2% OOD macro F1 on iWildCam).
Note that even though targeted augmentations can transform source data to more closely resemble target data, they are not perfect and do not describe the target distribution fully. Pretraining incorporates unlabeled data from the target to learn a better representation of the target data. Generative models such as Stable Diffusion could also be used to more easily design a targeted augmentation, and we view them as complementary to our framework.

Reviewer comment (Experiments):
It is not clear if pre-training in general is useful as the experimental results were only shown with one type of pre-training method and a different pre-training for different datasets. How sensitive are the results to the choice of pre-training method?

Author response:
> the experimental results were only shown with one type of pre-training method and a different pre-training for different datasets. How sensitive are the results to the choice of pre-training method?

- **Connect Later improves OOD accuracy with both masked autoencoding (AstroClassification, Redshifts) and contrastive pretraining (iWildCam-WILDS, Camelyon17-WILDS).**
- During the rebuttal period, we ran the standard fine-tuning baseline with SWaV contrastive pretraining on another benchmark, Camelyon17-WILDS, and found that **the choice of pretraining method is not the source of inconsistent OOD results from pretraining.** On Camelyon17-WILDS, contrastive pretraining with standard fine-tuning improves accuracy over both ERM+targeted augs and ERM. This is very different from iWildCam-WILDS, where contrastive pretraining with standard fine-tuning did not even improve accuracy over ERM. 

- On Camelyon17-WILDS, Connect Later also achieves the new state-of-the-art ID and OOD accuracy. **Overall, Connect Later achieves the state-of-the-art OOD accuracy on iWildCam-WILDS (for ResNet-50), Camelyon17-WILDS (for DenseNet-121), and AstroClassification.**

|                      | ID Val Avg Acc | OOD Test Avg Acc |
|---|---:|---:|
| ERM                  | 89.3 (0.9)     | 65.2 (1.1)       |
| Standard fine-tuning | 92.3 (0.2)     | 91.4 (0.9)       |
| ERM + targeted augs  | 96.7 (0.0)     | 90.5 (0.4)       |
| DANN                 | 86.1 (1.3)     | 64.5 (1.2)       |
| CORAL                | 92.3 (0.7)     | 62.3 (1.9)       |
| ICON (previous SoTA)          | 90.1 (0.3)     | 93.8 (0.2)       |
| Connect Later        | **98.5 (0.0)**     | **94.9 (0.4)**       |

- **Instead, we explain the inconsistency in OOD results of pretraining by differences in the connectivity of domains and classes in each dataset.** During the rebuttal period, we investigated why pretraining was much more effective for OOD performance for AstroClassification than for iWildCam by evaluating connectivity following Shen et al., 2022 [1]. We found that one of the assumptions made in [1] that governs when contrastive pretraining works for UDA was violated for iWildCam, but not for AstroClassification.

- Specifically, we evaluate the connectivity parameters alpha, beta, and gamma (alpha: across-domain (same class, different domain) connectivity, beta: across-class (same domain, different class) connectivity, gamma: across-both (different domain, different class) connectivity) for iWildCam and AstroClassification. We use the multi-crop augmentation strategy used in SwAV contrastive pretraining for iWildCam and the masking augmentation used in AstroClassification pretraining, respectively. Shen et al., 2022 requires that min(alpha, beta) > gamma for contrastive pretraining to work as a UDA method. We find that beta (across-class) < gamma (across-both) for iWildCam, which violates Shen et al.,’s assumption, whereas the assumption holds for AstroClassification. We also observe that both across-domain connectivity and across-class connectivity are much higher for AstroClassification compared to iWildCam, verifying our reasoning for the differing OOD downstream performance with pretraining.
| Dataset             | across-domain | across-class | across-both  |
|---|---:|---:|---:|
| iWildCam            | 0.116         | 0.071        | 0.076        |
| AstroClassification | 0.287         | 0.159        | 0.097        |

[1] Connect, Not Collapse: Explaining Contrastive Learning for Unsupervised Domain Adaptation, 2022.

Reviewer comment (Experiments):
From Table 2, it seems like the results are sensitive to the target augmentations. For iwilds most of the gains seem to come from the targeted augmentations. But for the astro datasets the ID test accuracy for astroclassification is better without targeted augmentations. On a new task or dataset, how should one choose the targeted augmentations?

Author response:
> From Table 2, it seems like the results are sensitive to the target augmentations. For iwilds most of the gains seem to come from the targeted augmentations. But for the astro datasets the ID test acc for astroclassification is better without targeted augmentations. On a new task or dataset, how should one choose the targeted augmentations?

**We describe a general procedure for designing targeted augmentations from the source and target data distributions in Section 4 of the paper**, as well as an example of designing the redshifting augmentation used for AstroClassification within this framework. Connect Later improves OOD accuracy on all datasets despite the varying effectiveness of targeted augmentations.

Reviewer comment (Presentation):
In Section 4, the technical details for creating augmentations are given before explaining the tasks. It may not be clear what the tasks for Astroclassification and Redshift are about, and so the technical details are not easy to understand. It is suggested to move Section 5 before Section 4.

Author response:
> Sec 4 gives the technical details for creating augmentations before explaining the tasks. It may not be clear what the tasks for AstroClassification and Redshift are about and so the technical details are not easy to understand. I would suggest moving Sec 5 before 4.

We thank the reviewer for this suggestion and will make the change in the final version.

Reviewer comment (Presentation):
In Figure 2, top row, what do the colors mean? Why does augmentation introduce errors?

Author response:
> In Fig 2, top row, what do the colors mean? Why does augmentation introduce errors?

Each color represents a set of flux measurements taken at a different wavelength. The redshifting augmentation increases the size of the error bars because this augmentation simulates placing each source object further away (i.e. at a higher redshift). Objects that are further away appear less bright and are not able to be resolved clearly (faraway objects look small and blurry), leading to larger uncertainties on each flux measurement. This is by design, since the target dataset objects look more like this (see top row, right of Fig 2).

Reviewer comment (Presentation):
It may be useful to have additional columns in the main table of results (e.g., Table 2) that describe what the different methods are, e.g., whether finetuning augmentations are used, with checks in the rows for the relevant methods.

Author response:
> It may be useful to have additional columns in the main table of results (e.g. Tab 2), that describes what the different methods are. E.g., whether finetuning augmentations are used, and with checks in the rows for the relevant methods.

We thank the reviewer for this suggestion and will make the change in the camera-ready version.

Reviewer comment (Experiments):
The authors refer several times to Shen et al. (2022), which argues that contrastive pre-training is beneficial as long as data augmentations ensure that augmented examples are more likely to change class or domain than changing both. The authors of this manuscript seem to argue in section 3 that the bad results they observe are explained by a violation of this assumption from Shen et al. (2022). However, this explanation is not tested rigorously, despite section 6 of Shen et al. (2022) providing examples on how to evaluate connectivity on real-world datasets.

Author response:
> Why does pre-training not always lead to increased OOD performance? This question remains largely unanswered. I think this paper would be better if it provided, for example, (1) clear evidence that the assumptions from Shen et al. (2022) are routinely violated in real-world datasets and commonly used augmentations	

**We empirically evaluated connectivity during the rebuttal period and show that the connectivity measurements of iWildCam-WILDS violates the assumptions from Shen et al, 2022 that govern the success of contrastive pretraining for domain adaptation, while AstroClassification has much higher connectivity** (details below). This validates our intuition that generic augmentations, such as the ones used in SwAV contrastive pretraining, are not enough to connect the domains for certain datasets and shifts. 

Specifically, we evaluate the connectivity parameters alpha, beta, and gamma (alpha: across-domain (same class, different domain) connectivity, beta: across-class (same domain, different class) connectivity, gamma: across-both (different domain, different class) connectivity) for iWildCam and AstroClassification.
We use the multi-crop augmentation strategy used in SwAV contrastive pretraining for iWildCam and the masking augmentation used in AstroClassification pretraining, respectively. Shen et al., 2022 requires that min(alpha, beta) > gamma for contrastive pretraining to work as a UDA method. **We find that beta (across-class) < gamma (across-both) for iWildCam, which violates Shen et al.,’s assumption, whereas the assumption holds for AstroClassification. We also observe that both across-domain connectivity and across-class connectivity are much higher for AstroClassification compared to iWildCam, verifying our reasoning for the differing OOD downstream performance with pretraining.** We thank the reviewer for this suggestion and will add this to the paper in the final revision.

| Dataset             | across-domain | across-class | across-both  |
|---|---:|---:|---:|
| iWildCam            | 0.116         | 0.071        | 0.076        |
| AstroClassification | 0.287         | 0.159        | 0.097        |

We follow Shen et al., 2022 to measure connectivity. First, we train a binary classifier to distinguish between class-domain pairs (e.g. for across-domain connectivity, the classifier distinguishes between augmented examples with the same class but different domain). We compute the connectivity as the test 0-1 error of these classifiers (1 - test_accuracy), averaged over class-domain pairs.

Reviewer comment (Experiments):
Different self-supervised methods were used for the iWildCam-WILDS dataset (SWaV contrastive learning) and the astronomical time-series dataset (masked autoencoding). This makes it difficult to draw conclusions from Table 1, since it is unclear which differences can be attributed to the use of different pre-training objectives.

Author response:
> different self-supervised methods were used for the iWildCam-WILDS dataset (SWaV contrastive learning) and the astronomical time-series (masked autoencoding). This makes it hard to draw conclusions from table 1, since it is unclear which differences can be attributed to the use of different pre-training objectives.

**During the rebuttal period, we ran the standard fine-tuning baseline with SWaV contrastive pretraining on another benchmark, Camelyon17-WILDS, and found that the choice of pretraining method is not the source of inconsistent OOD results from pretraining.** On Camelyon17-WILDS, contrastive pretraining with standard fine-tuning improves accuracy over both ERM+targeted augs and ERM. This is very different from iWildCam-WILDS, where contrastive pretraining with standard fine-tuning did not even improve accuracy over ERM. Our results on Camelyon17-WILDS and iWildCam-WILDS both use SwaV contrastive learning, showing that the choice of pretraining method is not the source of inconsistent OOD results. **On Camelyon17-WILDS, Connect Later also achieves the new state-of-the-art ID and OOD accuracy.**

|                      | ID Val Avg Acc | OOD Test Avg Acc |
|---|---:|---:|
| ERM                  | 89.3 (0.9)     | 65.2 (1.1)       |
| Standard fine-tuning | 92.3 (0.2)     | 91.4 (0.9)       |
| ERM + targeted augs  | 96.7 (0.0)     | 90.5 (0.4)       |
| DANN                 | 86.1 (1.3)     | 64.5 (1.2)       |
| CORAL                | 92.3 (0.7)     | 62.3 (1.9)       |
| ICON          | 90.1 (0.3)     | 93.8 (0.2)       |
| Connect Later        | **98.5 (0.0)**     | **94.9 (0.4)**       |

DANN and CORAL results are from Sagawa et al., 2022 (10 seeds) and the previous state-of-the-art results are from ICON (https://github.com/a-tea-guy/ICON, 3 seeds). All other results are averaged over 20 seeds. We report 95% confidence intervals in parentheses.

Reviewer comment (Evaluation):
The results on the iWildCam dataset are not very strong; given the means and standard deviations listed, there is a 20% chance that ERM + targeted augmentations outperforms Connect Later for any run. This leaves the improvements in the astronomical time-series datasets as the stronger evidence of Connect Later's performance. However, these results are somewhat confusing: standard fine-tuning appears to improve in-domain performance (Tables 1 and 2), suggesting that these datasets may be very small and benefit strongly from the regularization provided by pre-training. This raises the question of whether Connect Later is truly addressing a domain shift issue or merely mitigating a regular overfitting issue. The latter seems plausible given the small size of the dataset (674 objects). To draw stronger conclusions, the method should be tested on a wider variety of larger datasets.

Author response:
> The results on the iWildCam dataset are not very strong (given the means and standard deviations listed, there's a 20% chance that ERM+targeted augs outperforms Connect Later for any run). 

We ran additional seeds for iWildCam to clarify these results, for a total of 15 seeds for each method. We report 95% confidence intervals on the mean estimate in parentheses and demonstrate that **the mean F1 results of ERM+targeted augs are outside of the 95% confidence intervals of Connect Later. Overall, Connect Later achieves the state-of-the-art OOD accuracy on iWildCam-WILDS (for ResNet-50), Camelyon17-WILDS (for DenseNet-121), and AstroClassification.**

|                      | ID Test Macro F1 | OOD Test Macro F1 |
|---|---:|---:|
| ERM                  | 46.4 (0.5)       | 30.4 (0.6)        |
| Standard fine-tuning | 46.4 (0.8)       | 31.2 (0.6)        |
| ERM + targeted augs  | 51.4 (0.6)       | 36.1 (0.7)        |
| DANN                 | 48.5 (3.2)       | 31.9 (1.6)        |
| CORAL                | 40.5 (1.6)       | 27.9 (0.5)        |
| Noisy Student        | 47.5 (1.0)       | 32.1 (0.8)        |
| Connect Later        | 51.7 (0.8)       | 36.9 (0.7)        |



> is Connect Later really addressing a domain shift issue? Or is it just addressing a regular overfitting issue? The latter seems plausible given the very small size of the dataset (6,74 objects). I would argue that in order to draw conclusions, the method should be tested on a wider variety of larger datasets.

While pretraining could have a regularization effect for small datasets, **we tested Connect Later on two larger distribution shift benchmarks: iWildCam-WILDS (130k labeled examples) and Camelyon17-WILDS (300k labeled examples), where Connect Later still improves OOD performance**. We also tested standard fine-tuning with generic augmentations on AstroClassification (as a generic form of regularization), and this underperforms Connect Later by 7% OOD.

Reviewer comment (Experiments):
A baseline appears to be missing: the use of general augmentations during fine-tuning. Based on the reviewer's understanding, standard fine-tuning uses no augmentations at all. Including this baseline would clarify whether targeted augmentations during fine-tuning are essential or whether any form of augmentation would yield similar benefits.

Author response:
> A baseline that seems to be missing as well is the use of general augmentations during fine-tuning (my understanding is that standard fine-tuning uses no augmentations at all). This would clarify whether it is important to use targeted augmentations during fine-tuning, or whether any form of augmentations would be fine.

**During the rebuttal period, we ran this suggested baseline for AstroClassification and found that Connect Later outperforms standard fine-tuning+generic augmentations by 7% on OOD accuracy (79.9% vs. 72.7%).** The difference in ID accuracy between Connect Later and standard fine-tuning+generic augmentations is much smaller (80.54% vs 79.9%) in comparison.
This result demonstrates that targeted augmentations in Connect Later are key to OOD performance, and are not replaceable with generic augmentations. For the generic augmentations, we applied a 20% masking augmentation at fine-tuning time. We thank the reviewer for this suggestion and will add this to the paper in the final revision.

Reviewer comment (Experiments):
The insights provided by the paper seem limited, as it is not surprising that targeted augmentations would help with fine-tuning, given prior evidence that they help with regular ERM. The bigger question raised in section 3 — why pre-training does not always lead to increased OOD performance — remains largely unanswered. The paper would be stronger if it provided (1) clear evidence that the assumptions from Shen et al. (2022) are routinely violated in real-world datasets and commonly used augmentations, (2) a thorough evaluation of how different pre-training techniques (MAEs, contrastive learning) perform in the UDA setting, or (3) strong evidence that targeted augmentations are particularly important compared to general augmentations.

Author response:
> Why does pre-training not always lead to increased OOD performance? This question remains largely unanswered. I think this paper would be better if it provided, for example, (1) clear evidence that the assumptions from Shen et al. (2022) are routinely violated in real-world datasets and commonly used augmentations	

**We empirically evaluated connectivity during the rebuttal period and show that the connectivity measurements of iWildCam-WILDS violates the assumptions from Shen et al, 2022 that govern the success of contrastive pretraining for domain adaptation, while AstroClassification has much higher connectivity** (details below). This validates our intuition that generic augmentations, such as the ones used in SwAV contrastive pretraining, are not enough to connect the domains for certain datasets and shifts. 

Specifically, we evaluate the connectivity parameters alpha, beta, and gamma (alpha: across-domain (same class, different domain) connectivity, beta: across-class (same domain, different class) connectivity, gamma: across-both (different domain, different class) connectivity) for iWildCam and AstroClassification.
We use the multi-crop augmentation strategy used in SwAV contrastive pretraining for iWildCam and the masking augmentation used in AstroClassification pretraining, respectively. Shen et al., 2022 requires that min(alpha, beta) > gamma for contrastive pretraining to work as a UDA method. **We find that beta (across-class) < gamma (across-both) for iWildCam, which violates Shen et al.,’s assumption, whereas the assumption holds for AstroClassification. We also observe that both across-domain connectivity and across-class connectivity are much higher for AstroClassification compared to iWildCam, verifying our reasoning for the differing OOD downstream performance with pretraining.** We thank the reviewer for this suggestion and will add this to the paper in the final revision.

| Dataset             | across-domain | across-class | across-both  |
|---|---:|---:|---:|
| iWildCam            | 0.116         | 0.071        | 0.076        |
| AstroClassification | 0.287         | 0.159        | 0.097        |

We follow Shen et al., 2022 to measure connectivity. First, we train a binary classifier to distinguish between class-domain pairs (e.g. for across-domain connectivity, the classifier distinguishes between augmented examples with the same class but different domain). We compute the connectivity as the test 0-1 error of these classifiers (1 - test_accuracy), averaged over class-domain pairs. 

> different self-supervised methods were used for the iWildCam-WILDS dataset (SWaV contrastive learning) and the astronomical time-series (masked autoencoding). This makes it hard to draw conclusions from table 1, since it is unclear which differences can be attributed to the use of different pre-training objectives.

**During the rebuttal period, we ran the standard fine-tuning baseline with SWaV contrastive pretraining on another benchmark, Camelyon17-WILDS, and found that the choice of pretraining method is not the source of inconsistent OOD results from pretraining.** On Camelyon17-WILDS, contrastive pretraining with standard fine-tuning improves accuracy over both ERM+targeted augs and ERM. This is very different from iWildCam-WILDS, where contrastive pretraining with standard fine-tuning did not even improve accuracy over ERM. Our results on Camelyon17-WILDS and iWildCam-WILDS both use SwaV contrastive learning, showing that the choice of pretraining method is not the source of inconsistent OOD results. **On Camelyon17-WILDS, Connect Later also achieves the new state-of-the-art ID and OOD accuracy.**

|                      | ID Val Avg Acc | OOD Test Avg Acc |
|---|---:|---:|
| ERM                  | 89.3 (0.9)     | 65.2 (1.1)       |
| Standard fine-tuning | 92.3 (0.2)     | 91.4 (0.9)       |
| ERM + targeted augs  | 96.7 (0.0)     | 90.5 (0.4)       |
| DANN                 | 86.1 (1.3)     | 64.5 (1.2)       |
| CORAL                | 92.3 (0.7)     | 62.3 (1.9)       |
| ICON          | 90.1 (0.3)     | 93.8 (0.2)       |
| Connect Later        | **98.5 (0.0)**     | **94.9 (0.4)**       |

DANN and CORAL results are from Sagawa et al., 2022 (10 seeds) and the previous state-of-the-art results are from ICON (https://github.com/a-tea-guy/ICON, 3 seeds). All other results are averaged over 20 seeds. We report 95% confidence intervals in parentheses.