# Connect Later: Improving Fine-tuning for
Robustness with Targeted Augmentations

Helen Qu1, Sang Michael Xie2
<EMAIL_ADDRESS><EMAIL_ADDRESS>
1University of Pennsylvania
2Stanford University

###### Abstract

Models trained on a labeled source domain (e.g., labeled images from wildlife
camera traps) often generalize poorly when deployed on an out-of-distribution
(OOD) target domain (e.g., images from new camera trap locations). In the
domain adaptation setting where unlabeled target data is available, self-
supervised pretraining (e.g., masked autoencoding or contrastive learning) is
a promising method to mitigate this performance drop. Pretraining improves OOD
error when the generic data augmentations used (e.g., masking or cropping)
connect the source and target domains, which may be far apart in the input
space. In this paper, we show on real-world tasks that standard fine-tuning
after pretraining does not consistently improve OOD error over simply training
from scratch on labeled source data. To better leverage pretraining for
distribution shifts, we propose Connect Later: after pretraining with generic
augmentations, fine-tune with _targeted augmentations_ designed with knowledge
of the distribution shift. Pretraining learns good representations within the
source and target domains, while targeted augmentations connect the domains
better during fine-tuning. Connect Later improves average OOD error over
standard fine-tuning and supervised learning with targeted augmentations on 4
real-world datasets: Connect Later achieves the state-of-the-art on
astronomical time-series classification (AstroClassification) by 2.5%,
wildlife species identification (iWildCam-WILDS) with ResNet-50 by 0.9%, and
tumor identification (Camelyon17-WILDS) with DenseNet121 by 1.1%; as well as
best performance on a new dataset for astronomical time-series redshift
prediction (Redshifts) by 0.03 RMSE (11% relative).111Code and datasets are
available at https://github.com/helenqu/connect-later.

##  1 Introduction

In many real-world scenarios, machine learning models are deployed on data
that differ significantly from the training data [1, 2]. We focus on
unsupervised domain adaptation [3, 4, 5], where we have labeled data from a
source domain and unlabeled data from a target domain. We aim to learn a model
that generalizes well to these out-of-distribution (OOD) target domain inputs.
A real-world example is wildlife identification, where the task is to identify
animal species from static camera trap images. However, human labels are only
available for images from a small subset of these cameras. Images from labeled
cameras may not be representative of the habitats and characteristics of
unlabeled camera images.

Pretraining on broad unlabeled data has shown promising results on improving
OOD error in real-world problems[6, 7, 8, 9]. In particular, contrastive
pretraining has been shown to learn representations that transfer well across
domains [7, 10]. In contrast to conventional domain adaptation methods that
focus on learning domain-invariant features [11, 12, 13, 14, 15, 16],
contrastive pretraining learns representations that are not domain-invariant,
but instead decompose the class and domain information, facilitating transfer
across domains [7]. A favorable decomposition depends on the generic data
augmentations used during contrastive pretraining to connect the source and
target domains in a structured way. Intuitively, augmented (e.g. masked or
cropped) source and target inputs are more likely to look similar if they are
from the same class (e.g., cropping out the face of a lion in different
habitats) than from different classes (e.g., no body parts of elephants and
lions are alike). However, these generic augmentations must connect the
domains without knowledge of the distribution shift.

In this paper, we find on real-world benchmarks that standard fine-tuning
after contrastive pretraining is not always effective for improving OOD error
over purely supervised learning from scratch with labeled source data (Section
3). On the other hand, supervised learning with _targeted augmentations_ [17]
designed for the distribution shift improves OOD error over the supervised
learning baseline on all datasets without access to any target unlabeled data.
Thus, pretraining does not always learn representations that transfer across
domains with standard fine-tuning.

To better leverage pretraining for domain adaptation, we propose the Connect
Later framework (Figure 1): after pretraining with generic augmentations,
fine-tune with targeted augmentations (Section 4). Intuitively, pretraining
learns good representations within each domain, while targeted augmentations
better connect the domains. On a simple theoretical example where contrastive
pretraining fails, Connect Later generalizes well to the target domain. We
provide a general methodology for constructing these targeted augmentations,
where the augmented inputs match the target distribution on a feature space
where the domains differ.

We evaluate our framework on 4 real-world datasets: wildlife identification
[iWildCam-WILDS, 18, 9], tumor detection [Camelyon17-WILDS, 19, 9] and 2
astronomical time series tasks, AstroClassification and Redshifts, which we
curate from The PLAsTiCC team et al. [20]. In Section 5, we show that Connect
Later improves ID and OOD performance over standard fine-tuning or supervised
learning with targeted augmentations across all datasets. Although our
understanding stems from contrastive learning, we empirically apply Connect
Later to better leverage pretrained representations from both masked
autoencoding and contrastive learning. Connect Later achieves the state-of-
the-art on three benchmarks, improving OOD accuracy on AstroClassification by
3% [21], iWildCam-WILDS with ResNet-50 by 0.9%, and Camelyon17-WILDS with
DenseNet121 by 1.1%. We also contribute the Redshifts dataset, on which
Connect Later produces 11% and 14% relative improvement over standard fine-
tuning and supervised learning with targeted augmentations, respectively.

Figure 1: Overview of the Connect Later framework for domain adaptation
applied to a toy binary classification problem with two domains (filled and
unfilled points), showing the representations from contrastive pretraining in
$\mathbb{R}^{2}$. (Left) After contrastive pretraining with generic
augmentations, the classes within each domain are linearly separable in
representation space. Since the domains are far apart in input space, generic
augmentations may not connect the domains, resulting in misalignment in the
pretrained representations. In this case, a classifier (with a linearly
extrapolating decision boundary, dashed and solid line) learned on labeled
source data will misclassify the target data. (Right) Connect Later employs
targeted augmentations (filled points with black border), which are designed
with knowledge of the distribution shift, to connect the domains better,
resulting in a classifier that generalizes well to the target domain.

##  2 Setup

We consider a prediction problem from an input space $\mathcal{X}$ to a label
space $\mathcal{Y}$, where $\mathcal{Y}=\\{1,\dots,k\\}$ for classification
and $\mathcal{Y}\in\mathbb{R}$ for regression.

#### Domain adaptation.

Let $P_{S}$ and $P_{T}$ be the source and target input distributions over
$\mathcal{X}$, respectively. We consider unsupervised domain adaptation, where
we have access to source inputs $x\sim P_{S}$, with corresponding labels
$y\in\mathcal{Y}$ sampled from the label distribution $p^{*}(\cdot\mid x)$,
along with unlabeled target inputs sampled from the target distribution
$P_{T}$. Let the unlabeled distribution $P_{U}=\beta P_{S}+(1-\beta)P_{T}$ be
a mixture of the source and target, where $\beta\in[0,1]$. In practice,
$P_{U}$ may also be a broader unlabeled distribution. The goal is to learn a
model $f:\mathcal{X}\rightarrow\mathcal{Y}$ that minimizes error on the target
domain $L_{T}(f)=\mathbb{E}_{x\sim P_{T},y\sim p^{*}(\cdot\mid
x)}[\ell(f(x),y)]$. For example,
$\ell:\mathcal{Y}\times\mathcal{Y}\rightarrow\mathbb{R}$ is the 0-1 loss in
classification and squared loss in regression.

#### Augmentations.

Augmented inputs ${x^{\prime}}\in\mathcal{X}$ are drawn from an augmentation
distribution $\mathcal{A}(\cdot|x)$, given an input $x\in\mathcal{X}$.
Training with augmented inputs is often used to improve robustness [22, 23]
and is a crucial part of the objective in contrastive pretraining [6, 7, 24].
In this work, we define two distinct augmentation distributions,
$\mathcal{A}_{\text{pre}}$ and $\mathcal{A}_{\text{ft}}$, for the pretraining
and fine-tuning steps, respectively. Typically, the pretraining augmentations
$\mathcal{A}_{\text{pre}}$ are generic transformations, such as random
cropping in vision or masking in NLP [6, 25, 26, 8, 7, 27, 24]. Fine-tuning
augmentations $\mathcal{A}_{\text{ft}}$ have not been studied extensively and
are typically also generic or simply the identity transformation [9, 24].

#### Contrastive pretraining for domain adaptation.

Contrastive pretraining for domain adaptation consists of two steps: self-
supervised pretraining on unlabeled data, then supervised fine-tuning on
labeled source data [7]. For simplicity below, we consider the population
objectives. Contrastive learning aims to learn an encoder which maps augmented
views of the same input to similar features (“positive pairs”) and views of
different inputs to dissimilar features (“negative pairs”), according to some
distance metric. Formally, let $S_{+}(x,x^{+})=\mathbb{E}_{\bar{x}\sim
P_{U}}[\mathcal{A}_{\text{pre}}(x\mid\bar{x})\mathcal{A}_{\text{pre}}(x^{+}\mid\bar{x})]$
be the distribution over positive pairs, which are augmentations of a single
input $\bar{x}$. We pretrain an encoder
$\phi:\mathcal{X}\rightarrow\mathbb{R}^{k}$ to minimize the distance $d_{+}$
between positive pair embeddings and maximize the distance $d_{-}$ between
negative pair embeddings:

$\displaystyle\mathcal{L}_{\text{pretrain}}(\phi)=\mathbb{E}_{(x,x^{+})\sim
S_{+}}[d_{+}(\phi(x),\phi(x^{+}))]-\mathbb{E}_{x,{x^{\prime}}\sim
P_{U}}[d_{-}(\phi(x),\phi({x^{\prime}}))].$ (1)

The output of the pretraining step is a pretrained encoder
$\widehat{\phi}=\operatorname*{arg\,min}_{\phi}\mathcal{L}_{\text{pretrain}}(\phi)$.

Fine-tuning then learns a prediction head
$h:\mathbb{R}^{k}\rightarrow\mathbb{R}^{n}$ (for regression, we let $n=1$) on
top of the pretrained encoder using labeled source data with the objective

$\displaystyle\mathcal{L}_{\text{ft}}(h)=\mathbb{E}_{x\sim P_{S},y\sim
p^{*}(\cdot\mid
x),{x^{\prime}}\sim\mathcal{A}_{\text{ft}}(\cdot|x)}[\text{loss}_{\text{ft}}(h(\widehat{\phi}({x^{\prime}})),\;y;\;\theta)]$
(2)

where
$\text{loss}_{\text{ft}}:\mathbb{R}^{n}\times\mathcal{Y}\rightarrow\mathbb{R}$
is a fine-tuning objective such as softmax cross entropy loss for
classification or squared error for regression. The learned head is
$\widehat{h}=\operatorname*{arg\,min}_{h}\mathcal{L}_{\text{ft}}(h)$. In
practice, we jointly fine-tune the head $h$ and the encoder $\widehat{\phi}$.

#### Standard fine-tuning.

We refer to standard fine-tuning as the pretraining+fine-tuning procedure
where $\mathcal{A}_{\text{ft}}(x^{\prime}\mid x)=1$ if $x^{\prime}=x$ (no
fine-tuning augmentations). In our experiments, the standard fine-tuning
procedure is linear probing then fine-tuning (LP-FT) [28], which has been
shown to improve ID and OOD performance over vanilla fine-tuning. In LP-FT, we
first learn a linear predictor on top of frozen pretrained features before
fine-tuning all the parameters jointly.

#### ERM with augmentations.

As a baseline, we consider empirical risk minimization (ERM) with data
augmentation, which optimizes the fine-tuning objective (Equation 2) on
labeled source data with randomly initialized parameters. In this paper, we
refer to ERM as the instantiation where
$\mathcal{A}_{\text{ft}}(x^{\prime}\mid x)=1$ if $x^{\prime}=x$ (no
augmentations) and ERM + targeted augmentations as the instantiation with
$\mathcal{A}_{\text{ft}}$ that is designed with knowledge of the distribution
shift.

##  3 Pretraining may not improve OOD performance

Table 1: Contrastive pretraining with standard fine-tuning substantially improves OOD performance for Camelyon17-WILDS but is not as effective for iWildCam-WILDS. Results are averaged over 15 trials for iWildCam-WILDS and 20 trials for Camelyon17-WILDS, and we report the 95% confidence intervals on each mean estimate. | iWildCam (Macro F1, $\uparrow$) | Camelyon17 (Avg Acc, $\uparrow$)
---|---|---
| ID Test | OOD Test | ID Val | OOD Test
ERM | $46.4\pm 0.5$ | $30.4\pm 0.6$ | $89.3\pm 0.9$ | $65.2\pm 1.1$
Standard fine-tuning | $46.4\pm 0.8$ | $31.2\pm 0.6$ | $92.3\pm 0.2$ | $91.4\pm 0.9$

We compare ERM and standard fine-tuning on two benchmark datasets, iWildCam-
WILDS (wildlife species identification) and Camelyon17-WILDS (tumor
detection). In Table 1, we show that standard fine-tuning on a model
pretrained using SwAV contrastive learning [6] makes minimal gains over ERM on
iWildCam-WILDS ($46.4\rightarrow 46.4$ ID, $30.4\rightarrow 31.2$ OOD)
compared to Camelyon17-WILDS ($89.3\rightarrow 92.3$ ID, $65.2\rightarrow
91.4$ OOD). This result runs contrary to empirical results demonstrating that
contrastive pretraining is an effective domain adaptation method [6, 7, 8, 9].
We hypothesize that the generic pretraining augmentations connect the domains
better for some tasks and distribution shifts than others.

#### Simple example with misaligned connectivity structure.

To understand this phenomenon, we provide a simple binary classification
example of when contrastive pretraining fails for domain adaptation, following
a similar augmentation graph construction to Shen et al. [7], in Appendix E.
When the connectivity structure misaligns the source and target domains, such
that examples from the same class are less “connected” than examples from
different classes across the domains, a linear classifier trained on these
pretrained representations will not transfer from source to target. This could
happen, for example, when the source and target are far apart in input space
and connectivity is low between examples from the same class across different
domains.

###  3.1 Robustness gains from pretraining depend on dataset connectivity

To test this hypothesis, we empirically measure connectivity as defined in
Shen et al. [7]. We follow [7] and work in the augmentation graph setting,
where nodes are inputs and edge weights are the positive-pair probabilities
given by $S_{+}$. We define connectivity between a class-domain pair
$((y_{1},d_{1}),(y_{2},d_{2}))$ under four scenarios:

$\displaystyle\begin{cases}{\rho}&y_{1}=y_{2},d_{1}=d_{2}~{}\text{~{}~{}(same
class, same domain)}\\\ {\alpha}&y_{1}=y_{2},d_{1}\neq d_{1}\text{~{}~{}(same
class, different domain)}\\\ {\beta}&y_{1}\neq
y_{2},d_{1}=d_{2}\text{~{}~{}(different class, same domain)}\\\
{\gamma}&y_{1}\neq y_{2},d_{1}\neq d_{2}\text{~{}~{}(different class and
domain)}\\\ \end{cases},$ (3)

where each value is an average edge weight over the edges that satisfy each
case. Shen et al. [7] show in simple augmentation graphs that contrastive
pretraining theoretically learns transferable representations when
${\alpha}>{\gamma}$ and ${\beta}>{\gamma}$, and that the ratios
$\frac{{\alpha}}{{\gamma}}$ and $\frac{{\beta}}{{\gamma}}$ empirically
correlate well with OOD accuracy. Intuitively, the pretraining augmentations
are less likely to change both the domain and class of an input than changing
just domain or just class.

Table 2: Empirically estimated connectivity measures for iWildCam-WILDS and Camelyon17-WILDS. From Shen et al. [7], contrastive pretraining theoretically learns transferable representations for UDA when both across-domain ($\alpha$) and across-class ($\beta$) connectivity is greater than across-both ($\gamma$), using notation from [7]. In iWildCam-WILDS, across-both connectivity $>$ across-class, which violates the condition, while Camelyon17-WILDS satisfies the condition. | across-domain ($\alpha$) | across-class ($\beta$) | across-both ($\gamma$)
---|---|---|---
iWildCam-WILDS | 0.116 | 0.071 | 0.076
Camelyon17-WILDS | 0.16 | 0.198 | 0.152

#### Empirical evaluations of connectivity.

To better understand why contrastive pretraining performs differently on
iWildCam-WILDS and Camelyon17-WILDS, we empirically evaluate the connectivity
measures for iWildCam-WILDS and Camelyon17-WILDS, following Shen et al. [7].
Using augmented inputs from 2 class-domain pairs, we train a binary classifier
to predict the class-domain pair of each input, and interpret the test error
of the classifier as an estimate for connectivity. We average each
connectivity value over 15 class-domain pairs (see Appendix D for details).
Our results, summarized in Table 2, show that iWildCam-WILDS connectivity
measures violate the condition for contrastive pretraining in the UDA setting,
since across-both connectivity $>$ across-class ($\gamma>\beta$). This finding
is consistent with our observation that contrastive pretraining fails for
iWildCam-WILDS while producing gains for Camelyon17-WILDS, and further
underscores the need for domain adaptation methods that correct the misaligned
connectivity structure.

##  4 Connect Later: Pretrain First, Targeted Augmentations Later

Even when generic augmentations applied during pretraining misalign the
connectivity structure, the pretrained representations are still useful since
the classes are linearly separable within each domain. How do we leverage
these pretrained representations when they may not transfer well across
domains? In this work, we propose the Connect Later framework (Figure 1):

  1. 1.

Pretrain on unlabeled data with generic augmentations as in Equation 1,
producing a pretrained encoder $\widehat{\phi}$.

  2. 2.

Design a targeted augmentation $\mathcal{A}_{\text{ft}}$ (discussed below) and
use augmented source data to fine-tune the pretrained encoder $\widehat{\phi}$
jointly with a prediction head $h$ as in Equation 2.

Intuitively, pretraining learns good representations of the source and target
domains and allows us to reuse the pretrained model for multiple downstream
tasks, while targeted augmentations better connect the domains for the
particular task.

#### Simple example where Connect Later achieves 0 OOD error.

In our simple binary classification example in Appendix E, we show that when
the connectivity structure is misaligned, both standard fine-tuning with
contrastive pretraining and ERM + targeted augmentations have high OOD error,
while Connect Later achieves 0 OOD error. In this setting, ERM with targeted
augmentations is unable to achieve 0 OOD error since some target inputs are
“unreachable” via targeted augmentations of source inputs. The pretraining
step in Connect Later uses unlabeled target data to learn representations
where label information from source data can propagate to all target inputs.

###  4.1 Real-world examples of targeted augmentations

We design targeted augmentations for 4 real-world tasks: species
identification from wildlife camera trap images, tumor detection, and
astronomical time-series classification and redshift prediction. We show
examples from the source, augmented, and target datasets for these tasks in
Figure 2.

Figure 2: An example from the source dataset (left), an augmented version of
the source example (middle), and an example from the target dataset (right)
for our 3 tasks. (Top row) The target dataset in AstroClassification and
Redshifts is much higher redshift than the source dataset. We apply the
redshifting augmentation to simulate placing source objects at a higher
redshift to better match the target dataset. The flux errors and flux values
of the augmented example (middle) show much better resemblance of the target
example. (Bottom row) The iWildCam-WILDS target dataset is composed of images
from cameras unseen during fine-tuning, which are in potentially new habitats,
so we randomize the habitat background by applying the Copy-Paste Same Y
augmentation. This algorithm places source dataset animals into empty
backgrounds from other cameras that have observed the same species. iWildCam-
WILDS image examples shown here are from Gao et al. [17].

#### Wildlife species classification (iWildCam-WILDS).

For iWildCam-WILDS [18, 9], the task is to identify the wildlife species from
static camera trap images. These cameras are placed in a wide variety of
environments, which all have unique habitat conditions (e.g., African savannah
vs. tropical rainforest) and camera characteristics (e.g., angles,
resolutions). In this dataset, we use labeled data from 243 camera traps to
learn a model that can generalize to data from 48 unseen camera traps.

  * •

Source: 243 camera traps

  * •

Target: 48 unseen camera traps

  * •

Targeted Augmentation: We augment the labeled dataset with the Copy-Paste Same
Y algorithm, which uses image segmentation to copy-paste the animal onto
different background images from cameras that have observed the same species
[17].

  * •

Task: 182-class wildlife species classification

#### Tumor detection (Camelyon17-WILDS).

The task in Camelyon17-WILDS [19] is to classify whether a patch of a
histopathology slide contains a tumor. These slides are contributed from
multiple hospitals, which use different stain colors and also vary in
distributions of patient cancer stage.

  * •

Source: Hospitals 1-3.

  * •

Target: Hospitals 4 and 5.

  * •

Targeted Augmentation: We augment the labeled dataset with the Stain Color
Jitter algorithm, which jitters the color of the slide image in the
hematoxylin and eosin staining color space [29].

  * •

Task: Binary classification of whether a slide contains a tumor.

#### Astronomical object classification (AstroClassification).

Astronomical object classification [21, 30] involves predicting the object
type (e.g., type II supernova) from a time series of an object’s brightness at
multiple wavelengths (light curves). We curate this dataset from the
Photometric LSST Astronomical Time Series Classification Challenge [PLAsTiCC,
20] (details in Appendix A.1).

  * •

Source: Time-series of bright, nearby objects with expert labels

  * •

Target: Time-series of all observed objects from the telescope, often faint
and distant (higher redshift). Follow-up observation, which is required for
expert labeling, is too expensive for these objects.

  * •

Targeted Augmentation: We augment the labeled dataset by redshifting each
object, i.e., simulating its observed properties as if it were further away
(details in Appendix B.2).

  * •

Task: 14-class astronomical object classification

#### Redshift regression (Redshifts).

Similar to object type, redshift information is also available only for
bright, nearby objects. We predict the scalar redshift value of each object
and minimize mean squared error. This task has been studied for individual
object types, such as quasars [31] and type Ia supernovae [32], but we
consider a more realistic set of multiple object types. The labeled and
unlabeled data are derived from the PLAsTiCC dataset. Redshifts is a new
dataset that we contribute as part of this work.

  * •

Source: Time-series of bright, nearby labeled objects.

  * •

Target: Time-series of all observed objects from the telescope, often faint
and distant (higher redshift).

  * •

Targeted Augmentation: Redshifting (same as AstroClassification, Appendix
B.2).

  * •

Task: Redshift regression

###  4.2 Designing targeted augmentations

How do we design these targeted augmentations? We provide a general
methodology based on matching the target distribution on a feature space:

  1. 1.

Identify a feature space $\mathcal{Z}$. We assume that we can label
$z\in\mathcal{Z}$ for each input and that the source and target domains
largely differ on this feature space. One such example is the space of
spurious, domain-dependent features (e.g., camera angle or resolution for
iWildCam-WILDS), which is the approach followed by Gao et al. [17].

  2. 2.

Fit a transformed feature distribution $\hat{p_{T}}(z^{\prime}|z)$ to the
target feature distribution.

  3. 3.

Create a transformation distribution $T(x^{\prime}|x,z^{\prime})$ where
$x^{\prime}$ is the augmented version of $x$ with $z=z^{\prime}$. In this
paper, we define $T$ with domain knowledge.

  4. 4.

Given an input $x$, generate augmentations by sampling a new feature
$z^{\prime}$ from $\hat{p_{T}}(z^{\prime}\mid z)$, then sampling an
augmentation from $T(x^{\prime}|x,z^{\prime})$. The resulting targeted
augmentation probabilities are $\mathcal{A}_{\text{ft}}(x^{\prime}\mid
x)=\sum_{z^{\prime}}T(x^{\prime}\mid x,z^{\prime})\hat{p_{T}}(z^{\prime}\mid
z)$.

#### Targeted augmentation example.

We follow the procedure outlined above to design a targeted augmentation for
AstroClassification and Redshifts, two astronomical time-series datasets (see
Appendix B.2 for details). In these datasets, expert labels are only available
for time-series of bright, nearby astronomical objects, while the unlabeled
dataset contains mostly faint, distant objects. We describe distances in terms
of cosmological redshift: nearby objects have lower redshift values than
distant objects, causing the source and target redshift distributions to be
mismatched (Appendix Figure 4).

  1. 1.

The source and target domains primarily differ on their redshift
distributions, so we identify this scalar feature as $z$.

  2. 2.

We roughly fit the target redshift distribution while constraining the
transformed redshift value to not be too far from the original redshift $z$,
such that $\hat{p_{T}}(z^{\prime}\mid z)$ is distributed as
$\text{loguniform}(0.95z,\;\text{min}(1.5(1+z)-1,\;5z))$, following Boone
[21].

  3. 3.

We define a transformation distribution $T(x^{\prime}|x,z^{\prime})$, where
$x$ is a time-series of flux values at multiple wavelengths and $z^{\prime}$
is a new redshift value to transform to. We first fit a Gaussian process that
models $x$ as a function of time and wavelength. Given $z^{\prime}$, we
rescale the timestamps and wavelengths of the original input to account for
the physical effects of the new redshift value. Then, we sample
$\tilde{x^{\prime}}$ from the Gaussian process at these new timestamps and
wavelengths. Finally, we produce the transformed input $x^{\prime}$ by scaling
the flux values to account for $z^{\prime}$.

  4. 4.

We sample $z^{\prime}$ from $\hat{p_{T}}(z^{\prime}\mid z)$ and then sample
augmentations $x^{\prime}$ from $T(x^{\prime}|x,z^{\prime})$.

##  5 Experiments

We empirically test Connect Later with contrastive pretraining (iWildCam-
WILDS, Camelyon17-WILDS) as well as pretraining with masked autoencoding
(AstroClassification, Redshifts) to demonstrate Connect Later as a general
fine-tuning method.

#### Training procedure.

For iWildCam-WILDS, we use a ResNet-50 model pretrained on unlabeled ImageNet
data with SwAV contrastive learning [6]. We use a DenseNet121 pretrained on
the unlabeled data provided in [9] with SwAV for Camelyon17-WILDS. To test
Connect Later with other pretraining strategies, we use the masked
autoencoding objective for AstroClassification and Redshifts by masking 60% of
observations from each light curve (Appendix C). The same pretrained model is
used for both tasks to demonstrate the reusability of pretrained features. We
note that masked autoencoding has been linked to contrastive learning [33]. In
particular, Zhang et al. [33] show that the masked autoencoding objective
upper bounds the contrastive loss between positive pairs — thus, masked
autoencoding implicitly aligns the positive pairs induced by the masking
augmentations.

We fine-tune the pretrained models with linear probing then fine-tuning (LP-
FT) [28], which has been shown to improve OOD performance.

#### Baselines.

We evaluate our framework against three baselines: ERM, ERM+targeted augs, and
standard fine-tuning. We include Avocado [21], the previous state-of-the-art
model for AstroClassification. We also include a self-training baseline for
AstroClassification and Redshifts, which has been shown to perform well on
some real-world datasets [9]. For the self-training baseline, we pseudo-label
the target dataset with a trained ERM+targeted augs model, then perform the
same targeted augmentation on the pseudo-labeled target dataset. We then train
a model with the pseudo-labeled and augmented target dataset combined with the
labeled source dataset. We include additional domain adaptation baselines for
iWildCam-WILDS and Camelyon17-WILDS: domain-adversarial neural networks (DANN)
[11], correlation alignment (CORAL) [15], Noisy Student [34], and
ICON222https://github.com/a-tea-guy/ICON.

###  5.1 Main results

Table 3: ID and OOD accuracy (%) for AstroClassification and RMSE for Redshifts of each method. Results are averaged over 5 trials and rows with means within 1 STD of the best mean are bolded. | AstroClassification | Redshift
---|---|---
| ID Test Acc ($\uparrow$) | OOD Acc ($\uparrow$) | ID Test RMSE ($\downarrow$) | OOD RMSE ($\downarrow$)
ERM | $71.59\pm 1.10$ | $61.26\pm 1.10$ | $0.274\pm 0.016$ | $0.320\pm 0.009$
Standard fine-tuning | $78.84\pm 0.97$ | $67.84\pm 0.70$ | $\mathbf{0.246\pm 0.015}$ | $0.277\pm 0.004$
ERM + targeted augs | $68.75\pm 0.95$ | $67.54\pm 0.32$ | $0.310\pm 0.006$ | $0.286\pm 0.007$
Self-Training | $77.72\pm 0.59$ | $65.15\pm 0.67$ | $0.304\pm 0.010$ | $0.289\pm 0.003$
Avocado [21] | - | $77.40$ | - | -
Connect Later | $\mathbf{80.54\pm 1.20}$ | $\mathbf{79.90\pm 0.60}$ | $\mathbf{0.256\pm 0.005}$ | $\mathbf{0.247\pm 0.005}$
Table 4: ID and OOD performance for each method on iWildCam-WILDS and Camelyon17-WILDS. Results are averaged over 15 trials for iWildCam-WILDS and 20 trials for Camelyon17-WILDS, and we report 95% confidence intervals on each mean estimate. Rows with means within 1 interval of the best mean are bolded. | iWildCam (Macro F1, $\uparrow$) | Camelyon17 (Avg Acc, $\uparrow$)
---|---|---
| ID Test | OOD Test | ID Val | OOD Test
ERM | $46.4\pm 0.5$ | $30.4\pm 0.6$ | $89.3\pm 0.9$ | $65.2\pm 1.1$
Standard fine-tuning | $46.4\pm 0.8$ | $31.2\pm 0.6$ | $92.3\pm 0.2$ | $91.4\pm 0.9$
ERM + targeted augs | $\mathbf{51.4\pm 0.6}$ | $36.1\pm 0.7$ | $96.7\pm 0.0$ | $90.5\pm 0.4$
DANN [9] | $48.5\pm 3.2$ | $31.9\pm 1.6$ | $86.1\pm 1.3$ | $64.5\pm 1.2$
CORAL [9] | $40.5\pm 1.6$ | $27.9\pm 0.5$ | $92.3\pm 0.7$ | $62.3\pm 1.9$
Noisy Student [9] | $47.5\pm 1.0$ | $32.1\pm 0.8$ | - | -
ICON | $50.6\pm 1.3$ | $34.5\pm 1.4$ | $90.1\pm 0.4$ | $93.8\pm 0.3$
Connect Later | $\mathbf{51.7\pm 0.8}$ | $\mathbf{36.9\pm 0.7}$ | $\mathbf{98.5\pm 0.0}$ | $\mathbf{94.9\pm 0.4}$

Tables 3 and 4 compare the results of Connect Later with baseline methods. On
iWildCam-WILDS, Connect Later produces improvements in both ID and OOD
performance despite the failure of standard fine-tuning. Connect Later
substantially outperforms all other variants on the OOD metric, including
state-of-the-art performances on AstroClassification by 3% OOD, iWildCam-WILDS
by 0.8% OOD for ResNet-50, and Camelyon17-WILDS by 1.1% OOD for DenseNet121.

#### iWildCam.

On iWildCam-WILDS, standard fine-tuning does not improve over ERM in ID
performance and minimally improves over ERM in OOD performance, while
ERM+targeted augmentations improves by 6% ID and OOD over both ERM and
standard fine-tuning. While 3 of the domain adaptation baselines made ID and
OOD improvements over standard fine-tuning (DANN: $46.4\%\rightarrow 48.5\%$
ID, $31.2\%\rightarrow 31.9\%$ OOD; Noisy Student: $46.4\%\rightarrow 47.5\%$
ID, $30.4\%\rightarrow 32.1\%$ OOD; ICON: $46.4\%\rightarrow 50.6\%$ ID,
$30.4\%\rightarrow 34.5\%$), all fall short of ERM+targeted augs. Connect
Later improves over both standard fine-tuning ($30.4\%\rightarrow 37.2\%$) and
ERM+targeted augs ($36.3\%\rightarrow 37.2\%$) in OOD performance, achieving a
new state-of-the-art performance for ResNet-50 on the iWildCam-WILDS
benchmark.

#### Camelyon17.

On Camelyon17-WILDS, standard fine-tuning produces significant gains over ERM
in both ID ($89.3\%\rightarrow 92.3\%$) and OOD ($65.2\%\rightarrow 91.4\%$)
average accuracy. ERM+targeted augmentations outperforms standard fine-tuning
in ID accuracy ($92.3\%\rightarrow 96.7\%$), but does not improve OOD. DANN
underperforms ERM in both ID and OOD accuracy, while CORAL produces similar ID
accuracy as standard fine-tuning but poor OOD performance. Connect Later sets
a new state-of-the-art on Camelyon17-WILDS with DenseNet121, improving on the
best ID performance by 1.8% (ERM+targeted augs, $96.7\%\rightarrow 98.5\%$)
and OOD performance by 1.1% (ICON, $93.8\%\rightarrow 94.9\%$).

#### AstroClassification.

For AstroClassification, standard fine-tuning provides a significant
performance boost over ERM: $71.6\%\rightarrow 78.9\%$ ID, $61.3\%\rightarrow
67.8\%$ OOD. ERM+targeted augs underperforms in ID accuracy compared to ERM
alone ($71.6\%\rightarrow 68.8\%$) and standard fine-tuning
($78.9\%\rightarrow 68.8\%$), likely due to relatively strong targeted
augmentations toward the redshift distribution of distant, faint objects
present in the target distribution. However, OOD accuracy of ERM+targeted augs
is competitive with standard fine-tuning, outperforming ERM. Self-training
improves both ID and OOD performance compared to ERM but underperforms
standard fine-tuning in both domains. Connect Later outperforms the best
baseline, standard fine-tuning, by 12% OOD and 2% ID. The ID accuracy improves
over standard fine-tuning despite the drop in ID accuracy from adding targeted
augmentations to ERM, showing a complementary benefit between pretraining and
targeted augmentations. Connect Later sets a new state-of-the-art OOD
performance on AstroClassification by 3% over Avocado, a heavily tuned random
forest model with expert-designed features [21].

#### Redshifts.

Similarly to AstroClassification, standard fine-tuning significantly improves
over ERM in both ID ($0.27\rightarrow 0.25$, 7% relative) and OOD
($0.32\rightarrow 0.28$, 13% relative) RMSE. Self-training performs similarly
to ERM+targeted augs, while Connect Later outperforms the best baseline
variant, standard fine-tuning, by 0.03 RMSE (11% relative) with comparable ID
error. We use the same pretrained model for both AstroClassification and
Redshifts for standard fine-tuning and Connect Later, demonstrating the
reusability of pretrained representations.

###  5.2 Ablations

Table 5: Scaling up model size of Connect Later produces improvements in both ID and OOD performance on the AstroClassification task.  Number of Parameters | ID Accuracy ($\uparrow$) | OOD Accuracy ($\uparrow$)
---|---|---
21M (default) | 78.47 | 79.49
69M | 80.38 | 80.55
Table 6: Linear probing (LP) in addition to fine-tuning (FT) hurts performance for the ERM+targeted augs model but improves performance for Connect Later (tested on the AstroClassification task). | Connect Later | ERM+targeted augs
---|---|---
| ID Accuracy ($\uparrow$) | OOD Accuracy ($\uparrow$) | ID Accuracy ($\uparrow$) | OOD Accuracy ($\uparrow$)
FT only | 78.07 | 78.6 | 77.88 | 68.43
LP-FT | 78.47 | 79.49 | 65.68 | 67.07

We performed ablations on the model size, strength of pretraining
augmentations (masking percentage for masked autoencoding), and LP-FT on
AstroClassification. We find that downstream performance is quite robust to
masking percentage, while scaling up model size and LP-FT improve performance
for pretrained models.

Figure 3: On the AstroClassification task, Connect Later is relatively robust
to pretraining masking percentage both ID and OOD, but 60% masking performs
best out of the percentages we tested.

#### Model scale.

We tested Connect Later with a larger model ($\sim 3\times$ the parameters of
our model, $\text{21M}\rightarrow\text{69M}$), and find that scaling up model
size improves both ID and OOD accuracy (Table 5). This suggests that scaling
up the model is a promising way to further improve performance with Connect
Later.

#### Strength of pretraining augmentations (masking percentage).

We vary the strength of pretraining augmentations, which changes the
connectivity between domains. This is most straightforward with the MAE
objective, as augmentation strength is parameterized solely by masking
percentage. We tested pretraining masking percentages {20, 40, 60, 80, 90}%
while keeping the masking strategy unchanged (replace 10% of masked indices
with random values from the lightcurve, another 10% are kept unchanged, and
80% are replaced with the mask token, which we choose to be 0). We show the ID
and OOD test accuracy of each variant in Figure 3. Both ID and OOD performance
peak at 60% masking, although we find that the performance of Connect Later is
quite robust to the masking percentage, particularly for OOD performance. All
of the masking percentages we tried improve on OOD performance over standard
fine-tuning or ERM with targeted augmentations. Particularly, even with the
strongest pretraining augmentations (90% masking), which should connect the
domains more, the OOD performance did not improve over weaker augmentations.
We hypothesize that increasing the strength of generic augmentations may
indiscriminately increase the connectivity between all source and target
examples, including examples from different classes that should not be
strongly connected.

#### Linear probing then fine-tuning.

Kumar et al. [28] showed that linear probing (with fixed neural embeddings)
and then fine-tuning (LP-FT) the entire model improves both ID and OOD
performance. Intuitively, full fine-tuning with a randomly initialized linear
probe can destroy the pretrained features, and training the linear probe first
mitigates this. We test LP-FT against FT only (all model weights are fine-
tuned) with the Connect Later model and the ERM+targeted augs baseline. We
find that LP-FT improves OOD accuracy by 0.9% over FT only when applied to
Connect Later on AstroClassification (Table 6). On the other hand, LP-FT
decreased OOD accuracy by 1.4% when applied to ERM+targeted augs, which uses
random initialization (no pretraining). As a result, we use LP-FT on
pretrained models but not on ERM or ERM+targeted augs.

##  6 Discussion and Related Work

#### Augmentations for pretraining.

Data augmentations such as cropping or masking have been vital to semi- and
self-supervised learning objectives. Masking or noising the data and training
a model to reconstruct the original inputs have been shown to produce useful
pretrained representations across multiple modalities [24, 35, 27, 36, 25, 26,
6]. In contrastive learning, models are trained to distinguish augmented
“views” of the same input from views of a different input [25, 6, 26]. Our
results demonstrating inconsistent OOD performance across datasets brings up
the important future question of how to choose the best pretraining
augmentation and algorithm for learning transferable representations.

#### Augmentations for robustness.

Data augmentation has been used to improve model robustness and avoid
catastrophic failures due to spurious, label-independent changes (e.g.
translation or rotation in vision) [22, 37, 38]. The augmentation strategies
used in prior work are generic perturbations that aim to increase the
diversity of inputs [e.g., 39, 40, 41, 42, 43, 44], though a number of studies
have shown that the type of data augmentations matters for performance [25,
45]. Augmentations have also been leveraged in the self-training paradigm,
which improves generalization to unseen data by training on the pseudo-labeled
full dataset [34, 46, 47]. We show that a self-training baseline with pseudo-
labels from an ERM+targeted augs model does not outperform Connect Later,
indicating that pretraining is an important component of the framework.
Connect Later exposes targeted augmentations as a design interface for
improving robustness with knowledge of the distribution shift, while still
leveraging pretrained representations.

#### Targeted augmentations.

In problems with domain shift, Gao et al. [17] show that targeted
augmentations outperform generic augmentations on unseen data. They identify
spurious domain-dependent, label-independent features in the source dataset
and construct targeted augmentations by randomizing these features. Gao et al.
[17] consider the domain generalization setting, in which no data from the
target dataset is available. We consider targeted augmentations in the domain
adaptation setting, in which we can model the target distribution of these
spurious features with the unlabeled target data. In general, designing
targeted augmentations specific to each distribution shift may be difficult
and require expert guidance. As part of the Connect Later framework, we
provide a general methodology for the design of such augmentations. Certain
aspects, such as the selection of feature space $z$ and transformation
distribution $T$ could be learned from the unlabeled data itself, which we
leave for future work. We show that targeted augmentations better leverage
pretrained representations for complementary gains in OOD performance.

##  7 Conclusion

We show that pretraining with generic augmentations is not a panacea for all
distribution shifts and tasks, and sometimes does not outperform supervised
learning on labeled source data. Pure supervised learning, however, does not
use the unlabeled data or produce reusable representations. The Connect Later
framework allows for better leverage of pretrained representations for OOD
performance by applying targeted augmentations at fine-tuning time. Future
work could focus on learning targeted augmentations from the source and target
data distributions as well as further understanding of how the choice of
pretraining augmentations affects downstream ID/OOD performance.

## References

  * Quiñonero-Candela et al. [2009] Joaquin Quiñonero-Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D. Lawrence.  _Dataset shift in machine learning_.  The MIT Press, 2009.
  * Koh et al. [2021] Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton A. Earnshaw, Imran S. Haque, Sara Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang.  WILDS: A benchmark of in-the-wild distribution shifts.  In _International Conference on Machine Learning (ICML)_ , 2021.
  * Shimodaira [2000] Hidetoshi Shimodaira.  Improving predictive inference under covariate shift by weighting the log-likelihood function.  _Journal of Statistical Planning and Inference_ , 90:227–244, 2000.
  * Blitzer et al. [2006] John Blitzer, Ryan McDonald, and Fernando Pereira.  Domain adaptation with structural correspondence learning.  In _Empirical Methods in Natural Language Processing (EMNLP)_ , 2006.
  * Sugiyama et al. [2007] Masashi Sugiyama, Matthias Krauledat, and Klaus-Robert Muller.  Covariate shift adaptation by importance weighted cross validation.  _Journal of Machine Learning Research (JMLR)_ , 8:985–1005, 2007.
  * Caron et al. [2020] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin.  Unsupervised learning of visual features by contrasting cluster assignments.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , volume 33, pages 9912–9924, 2020.
  * Shen et al. [2022] Kendrick Shen, Robbie Jones, Ananya Kumar, Sang Michael Xie, Jeff Z. HaoChen, Tengyu Ma, and Percy Liang.  Connect, not collapse: Explaining contrastive learning for unsupervised domain adaptation.  In _International Conference on Machine Learning (ICML)_ , 2022.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.  Learning transferable visual models from natural language supervision.  In _International Conference on Machine Learning (ICML)_ , volume 139, pages 8748–8763, 2021.
  * Sagawa et al. [2022] Shiori Sagawa, Pang Wei Koh, Tony Lee, Irena Gao, Sang Michael Xie, Kendrick Shen, Ananya Kumar, Weihua Hu, Michihiro Yasunaga, H. Marklund, Sara Beery, E. David, I. Stavness, Wei Guo, J. Leskovec, Kate Saenko, Tatsunori B. Hashimoto, S. Levine, Chelsea Finn, and Percy Liang.  Extending the WILDS benchmark for unsupervised adaptation.  In _International Conference on Learning Representations (ICLR)_ , 2022.
  * HaoChen et al. [2022] Jeff Z. HaoChen, Colin Wei, Ananya Kumar, and Tengyu Ma.  Beyond separability: Analyzing the linear transferability of contrastive representations to related subpopulations.  _arXiv_ , 2022.
  * Ganin et al. [2016] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, Francois Laviolette, Mario March, and Victor Lempitsky.  Domain-adversarial training of neural networks.  _Journal of Machine Learning Research (JMLR)_ , 17, 2016.
  * Kang et al. [2019] Guoliang Kang, Lu Jiang, Yi Yang, and Alexander G Hauptmann.  Contrastive adaptation network for unsupervised domain adaptation.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 4893–4902, 2019.
  * Tzeng et al. [2017] Eric Tzeng, Judy Hoffman, Kate Saenko, and Trevor Darrell.  Adversarial discriminative domain adaptation.  In _Computer Vision and Pattern Recognition (CVPR)_ , 2017.
  * Saenko et al. [2010] Kate Saenko, Brian Kulis, Mario Fritz, and Trevor Darrell.  Adapting visual category models to new domains.  In _European conference on computer vision_ , pages 213–226, 2010.
  * Sun et al. [2016] Baochen Sun, Jiashi Feng, and Kate Saenko.  Return of frustratingly easy domain adaptation.  In _Association for the Advancement of Artificial Intelligence (AAAI)_ , 2016.
  * Hoffman et al. [2018] Judy Hoffman, Eric Tzeng, Taesung Park, Jun-Yan Zhu, Phillip Isola, Kate Saenko, Alexei A. Efros, and Trevor Darrell.  Cycada: Cycle consistent adversarial domain adaptation.  In _International Conference on Machine Learning (ICML)_ , 2018.
  * Gao et al. [2023] Irena Gao, Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang.  Out-of-domain robustness via targeted augmentations.  In _International Conference on Machine Learning (ICML)_ , 2023.
  * Beery et al. [2020] Sara Beery, Elijah Cole, and Arvi Gjoka.  The iwildcam 2020 competition dataset.  _arXiv preprint arXiv:2004.10340_ , 2020.
  * Bandi et al. [2018] Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, et al.  From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge.  _IEEE transactions on medical imaging_ , 38(2):550–560, 2018.
  * The PLAsTiCC team et al. [2018] The PLAsTiCC team, Tarek Allam Jr. au2, Anita Bahmanyar, Rahul Biswas, Mi Dai, Lluís Galbany, Renée Hložek, Emille E. O. Ishida, Saurabh W. Jha, David O. Jones, Richard Kessler, Michelle Lochner, Ashish A. Mahabal, Alex I. Malz, Kaisey S. Mandel, Juan Rafael Martínez-Galarza, Jason D. McEwen, Daniel Muthukrishna, Gautham Narayan, Hiranya Peiris, Christina M. Peters, Kara Ponder, Christian N. Setzer, The LSST Dark Energy Science Collaboration, The LSST Transients, and Variable Stars Science Collaboration.  The photometric lsst astronomical time-series classification challenge (plasticc): Data set, 2018.
  * Boone [2019] Kyle Boone.  Avocado: Photometric classification of astronomical transients with gaussian process augmentation.  _The Astronomical Journal_ , 158(6):257, dec 2019.  doi: 10.3847/1538-3881/ab5182.  URL https://doi.org/10.3847%2F1538-3881%2Fab5182.
  * Hendrycks et al. [2019] Dan Hendrycks, Norman Mu, Ekin D Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshminarayanan.  Augmix: A simple data processing method to improve robustness and uncertainty.  In _International Conference on Learning Representations (ICLR)_ , 2019.
  * Hendrycks et al. [2020] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer.  The many faces of robustness: A critical analysis of out-of-distribution generalization.  _arXiv preprint arXiv:2006.16241_ , 2020.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In _Association for Computational Linguistics (ACL)_ , pages 4171–4186, 2019.
  * Chen et al. [2020] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In _International Conference on Machine Learning (ICML)_ , pages 1597–1607, 2020.
  * He et al. [2020] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick.  Momentum contrast for unsupervised visual representation learning.  In _Computer Vision and Pattern Recognition (CVPR)_ , 2020.
  * He et al. [2022] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick.  Masked autoencoders are scalable vision learners.  In _Computer Vision and Pattern Recognition (CVPR)_ , 2022.
  * Kumar et al. [2022] Ananya Kumar, Aditi Raghunathan, Robbie Jones, Tengyu Ma, and Percy Liang.  Fine-tuning can distort pretrained features and underperform out-of-distribution.  In _International Conference on Learning Representations (ICLR)_ , 2022.
  * Tellez et al. [2018] David Tellez, Maschenka Balkenhol, Irene Otte-H
"oller, Rob van de Loo, Rob Vogels, Peter Bult, Carla Wauters, Willem Vreuls,
Suzanne Mol, Nico Karssemeijer, et al.  Whole-slide mitosis detection in h&e
breast histology using phh3 as a reference to train distilled stain-invariant
convolutional networks.  _IEEE transactions on medical imaging_ ,
37(9):2126–2136, 2018.

  * Allam Jr. and McEwen [2022] Tarek Allam Jr. and Jason D. McEwen.  Paying attention to astronomical transients: Introducing the time-series transformer for photometric classification, 2022.
  * Nakoneczny et al. [2021] SJ Nakoneczny, M Bilicki, Agnieszka Pollo, M Asgari, A Dvornik, T Erben, B Giblin, C Heymans, H Hildebrandt, A Kannawadi, et al.  Photometric selection and redshifts for quasars in the kilo-degree survey data release 4.  _Astronomy & Astrophysics_, 649:A81, 2021.
  * Qu and Sako [2023] Helen Qu and Masao Sako.  Photo-zSNthesis: Converting Type Ia Supernova Lightcurves to Redshift Estimates via Deep Learning.  _Astrophysical Journal_ , 954(2):201, September 2023.  doi: 10.3847/1538-4357/aceafa.
  * Zhang et al. [2022] Qi Zhang, Yifei Wang, and Yisen Wang.  How mask matters: Towards theoretical understandings of masked autoencoders.  _Advances in Neural Information Processing Systems_ , 35:27127–27139, 2022.
  * Xie et al. [2020a] Qizhe Xie, Minh-Thang Luong, Eduard Hovy, and Quoc V. Le.  Self-training with noisy student improves imagenet classification.  _arXiv_ , 2020a.
  * Lewis et al. [2020] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer.  Bart: Denoising sequence-to-sequence pre-training for natural language generation, translation, and comprehension.  In _Association for Computational Linguistics (ACL)_ , 2020.
  * Raffel et al. [2019] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _arXiv preprint arXiv:1910.10683_ , 2019.
  * Rebuffi et al. [2021] Sylvestre-Alvise Rebuffi, Sven Gowal, Dan Andrei Calian, Florian Stimberg, Olivia Wiles, and Timothy A Mann.  Data augmentation can improve robustness.  _Advances in Neural Information Processing Systems_ , 34:29935–29948, 2021.
  * Ng et al. [2020] Nathan Ng, Kyunghyun Cho, and Marzyeh Ghassemi.  Ssmba: Self-supervised manifold based data augmentation for improving out-of-domain robustness.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)_ , pages 1268–1283, 2020.
  * Simard et al. [2003] Patrice Y. Simard, Dave Steinkraus, and John C. Platt.  Best practices for convolutional neural networks applied to visual document analysis.  _International Conference on Document Analysis and Recognition_ , 2:958–964, 2003.
  * Krizhevsky et al. [2012] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , pages 1097–1105, 2012.
  * Cubuk et al. [2019] Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le.  Autoaugment: Learning augmentation policies from data.  In _Computer Vision and Pattern Recognition (CVPR)_ , 2019.
  * Cubuk et al. [2020] Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space.  In _Computer Vision and Pattern Recognition (CVPR)_ , pages 702–703, 2020.
  * DeVries and Taylor [2017] Terrance DeVries and Graham W Taylor.  Improved regularization of convolutional neural networks with cutout.  _arXiv preprint arXiv:1708.04552_ , 2017.
  * Zhang et al. [2017] Hongyi Zhang, Moustapha Cissé, Yann N. Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  _ICLR_ , 2017.
  * Xie et al. [2020b] Qizhe Xie, Zihang Dai, Eduard Hovy, Minh-Thang Luong, and Quoc V Le.  Unsupervised data augmentation for consistency training.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2020b.
  * Sohn et al. [2020] Kihyuk Sohn, David Berthelot, Chun-Liang Li, Zizhao Zhang, Nicholas Carlini, Ekin D. Cubuk, Alex Kurakin, Han Zhang, and Colin Raffel.  Fixmatch: Simplifying semi-supervised learning with consistency and confidence.  _arXiv_ , 2020.
  * Yang et al. [2021] Qize Yang, Xihan Wei, Biao Wang, Xian-Sheng Hua, and Lei Zhang.  Interactive self-training with mean teachers for semi-supervised object detection.  In _2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 5937–5946, 2021.  doi: 10.1109/CVPR46437.2021.00588.
  * Ivezić et al. [2019] Željko Ivezić, Steven M. Kahn, J. Anthony Tyson, Bob Abel, Emily Acosta, Robyn Allsman, David Alonso, Yusra AlSayyad, Scott F. Anderson, John Andrew, James Roger P. Angel, George Z. Angeli, Reza Ansari, Pierre Antilogus, Constanza Araujo, Robert Armstrong, Kirk T. Arndt, Pierre Astier, Éric Aubourg, Nicole Auza, Tim S. Axelrod, Deborah J. Bard, Jeff D. Barr, Aurelian Barrau, James G. Bartlett, Amanda E. Bauer, Brian J. Bauman, Sylvain Baumont, Ellen Bechtol, Keith Bechtol, Andrew C. Becker, Jacek Becla, Cristina Beldica, Steve Bellavia, Federica B. Bianco, Rahul Biswas, Guillaume Blanc, Jonathan Blazek, Roger D. Blandford, Josh S. Bloom, Joanne Bogart, Tim W. Bond, Michael T. Booth, Anders W. Borgland, Kirk Borne, James F. Bosch, Dominique Boutigny, Craig A. Brackett, Andrew Bradshaw, William Nielsen Brandt, Michael E. Brown, James S. Bullock, Patricia Burchat, David L. Burke, Gianpietro Cagnoli, Daniel Calabrese, Shawn Callahan, Alice L. Callen, Jeffrey L. Carlin, Erin L. Carlson, Srinivasan Chandrasekharan, Glenaver Charles-Emerson, Steve Chesley, Elliott C. Cheu, Hsin-Fang Chiang, James Chiang, Carol Chirino, Derek Chow, David R. Ciardi, Charles F. Claver, Johann Cohen-Tanugi, Joseph J. Cockrum, Rebecca Coles, Andrew J. Connolly, Kem H. Cook, Asantha Cooray, Kevin R. Covey, Chris Cribbs, Wei Cui, Roc Cutri, Philip N. Daly, Scott F. Daniel, Felipe Daruich, Guillaume Daubard, Greg Daues, William Dawson, Francisco Delgado, Alfred Dellapenna, Robert de Peyster, Miguel de Val-Borro, Seth W. Digel, Peter Doherty, Richard Dubois, Gregory P. Dubois-Felsmann, Josef Durech, Frossie Economou, Tim Eifler, Michael Eracleous, Benjamin L. Emmons, Angelo Fausti Neto, Henry Ferguson, Enrique Figueroa, Merlin Fisher-Levine, Warren Focke, Michael D. Foss, James Frank, Michael D. Freemon, Emmanuel Gangler, Eric Gawiser, John C. Geary, Perry Gee, Marla Geha, Charles J. B. Gessner, Robert R. Gibson, D. Kirk Gilmore, Thomas Glanzman, William Glick, Tatiana Goldina, Daniel A. Goldstein, Iain Goodenow, Melissa L. Graham, William J. Gressler, Philippe Gris, Leanne P. Guy, Augustin Guyonnet, Gunther Haller, Ron Harris, Patrick A. Hascall, Justine Haupt, Fabio Hernandez, Sven Herrmann, Edward Hileman, Joshua Hoblitt, John A. Hodgson, Craig Hogan, James D. Howard, Dajun Huang, Michael E. Huffer, Patrick Ingraham, Walter R. Innes, Suzanne H. Jacoby, Bhuvnesh Jain, Fabrice Jammes, M. James Jee, Tim Jenness, Garrett Jernigan, Darko Jevremović, Kenneth Johns, Anthony S. Johnson, Margaret W. G. Johnson, R. Lynne Jones, Claire Juramy-Gilles, Mario Jurić, Jason S. Kalirai, Nitya J. Kallivayalil, Bryce Kalmbach, Jeffrey P. Kantor, Pierre Karst, Mansi M. Kasliwal, Heather Kelly, Richard Kessler, Veronica Kinnison, David Kirkby, Lloyd Knox, Ivan V. Kotov, Victor L. Krabbendam, K. Simon Krughoff, Petr Kubánek, John Kuczewski, Shri Kulkarni, John Ku, Nadine R. Kurita, Craig S. Lage, Ron Lambert, Travis Lange, J. Brian Langton, Laurent Le Guillou, Deborah Levine, Ming Liang, Kian-Tat Lim, Chris J. Lintott, Kevin E. Long, Margaux Lopez, Paul J. Lotz, Robert H. Lupton, Nate B. Lust, Lauren A. MacArthur, Ashish Mahabal, Rachel Mandelbaum, Thomas W. Markiewicz, Darren S. Marsh, Philip J. Marshall, Stuart Marshall, Morgan May, Robert McKercher, Michelle McQueen, Joshua Meyers, Myriam Migliore, Michelle Miller, David J. Mills, Connor Miraval, Joachim Moeyens, Fred E. Moolekamp, David G. Monet, Marc Moniez, Serge Monkewitz, Christopher Montgomery, Christopher B. Morrison, Fritz Mueller, Gary P. Muller, Freddy Muñoz Arancibia, Douglas R. Neill, Scott P. Newbry, Jean-Yves Nief, Andrei Nomerotski, Martin Nordby, Paul O’Connor, John Oliver, Scot S. Olivier, Knut Olsen, William O’Mullane, Sandra Ortiz, Shawn Osier, Russell E. Owen, Reynald Pain, Paul E. Palecek, John K. Parejko, James B. Parsons, Nathan M. Pease, J. Matt Peterson, John R. Peterson, Donald L. Petravick, M. E. Libby Petrick, Cathy E. Petry, Francesco Pierfederici, Stephen Pietrowicz, Rob Pike, Philip A. Pinto, Raymond Plante, Stephen Plate, Joel P. Plutchak, Paul A. Price, Michael Prouza, Veljko Radeka, Jayadev Rajagopal, Andrew P. Rasmussen, Nicolas Regnault, Kevin A. Reil, David J. Reiss, Michael A. Reuter, Stephen T. Ridgway, Vincent J. Riot, Steve Ritz, Sean Robinson, William Roby, Aaron Roodman, Wayne Rosing, Cecille Roucelle, Matthew R. Rumore, Stefano Russo, Abhijit Saha, Benoit Sassolas, Terry L. Schalk, Pim Schellart, Rafe H. Schindler, Samuel Schmidt, Donald P. Schneider, Michael D. Schneider, William Schoening, German Schumacher, Megan E. Schwamb, Jacques Sebag, Brian Selvy, Glenn H. Sembroski, Lynn G. Seppala, Andrew Serio, Eduardo Serrano, Richard A. Shaw, Ian Shipsey, Jonathan Sick, Nicole Silvestri, Colin T. Slater, J. Allyn Smith, R. Chris Smith, Shahram Sobhani, Christine Soldahl, Lisa Storrie-Lombardi, Edward Stover, Michael A. Strauss, Rachel A. Street, Christopher W. Stubbs, Ian S. Sullivan, Donald Sweeney, John D. Swinbank, Alexander Szalay, Peter Takacs, Stephen A. Tether, Jon J. Thaler, John Gregg Thayer, Sandrine Thomas, Adam J. Thornton, Vaikunth Thukral, Jeffrey Tice, David E. Trilling, Max Turri, Richard Van Berg, Daniel Vanden Berk, Kurt Vetter, Francoise Virieux, Tomislav Vucina, William Wahl, Lucianne Walkowicz, Brian Walsh, Christopher W. Walter, Daniel L. Wang, Shin-Yawn Wang, Michael Warner, Oliver Wiecha, Beth Willman, Scott E. Winters, David Wittman, Sidney C. Wolff, W. Michael Wood-Vasey, Xiuqin Wu, Bo Xin, Peter Yoachim, and Hu Zhan.  LSST: From Science Drivers to Reference Design and Anticipated Data Products.  _The Astrophysical Journal_ , 873(2):111, March 2019.  doi: 10.3847/1538-4357/ab042c.
  * Kessler et al. [2009] Richard Kessler, Joseph P. Bernstein, David Cinabro, Benjamin Dilday, Joshua A. Frieman, Saurabh Jha, Stephen Kuhlmann, Gajus Miknaitis, Masao Sako, Matt Taylor, and Jake Vanderplas.  SNANA: A Public Software Package for Supernova Analysis.  _Proceedings of the Astronomical Society of the Pacific_ , 121(883):1028, September 2009.  doi: 10.1086/605984.
  * Kessler et al. [2019] R. Kessler, G. Narayan, A. Avelino, E. Bachelet, R. Biswas, P. J. Brown, D. F. Chernoff, A. J. Connolly, M. Dai, S. Daniel, R. Di Stefano, M. R. Drout, L. Galbany, S. Gonzá lez-Gaitán, M. L. Graham, R. Hložek, E. E. O. Ishida, J. Guillochon, S. W. Jha, D. O. Jones, K. S. Mandel, D. Muthukrishna, A. O’Grady, C. M. Peters, J. R. Pierel, K. A. Ponder, A. Prša, S. Rodney, and V. A. Villar and.  Models and simulations for the photometric LSST astronomical time series classification challenge (PLAsTiCC).  _Publications of the Astronomical Society of the Pacific_ , 131(1003):094501, jul 2019.  doi: 10.1088/1538-3873/ab26f1.  URL https://doi.org/10.1088%2F1538-3873%2Fab26f1.
  * Donoso-Oliva et al. [2023] C. Donoso-Oliva, I. Becker, P. Protopapas, G. Cabrera-Vives, M. Vishnu, and H. Vardhan.  ASTROMER.  _Astronomy & Astrophysics_, 670:A54, feb 2023.  doi: 10.1051/0004-6361/202243928.  URL https://doi.org/10.1051%2F0004-6361%2F202243928.
  * Zhou et al. [2021] Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang.  Informer: Beyond efficient transformer for long sequence time-series forecasting.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 35, pages 11106–11115, 2021.
  * Li et al. [2021] Yang Li, Si Si, Gang Li, Cho-Jui Hsieh, and Samy Bengio.  Learnable fourier features for multi-dimensional spatial positional encoding.  _Advances in Neural Information Processing Systems_ , 34:15816–15829, 2021.
  * HaoChen et al. [2021] Jeff Z. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma.  Provable guarantees for self-supervised deep learning with spectral contrastive loss.  _arXiv preprint arXiv:2106.04156_ , 2021.

##  Appendix A Additional Dataset Details

###  A.1 AstroClassification, Redshifts Datasets

The AstroClassification and Redshifts datasets were adapted from the 2019
Photometric LSST Astronomical Time-Series Classification Challenge [20]
333https://zenodo.org/record/2539456. This diverse dataset contains 14 types
of astronomical time-varying objects, simulated using the expected instrument
characteristics and survey strategy of the upcoming Legacy Survey of Space and
Time [LSST 48] conducted at the Vera C. Rubin Observatory. It includes two
overall categories of time-series objects: transients, short-lived events such
as supernovae, and variable sources, those with fluctuating brightness such as
pulsating stars. Specifically, the dataset includes the following transients:
type Ia supernovae (SNIa), SNIax, SNIa-91bg, SNIbc, SNII, superluminous
supernovae (SLSN), tidal disruption events (TDE), and single lens microlensing
events ($\mu$Lens-Single); and the following variable objects: active galactic
nuclei (AGN), Mira variables, eclipsing binary systems (EB), and RR Lyrae
(RRL).

Millions of potential new objects are discovered per observing night, and
important metadata such as object type, redshift, or other physical
parameters, require astronomers to take time-intensive spectra of each object.
Spectra are a granular brightness vs. wavelength measurement at a single point
in time, and are typically only taken for bright, nearby objects which require
less exposure time than faint, faraway objects. The vast majority of
discovered objects, however, will not have spectra but instead a time series
of imaging data taken in 6 broad wavelength ranges, or photometric bands. The
time-varying behavior of these objects in these coarse wavelength bands does
offer important clues about these physical parameters, but expert
interpretation of spectra are traditionally required for confident labeling.
Thus, our labeled training data for both AstroClassification and Redshifts
come from the unrepresentative subset of objects with spectra.

In these tasks, we are specifically interested in predicting the object type
(e.g. type II supernova) and the cosmological redshift of objects in the
unlabeled dataset. Cosmological redshift is a proxy for distance in the
universe, and an important piece of metadata for understanding an object’s
physical processes as well as other applications, such as estimating the
expansion rate of the universe with type Ia supernovae.

#### Problem Setting.

The task is to predict object type for AstroClassification (redshift for
Redshifts) from time-series of object brightness. The input $x$ consists of
flux measurements and associated uncertainties at times ${\bm{t}}$ and
photometric band that each measurement was taken in ${\bm{b}}$:
$\\{F(t_{i},b_{j})\\}_{i=1,j=1}^{T,W},\\{F_{\text{err}}(t_{i},b_{j})\\}_{i=1,j=1}^{T,W}$.
For this work, we map each $b\in{\bm{b}}$ to the central wavelength of the $b$
band, which we denote ${\bm{w}}$. The domain $d$ is binary, corresponding to
whether the object has a spectrum (and thus a label). The labels $y$ are
available only for objects with spectra, and are one of 14 types of
astronomical time-varying objects for AstroClassification (redshift of the
object for Redshifts). We seek to optimize performance on the unlabeled data,
which are generally fainter and further away than the labeled subset. We
evaluate on these examples as well as held-out examples from the labeled
subset.

#### Data.

The training set of 7,846 objects is designed to emulate a sample of objects
with spectra and thus biased toward brighter, more nearby objects compared to
the test set of 3,492,888 objects. A random subset of 10,000 test set objects
was selected for evaluation.

  1. 1.

Source: 6,274 objects

  2. 2.

ID Test: 782 objects

  3. 3.

OOD Test: 10,000 objects

All data were simulated with the SuperNova ANAlysis [SNANA, 49] software
library. Further details about the astrophysical models and LSST instrument
characteristics used in the simulation can be found in Kessler et al. [50].

##  Appendix B Data Augmentations

###  B.1 Generic Augmentations for Pretraining

#### AstroClassification and Redshifts.

For the AstroClassification and Redshifts datasets, we randomly mask a subset
of the input sequence using the masked language modeling paradigm introduced
by [24]. Given an unlabeled input sequence $x$, a training input $x^{\prime}$
can be generated by randomly masking elements of $x$ while the associated
label $y$ consists of the original, unmasked values. The model is trained to
use contextual information (unmasked elements) to successfully reconstruct
most of the sequence. From our ablation experiments, we find that a masking
percentage of 60% produces the best downstream results. We follow an existing
implementation for astronomical time-series [51] and set 80% of the masked
elements to 0, replace 10% with a random element from the sequence, and keep
the remaining 10% unchanged.

#### iWildCam and Camelyon17.

For iWildCam-WILDS, we use a ResNet-50 model pretrained on ImageNet with SwAV,
a contrastive learning algorithm [6]. For Camelyon17-WILDS, we use a
DenseNet121 pretrained with SwAV on the unlabeled Camelyon17-WILDS dataset
from Sagawa et al. [9]. SwAV uses random cropping augmentations of different
resolutions.

###  B.2 Targeted Augmentations for Fine-Tuning

#### Redshifting for AstroClassification and Redshifts.

Figure 4: Redshift distributions of source, augmented, and target datasets for
the AstroClassification and Redshifts tasks.

The OOD test set of the AstroClassification and Redshifts datasets have many
more high redshift objects than the source dataset, leading us to adopt an
augmentation scheme to alleviate this shift. Figure 4 shows the redshift
distributions of the source, augmented, and target datasets. Redshifting
places each object at a new redshift and recomputes its light curve sampling,
fluxes, and flux uncertainties accordingly. This augmentation algorithm was
adapted from Boone [21].

An input ${\bm{X}}\in\mathbb{R}^{T\times W}$ is a multivariate time series of
flux values at specified times and observed wavelengths,
$\\{F(t_{i},w_{j})\\}_{i=1,j=1}^{T,W}$. We also have
${\bm{X}}_{\text{err}}\in\mathbb{R}^{T\times W}$, representing the flux errors
corresponding to each element of ${\bm{X}}$. We denote the elements of
${\bm{X}}^{\prime}_{\text{err}}$ by
$\\{F_{\text{err}}(t_{i},w_{j})\\}_{i=1,j=1}^{T,W}$.Our goal is to model
$F,F_{\text{err}}:\mathbb{R}\times\mathbb{R}\rightarrow\mathbb{R}$ at a new
chosen redshift, $z^{\prime}$, to produce augmented inputs
${\bm{X}}^{\prime},{\bm{X}}^{\prime}_{\text{err}}$.

  * •

We first construct a distribution from which to sample the new redshift,
taking into account the current redshift of the object $z_{\text{orig}}$ as
well as the target redshift distribution. We then sample a new redshift,
$z^{\prime}\sim\text{loguniform}(0.95z_{\text{orig}},\;\text{min}(1.5(1+z_{\text{orig}})-1,\;5z_{\text{orig}}))$.

  * •

We fit a Gaussian process (GP) model for $F$ with training observations
${\bm{X}}$ queried at the training input values $({\bm{t}},{\bm{w}})$, and
denote the predictive mean and variance of the GP as
$F^{\prime},F^{\prime}_{\text{err}}$.

  * •

Given the new redshift value $z^{\prime}$, we rescale the timestamps and
wavelengths of the original observations to account for the physical effects
of the new redshift value:
${\bm{t}}_{\text{new}}=\frac{1+z^{\prime}}{1+z_{\text{orig}}}{\bm{t}}$,
${\bm{w}}_{\text{new}}=\frac{1+z^{\prime}}{1+z_{\text{orig}}}{\bm{w}}$. We
also randomly drop out 10% as well as a large swath of
$({\bm{t}}_{\text{new}},{\bm{w}}_{\text{new}})$ to simulate distinct observing
seasons (telescope observing only occurs in the winter).

  * •

We obtain GP predictions at test inputs
$\\{F^{\prime}(t_{\text{new},i},w_{\text{new},j})\\}_{i=1,j=1}^{T,W}$,
$\\{F^{\prime}_{\text{err}}(t_{\text{new},i},w_{\text{new},i})\\}_{i=1,j=1}^{T,W}$
and scale them by the log ratio of the new and original distances:

$\tilde{{\bm{X}}^{\prime}}=10^{0.4(d(z^{\prime})-d(z_{\text{orig}}))}\\{F^{\prime}(t_{\text{new},i},w_{\text{new},j})\\}_{i=1,j=1}^{T,W},$
$\tilde{{\bm{X}}^{\prime}}_{\text{err}}=10^{0.4(d(z^{\prime})-d(z_{\text{orig}}))}\\{F^{\prime}_{\text{err}}(t_{\text{new},i},w_{\text{new},j})\\}_{i=1,j=1}^{T,W},$

where $d(z)$ is the distance corresponding to redshift $z$.

  * •

We roughly model the observational noise of the telescope from the target data
as a function of wavelength and sample $\epsilon\in\mathbb{R}^{W}$ from it. We
define

${\bm{X}}^{\prime}=\\{\tilde{{\bm{X}}^{\prime}}_{:,j}+\epsilon_{j}\\}_{j=1}^{W},{\bm{X}}^{\prime}_{\text{err}}=\left\\{\sqrt{\tilde{{\bm{X}}^{\prime}}_{\text{err},:,j}^{2}+\epsilon_{j}^{2}}\right\\}_{j=1}^{W}.$

  * •

We model the observational capabilities of the telescope to ensure that our
augmented input ${\bm{X}}^{\prime},{\bm{X}}^{\prime}_{\text{err}}$ does not
fall below the threshold of detection. We “accept" an augmented input
${\bm{X}}^{\prime},{\bm{X}}^{\prime}_{\text{err}}$ if the signal-to-noise
ratio (SNR) of at least two observations is over 5, i.e.
$\text{SNR}({\bm{X}}^{\prime}_{i,j},{\bm{X}}^{\prime}_{\text{err},i,j})\geq 5$
for at least 2 of $i\in\\{1,...,T\\},j\in\\{1,...,W\\}$. We define
$\text{SNR}(x,x_{\text{err}})=\frac{|x|}{x_{\text{err}}}$.

#### Copy-Paste (Same Y) for iWildCam.

This augmentation strategy randomizes the backgrounds of wildlife images to
reduce the model’s dependence on these spurious features for species
classification. Specifically, a segmentation mask is applied to each image to
separate the animal from the background, and the animal is “copy-pasted" into
a new background from a camera that has observed that animal species. This was
the best performing augmentation strategy from Gao et al. [17].

#### Stain Color Jitter for Camelyon17.

This augmentation, originally from Tellez et al. [29], alters the pixel values
of the slide images to emulate different staining procedures used by different
hospitals. The augmentation uses a pre-specified Optical Density (OD) matrix
to project images from RGB space to a three-channel hematoxylin, eosin, and
DAB space before applying a random linear combination. This was the best
performing augmentation strategy from Gao et al. [17].

##  Appendix C Experimental Details

#### AstroClassification and Redshifts.

For AstroClassification and Redshifts, we pretrain with a masked autoencoding
objective:

$\displaystyle\mathcal{L}_{\text{MAE}}(\phi)=\mathbb{E}_{x\sim
P_{U},{x^{\prime}}\sim\mathcal{A}_{\text{pre}}(\cdot\mid
x)}[(\phi({x^{\prime}})-x)^{2}]$ (4)

We use an encoder-only Informer model [52] with 8 encoder layers of 12
attention heads each. The model hidden dimension was chosen to be 768 and the
layer MLPs have hidden dimension 256. Due to the 2-dimensional position data
(each element of the time-series has an associated time and photometric
band/wavelength) and irregular sampling of our dataset, we train a positional
encoding based on learnable Fourier features following Li et al. [53]. We also
select a random window of length 300 from each example (and zero-pad examples
with fewer than 300 observations) to produce inputs of uniform shape. We
perform pretraining with a batch size of 256 and learning rate 1e-4 (selected
from 1e-3 $\sim$ 1e-6) for 75,000 steps. We finetune the pretrained model with
linear probing for 20,000 steps (for pretrained models only) and learning rate
1e-4, then fine-tuning for 10,000 steps at learning rate of 4e-5. We increase
the learning rate for models without pretraining to 1e-4 for FT. The Redshifts
task uses LP learning rate of 5e-4 and FT learning rate of 1e-4. We decrease
the learning rate per step with a linear scheduler.

#### iWildCam.

For pretraining, we use ResNet-50 pretrained on ImageNet with SwAV [6]. During
fine-tuning, we train all models for 15 epochs with early stopping on OOD
validation performance, following Gao et al. [17]. For pretrained models, we
also do 10 epochs of linear probing before fine-tuning [LP-FT, 28] for 15
epochs, where the linear probe is trained with Adam and the linear probe
weights used to initialize the fine-tuning stage is chosen with OOD validation
performance. To reduce the noise in OOD results, for all methods we select the
epoch in the last 5 epochs with the best OOD validation performance and report
OOD test results with that version of the model. Following Gao et al. [17], we
allow for 10 hyperparameter tuning runs, where we sample the following
hyperparameters independently from the following distributions: the linear
probe learning rate ($10^{\text{Uniform}[-3,-2]}$), fine-tuning learning rate
($10^{\text{Uniform}[-5,-2]}$), and probability of applying the augmentation
($\text{Uniform}[0.5,0.9]$) and pick the hyperparameter configuration with the
best OOD validation performance. For ERM and ERM+targeted augmentations, we
use the tuned hyperparameters from Gao et al. [17]. To decrease the confidence
interval due to an outlier seed, the reported performance of Connect Later is
averaged over 15 seeds. All other results are averaged over 5 seeds.

#### Camelyon17.

For pretraining, we use DenseNet121 pretrained on the unlabeled
Camelyon17-WILDS dataset presented in Sagawa et al. [9] with SwAV [6]. During
fine-tuning, we train all models for 15 epochs with early stopping on OOD
validation performance, following Gao et al. [17]. For pretrained models, we
also do 10 epochs of linear probing before fine-tuning [LP-FT, 28] for 15
epochs, where the linear probe is trained with Adam and the linear probe
weights used to initialize the fine-tuning stage is chosen with OOD validation
performance. To reduce the noise in OOD results, for all methods we select the
epoch with the best OOD validation performance and report OOD test results
with that version of the model. Following Gao et al. [17], we allow for 10
hyperparameter tuning runs, where we sample the following hyperparameters
independently from the following distributions: the linear probe learning rate
($10^{\text{Uniform}[-3,-2]}$), fine-tuning learning rate
($10^{\text{Uniform}[-5,-2]}$), probability of applying the augmentation
($\text{Uniform}[0.5,0.9]$), and augmentation strength
($\text{Uniform}[0.05,0.1]$), and pick the hyperparameter configuration with
the best OOD validation performance. All results are averaged over 20 seeds.

##  Appendix D Empirical Estimates of Connectivity

Table 7: Empirically estimated connectivity measures for iWildCam-WILDS, AstroClassification, and Camelyon17-WILDS. iWildCam-WILDS and AstroClassification results are averaged over 15 randomly selected class-domain pairs, while Camelyon17-WILDS results are averaged over all possible class-domain pairs. | across-domain | across-class | across-both
---|---|---|---
iWildCam-WILDS | 0.116 | 0.071 | 0.076
AstroClassification | 0.287 | 0.159 | 0.097
Camelyon17-WILDS | 0.16 | 0.198 | 0.152

We empirically estimate connectivity measures for all of the datasets we
tested on following the procedure outlined in Appendix D of Shen et al. [7].
Specifically, we train binary classifiers from scratch to predict the class-
domain pair of a given input example. We randomly select 15 class-domain pairs
for iWildCam-WILDS and AstroClassification, while for Camelyon17-WILDS we use
all class-domain pairs since Camelyon17-WILDS is a binary classification task.
We label these class-domain examples following Appendix D of Shen et al. [7]
and create a dataset with 80/10/10 train/validation/test split. We train using
the same hyperparameters described in Appendix C for 3,000 steps with early
stopping on the validation accuracy. Our results are presented in Table 7.

##  Appendix E Simple construction where Connect Later improves over
pretraining or targeted augmentations alone

We give a simple construction for constrastive pretraining based on the
construction in Proposition 3 (Appendix A.2) of Shen et al. [7], where Connect
Later improves over pretraining (standard fine-tuning) or targeted
augmentations alone.

#### Data distribution.

We consider binary classification with 2 domains. Let
$\mathcal{S}=\\{x\in\mathcal{X}:d_{x}=1\\}$ and
$\mathcal{T}=\\{x\in\mathcal{T}:d_{x}=2\\}$, and assume that $P_{S}$ and
$P_{T}$ are uniform over $\mathcal{S}$ and $\mathcal{T}$. The unlabeled
distribution for pretraining is the uniform distribution over $\mathcal{X}$.
The source domain $\mathcal{S}=\\{1,2\\}$ contains 2 points and the target
domain $\mathcal{T}=\\{3,4,5,6,7,8\\}$ contains 6 points. For simplicity, we
let the labels $y_{x}$ be a deterministic function of the input $x$. The label
space is $\mathcal{Y}=\\{-1,1\\}$. The label for $x\in\\{1,3,5,7\\}$ is
$y_{x}=1$ and the label for $x\in\\{2,4,6,8\\}$ is $y_{x}=-1$. Only the source
data is labeled.

#### ERM with targeted augmentations.

ERM with targeted augmentations learns a model on source labeled data. To
specialize to this section, the ERM objective is

$\displaystyle\mathcal{L}_{\text{ERM}}(f)=\mathbb{E}_{x\sim
P_{S},x^{\prime}\sim\mathcal{A}_{\text{ft}}(\cdot\mid
x)}[\ell(f(x^{\prime}),y_{x})].$ (5)

ERM returns a classifier
$\widehat{f}_{\text{erm}}\in\operatorname*{arg\,min}_{f}\mathcal{L}_{\text{ERM}}(f)$.

#### Spectral contrastive learning.

Following HaoChen et al. [54] and Shen et al. [7], we analyze contrastive
learning from an augmentation graph perspective, where inputs $x$ are
connected via augmentations with edge weights $S_{+}(x,{x^{\prime}})$, which
represent the probability of $x,{x^{\prime}}$ being a positive pair
(augmentations of the same input $x$). For theoretical analysis, we analyze
the spectral contrastive learning objective:

$\displaystyle\mathcal{L}_{\text{pretrain}}(\phi)=-2\cdot$
$\displaystyle\mathbb{E}_{(x,x^{+})\sim
S_{+}}\left[\phi(x)^{\top}\phi(x^{+})\right]+\mathbb{E}_{x,x^{\prime}\sim
P_{U}}\left[\left(\phi(x)^{\top}\phi(x^{\prime})\right)^{2}\right].$ (6)

The result of pretraining to optimize the above objective is an encoder
$\widehat{\phi}:\mathcal{X}\to\mathbb{R}^{k}$.

Figure 5: Example distribution of data and augmentations for contrastive
learning where Connect Later improves OOD performance over contrastive
pretraining+standard fine-tuning and ERM+targeted augmentations. The
augmentation graph is similar to Shen et al. [7] except the edge weights
connecting 1,2 and 3,4 are swapped. The shapes represent classes, while the
labeled data is shaded in green. The generic augmentation probabilities are
marked as edge weights, where we assume that $\alpha>\gamma+\beta$. Here,
targeted augmentations which first swap inputs 1 and 2 before applying a
generic augmentation help to align the source and target. However, some target
inputs are not reachable via augmentations from source inputs. Standard fine-
tuning can generalize throughout the target domain, but only in conjunction
with targeted augmentations that align the source and target. The orange
dotted lines on the far ends connect to each other (the graph wraps around).

#### Linear probing (fine-tuning step).

Instead of analyzing fine-tuning, we follow Shen et al. [7] and analyze linear
probing on top of the pretrained representations from the encoder. We train a
linear model with parameters $B\in\mathbb{R}^{r\times k}$, where $r$ is the
number of classes. We minimize the objective:

$\displaystyle\mathcal{L}(B)=\mathbb{E}_{x\sim
P_{S}}\left[\ell(B\widehat{\phi}(x),y_{x})\right]+\eta\|B\|_{F}^{2},$ (7)

where $\ell$ is the squared loss and we take $y_{x}\in\mathbb{R}^{k}$ to be a
one-hot encoding of the class label. The resulting classifier is
$\widehat{f}(x)=\operatorname*{arg\,max}_{i\in[r]}(\widehat{B}\widehat{\phi}(x))_{i}$.

#### Pretraining augmentations (Figure 5)

We define the pretraining augmentation distribution
$\mathcal{A}_{\text{pre}}(\cdot\mid x)$ to be

$\displaystyle\mathcal{A}_{\text{pre}}({x^{\prime}}\mid
x)=\begin{cases}{\rho^{\prime}}&x={x^{\prime}}\\\
{\alpha^{\prime}}&\\{x^{\prime},x\\}\in\\{\\{1,4\\},\\{3,5\\},\\{5,7\\},\\{2,5\\},\\{4,6\\},\\{6,8\\},\\{1,8\\},\\{2,7\\}\\\
{\beta^{\prime}}&\\{{x^{\prime}},x\\}\in\\{\\{1,2\\},\\{3,4\\},\\{5,6\\},\\{7,8\\}\\}\\\
{\gamma^{\prime}}&\\{{x^{\prime}},x\\}\in\\{\\{1,3\\},\\{2,4\\},\\{3,6\\},\\{4,5\\},\\{5,8\\},\\{6,7\\},\\{1,7\\},\\{2,8\\}\\}\\\
\end{cases}.$ (8)

Notice that the weight between 1,3 is ${\gamma^{\prime}}$ and the weight
between 1,4 is ${\alpha^{\prime}}$, and the weights are similarly swapped for
2,4, and 2,5. We assume that
${\rho^{\prime}},{\alpha^{\prime}},{\beta^{\prime}}$, and ${\gamma^{\prime}}$
are in $(0,1)$ and are distinct. We also assume that the augmentation
probabilities satisfy
${\rho^{\prime}}>\max\\{{\alpha^{\prime}},{\beta^{\prime}}\\}$ and
$\min\\{{\alpha^{\prime}},{\beta^{\prime}}\\}>{\gamma^{\prime}}$. Following
Shen et al. [7], we can convert these to positive pair probabilities
${\rho},{\alpha},{\beta},{\gamma}$ with similar properties by renormalizing.

Given the above setting, the following is a simplified form of Proposition 3
from Shen et al. [7], if we instead use the following augmentation
distribution, which swaps the edge weight magnitudes that involve nodes 1 and
2:

$\displaystyle\mathcal{A}_{\text{prop}}({x^{\prime}}\mid
x)=\begin{cases}{\rho^{\prime}}&x={x^{\prime}}\\\
{\alpha^{\prime}}&\\{x^{\prime},x\\}\in\\{\\{1,3\\},\\{3,5\\},\\{5,7\\},\\{2,4\\},\\{4,6\\},\\{6,8\\},\\{1,7\\},\\{2,8\\}\\\
{\beta^{\prime}}&\\{{x^{\prime}},x\\}\in\\{\\{1,2\\},\\{3,4\\},\\{5,6\\},\\{7,8\\}\\}\\\
{\gamma^{\prime}}&\\{{x^{\prime}},x\\}\in\\{\\{1,4\\},\\{2,3\\},\\{3,6\\},\\{4,5\\},\\{5,8\\},\\{6,7\\},\\{1,8\\},\\{2,7\\}\\}\\\
\end{cases}.$ (9)

######  Proposition 1 (Shen et al. [7]).

With the above construction for the input space $\mathcal{X}$, unlabeled
distribution $P_{U}$, and data augmentation $\mathcal{A}_{\text{prop}}$, for
some feature dimension $k\in\mathbb{Z}^{+}$ a linear probe trained on
contrastive pre-trained features achieves 0 target error:
$\mathcal{L}_{0-1}(\widehat{f})=0$. However, for all $k\in\mathbb{Z}^{+}$,
there exists a minimizer $\widehat{f}_{\text{erm}}$ of the ERM objective (with
data augmentations according to $\mathcal{A}_{\text{prop}}$) that has non-zero
error: $\mathcal{L}_{0-1}(\widehat{f}_{\text{erm}})=1/3$.

#### ERM with targeted augmentations can get high OOD error.

In general, we proceed by defining the following targeted augmentation, which
allows us to reduce to the setting of Proposition 1:

$\displaystyle\mathcal{A}_{\text{ft}}({x^{\prime}}\mid
x)=\begin{cases}1&\\{{x^{\prime}},x\\}\in\\{1,4\\},\\{2,3\\}\\\
1&x={x^{\prime}}\text{ and }x\notin\\{1,2\\}\\\ 0&\text{otherwise}\end{cases}$
(10)

which transforms input 1 to 4 and the input 2 to 3, while keeping all other
inputs the same. Since the ERM with augmentations objective will not contain a
term involving inputs 5,6,7, or 8 and thus the prediction on these inputs do
not affect the objective, there exists a minimizer of the ERM objective
(Equation 5) that predicts the wrong label for inputs 5,6,7,8 and has target
error 2/3. This is because these nodes are unreachable via augmentations of
the source inputs, and thus the ERM objective can be minimized with any
arbitrary prediction on these inputs.

#### Standard fine-tuning has high OOD error.

By Proposition 1, standard fine-tuning after contrastive pretraining has zero
target (OOD) error when the pretraining augmentations do not have swapped
edges. By symmetry, standard fine-tuning (contrastive pretraining + linear
probing) on our augmentation graph with pretraining augmentations
$\mathcal{A}_{\text{pre}}$ outputs the opposite label for all target inputs,
resulting in an OOD error of 1. This is because the source and target domains
are misaligned in our augmentation graph.

#### Connect Later achieves zero OOD error.

Connect Later applies targeted augmentations $\mathcal{A}_{\text{ft}}$ during
the linear probing step (on top of contrastive pretrained representations).
This choice of targeted augmentations reduces to the setting of Proposition 1
where the labeled source domain consists of the inputs 3,4 instead. By the
symmetry of the graph and applying Proposition 1, Connect Later achieves 0 OOD
error.