Okapi: Generalising Better by Making Statistical Matches Match
Authors: Myles Bartlett
Sara Romiti
Viktoriia Sharmanska
Novi Quadrianto
Venue: NeurIPS 2022 Accept (2022)
Keywords: Domain Generalisation
Semi-Supervised Learning
Statistical Matching

Abstract: We propose Okapi, a simple, efficient, and general method for robust semi-supervised learning based on online statistical matching. Our method uses a nearest-neighbours-based matching procedure to generate cross-domain views for a consistency loss, while eliminating statistical outliers. In order to perform the online matching in a runtime- and memory-efficient way, we draw upon the self-supervised literature and combine a memory bank with a slow-moving momentum encoder. The consistency loss is applied within the feature space, rather than on the predictive distribution, making the method agnostic to both the modality and the task in question. We experiment on the WILDS 2.0 datasets Sagawa et al., which significantly expands the range of modalities, applications, and shifts available for studying and benchmarking real-world unsupervised adaptation. Contrary to Sagawa et al., we show that it is in fact possible to leverage additional unlabelled data to improve upon empirical risk minimisation (ERM) results with the right method. Our method outperforms the baseline methods in terms of out-of-distribution (OOD) generalisation on the iWildCam (a multi-class classification task) and PovertyMap (a regression task) image datasets as well as the CivilComments (a binary classification task) text dataset. Furthermore, from a qualitative perspective, we show the matches obtained from the learned encoder are strongly semantically related. Code for our paper is publicly available at https://github.com/wearepal/okapi/.

Okapi: Generalising Better by Making Statistical Matches Match
Myles B
Sara Romiti
Viktoriaa Sharmanska 1.2
Novi Quadrianto 13,44
Predictive Analytios Lab, University of Sussex Imperial College London BCAM Severo Ochoa Strategic Lab on Trustworthy Machine Learning
4 *Monash University, Indonesia
Abstract
We propose Okapi, simple, efficient, and general method for robust semi- supervised learning based on online statistical matching. Our method uses a nearest-neighhours- based matching procedure to generate cross-domain views for a consistency loss, while eliminating statistical outliers. In order to perform the online matching in runtime- and memory-efficient way, draw upon the self- supervised literature and combine memory bank with slow-moving momentum encoder. The consistency loss is applied within the feature space, rather than on the predictive distribution, making the method agnostic to both the modality and the task in question. We experiment on the WILDS 2.0 datasets [63), which significantly expands the range of modalities, applications, and shifts available for studying and benchmarking real-world unsupervised adaptation. Contrary to [63], we show that it is in fact possible to leverage additional unlabelled data to improve upon empirical risk minimisation (ERM) results with the right method. Our method outperforms the baseline methods in terms of out-of-distributio (OOD) generalisa- tion on the iWildCam (a multi-class classification task) and PovertyMap regres- sion task) image datasets as well as the CivilComments (a binary classification task) text dataset. Furthermore, from qualitative perspective, we show the matches obtained from the learned encoder are strongly semantically related. Code for our paper is publicly available at hhtps:: thub. com/ cwm/ http//eereaa//opp...
Introduction
Machine learning models have been deployed for safety-critical applications such as disease diagnosis 73 and self-driving cars (77), and in socially important contexts such as the allocation of healthcare, education, and credit (c.g. 23 35 ). Many machine learning algorithms, however, rely on supervision from large amount of labelled data, and are typically trained to exploit complex relationships and distant correlations present in the training dataset. This strategy has proven to be effective in the setting when have training (source) and test (target) data that are i..dd
In reality, machine learning models are often deployed on target data whose distribution different from the source distribution they were trained on. For example, in the task of classifying animal species in camera trap image, one aims learn model that can generalise to new camera trap locations despite variations in illumination, background, and label frequencies, given training examples from a limited set of camera trap locations. Exploiting correlations that only hold in these limited locations but not in the new locations can hurt -uittf-diitribution (OOD) generalisation. While we only have a small subset camera traps that have their images labelled, we have large
Coorrspponiing author: m. bat.lett@suseex.eet ac. uk.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

amount of unlabelled data from the other camera traps that capture diverse operating conditions. In general, unlabelled data is much more readily available than labelled data and can often be obtained from distributions beyond the source distribution. Taking advantage of these unlabelled data during training is a key element to build robust models that have good OOD performance without sacrificing in-distribution (ID) performance.
Our work is a direct response to the empirical conclusion of 663 for the WILDS 2.0 dataset that existing semi-supervised methods- - -everaging the unlabelled data provided in this extended version of the WILDS benchmark [4] fail to provide consistent benefit over the combination of judicious aattaaaammetttion and standard empirical risk minimisation (ERM). We show that with the right method, however, it is in fact possible to make effective use of large volumes of unlabelled data as supplement to a smaller set of labelled data, from a limited set of domains, to achieve strong generalisation to data from domains outside the training distribution. To develop this method, we turn to a statistical matching (SM) framework [57] (5, 60], a model-based approach for providing joint information on variables and indicators collected through multiple sources. SM has been widely utilised to assess the effect of interventions in numerous fields, such as education, medical and community policies (e.g. 10 In SM, intervened units are paired with control units and those units without sufficiently-gon estimating the treatment effect. In the running example of moiiall-spccess classification, intervened units may correspond to the limited set of camera trap locations that are aanttaenno,,,ed, while control units refer to the many more camera trap locations that are only partially annotated. Pairing is beneficial for capturing diverse operating conditions, yet the ability to drop unpaired units is crucial for mitigating the risk of statistically- -ptoo -poor matches corrupting the training signal.
match according to a given statistical criteria are excluded when
In designing an online method for statistically matching samples from different domains (camera-trap locations) and using this to define a consistency loss, we arrive at our proposed semi-supervised method, Okapi. This consistency loss is predicated on the simple idea of pulling together similar samples from different domains within the latent space of the encoder, and using this to bootstrap said encoder such that the distributions become progressively more aligned over the course of training. Since matching samples using the full dataset at each step of training is computationally infeasible, we instead approximate it using combination of wwwmmnntu---nnoodin has been well-proven in self-supervised learning [33 42). Compared with other consistency-bet based methods such FixMatch (67). Okapi has the advantage of being agnostic both the task and the modality, in addition to being distributionally robust. Contrary, to Sagawa et al. [63], we show that the supplementary unlabelled data and domain information can be leveraged by Okapi improve upon standard ERM on datasets from the WILDS 2.0 benchmark.
and memory-bank that
2 Preliminaries
2.1 Problem setting
In the standard supervised setting, one is given a dataset, Dj A {I J and trains a model, parameterised by 6, to well-approximate the empirical distribution as (o(g)a). Labelled data is limited by the cost of annotation yet one often has access far larger corpus of unlabelled data, Du & ( which can be used to supplement Dj. Semi-supervised learning is motivated by the idea that this additional data can often be used to improve the ID and/or OOD performance of po(y)æ:). We can view unsupervised domain adaptation (UDA) as a special case of semi cumissupervieed learning, where there is assumed to be some distribution shift (adverse to mnively-trained predictor) between D, and D,,. Here, D,, comes from the domain on which po(y|z) is to be evaluated, such that we have D. & DOOD. where DOoD denotes the target domain, that is OOD w.r.t. D;. In the most general sense, a domain, or environment 19 describes some partitioning of the data according to its source or some secondary characteristic, such as time of day, weather, location, lighting, or the model of the device used to collect said data; one would hope that a predictor trained under one set of conditions (e.g. day) would perform with minimal degradation under another set of conditions (e.g. night) when those conditions irrelevant the task hand.
Assuming the data follows the conditional generative distribution x ~ ((]]s), where s is the domain label, one would ideally use DOoD to learn invariance to the marginal distribution, p(s), and thereby achieve the equivalence po (y) (r) = Po (y) ar, 8). practice, one typically does not have access to DOOD but does have access to training data sourced from a mixture of domains which can be

---

leveraged to learn more general invariance that extends to those domains outside the training distribution 4. Such a learning paradigm is referred to as domain generalisation (DG). While some DG works consider the more extreme case of being unobserved [19] we follow the more conventional setup 4
[ in which the domain(s) associated with each sample (labelled and unlabelled) is indicated by the discrete label (set of labels) s. We denote the set of possible domains for the in-distribution labelled and unlabeled data, as St and S., respectively, and their union as S - S U S., Following the setup established in 633. D., is assumed to be unlabelled only w.r.t the targets and not w.r.t the domain labels and thus that both Dr and Dr can be augmented with the latter N2
to give the re-definitions Dj & {zi, Yi, 8: 3
No and D. {z1, 81J==11
2.2 Statistical matching
Statistical matching is a sampling strategy which aims to balance the distribution of the observed covariates in the treated and control groups. In general terms, observed covariates are measured characteristics of the samples; in our work we refer to the encodings generated by a deep neural network as covariates instead of the original characteristics. The treated and control groups are two partitions of the data; specifically, the treated group is the set of samples having a specific value of a variable of interest (here, the domain indicator, 8) and the control group is its complement.
In this work we utilise Nearest Neighbour (NN) matching, a distance-basee matching method that pairs sample i of the treated group with the closest sample, j belonging to the control group. A distance measure is used to define how close two samples, i and j. are, with propensity score distance (PSD) and Euclidean distance being two widely-used distances that we employ here indirectly (as means of filtering) and directly, respectively.
The propensity score distance is defined as the difference between propensity scores, and ej, of samples and j. i.e. PSD(E,, eạ) A lee ejl- In causal inference, the propensity score refers to the probability of sample belonging to the treated group, given its covariates z; [58]; in practice, this conditional probability is rarely known priori and thus requires estimation, typically via logistic regression [68]. We generalise the notion of a propensity score to categorical domains simply by modelling the conditional probability for each domain, with e, instead -im-eenmsii probability vector. The Suuclideanddistanee approach, in contrast, computes the distance between the covariates, Iy and 2tj, of given pair of samples. Despite PSD being the more prevalent of the two distances, it is ill-suited to cases in where pairs are close in value w.r.t. all covariates and in such cases Fuclidean distance should be preferred [40). Nevertheless, propensity scores remain a relevant component of NN-based matching for defining calipers that can reduce the likelihood of false-positive matches.
We make use of two types of caliper, fixed and standard deviation. The fixed caliper (20), ty. defines a region of common support between the estimated propensity score distribution of the two groups; only those samples within the feasible region are admissible for matching. For binary problems, the feasible region is symmetric such that we have {il e; E (1 (1-t) ty,ty)) whereas in the more general, categorical case the constraint is one-sided, i.e. {i ||e/|l>o < <,)). This selection rule helps by removing samples with extreme propensity scores. The standard deviation-based caliper (std-caliper), on the other hand, (59] defines the maximum discrepancy permitted between paired two samples. The discrepancy is usually expressed in terms of estimated PSD as |e, -e--e, eyl < a. tt. where a denotes the mean of the group-wise standard deviations of the propensity scores and to controls the percentage bias-reduction of the covariates, In the categorical case, we can simply substitute the absolute value for the infinity norm: l|e, llo a l.. In the following section, we describe how one can leverage this matching framework to define consistency loss encouraging inter-domain robustness.
3 Method
Here, we introduce Okapi, simple, efficient, and general (in the sense that it is applicable to any task or modality) method for robust semi-supervised learning based on online statistical matching. Our method belongs to the broad family of consistency- based methods, characterised by methods such as FixMatch [67], where the idea is to enforce similarity between a model's outputs for two views of an unlabelled image. These semi-supervised approaches based on minimising the discrepancy between two views of given sample are closely related with self-supervised methods based on instance discrimination [15] and self-distillation
6 [4 30). Many of the methods within this family, however,
3

---

are limited in applicability due to their dependence on recently has research into sulf-suuer sought to redress this problem with -ggality
transformations and only
alternatives such as MixUp 722, masking [6], and nearest-neighhours [24] 42 71. Approaches
such as FixMatch, AlphaMatch 28 and CSSL. 46 that enforce consistency between the predictive distributions suffer further from not being directly generalisable to tasks other than classification. Okapi addresses both of these aforementioned issues through 1) the use of a statistical matching procedure that call CaliperNN and detail in Sec. 3.2 to generate multiple views for given sample; 2) enforcing consistency between encodings rather than between predictive distributions.
We show that models trained to maximise the similarity between the encoding of given sample and those of its CaliperNNN-generatee match are significantly more robust to real-world distribution shifts than the baseline methods, while having the advantage of being both computationally efficient and agnostic to the modality and task question. Qualitatively speaking, we see that matches produced with the final model are related in semantically weatnnicclll--mmaaae eemantcallyymmaninnffuu ways. Furthermore, since the only constraint that samples be from different domains, the method is applicable whether information about the domain is coarse or fine-grained.
In the following subsections, we begin by giving general formulation of our proposed semi- supervised loss employing a generic cross-domain k-NN algorithm. We then explain how we can replace this algorithm with CaliperNN in order to mitigate the risk of poorly- matched samples, and how the loss may be computed in an online fashion to give our complete algorithm.
3.1 Enforcing consistency between cross-domain pairs
We view our predictor being composed of an encoder (or backbone) network, fo X - Rd generating intermediary outputs (features) & fo(z), and prediction head, 9d, such that the prediction for sample æ is given by ÿ _ Io o fe(x). We similarly consider the aggregate loss L as having a two-part decomposition given by
C c Loup + A Lunsups
(1)
where Lzup is the supervised component measuring the discrepancy (as computed, for example, by the MSE loss) between ÿ and the ground ground-truth label Cunsup is the unsupervised component based on some kind of pretext task, such as cross-view consistency, and À is a positive pre-factor determining the trade-off between the two components. For our method, we do not assume any particular form for Lxup and focus solely on Lunsup-
Given a pair of datasets Dj and Du, sourced from the labelled domain S₁, and unlabelled domain S, respectively, along with their union D - DD U D., our goal is to train a predictor that is robust (invariant) to changes in domain, including those unseen during training. To do this, we propose to regularise A (x) to be smooth (consistent) within local, cross-domain neighbourhoods. At a high-level, for any given query sample T.g sourced from domain 8q. we compute Lunsup as the mean distance between its encoding 2n A fo(x,,) and that of its k-nearest neighbours, V& (zz)) with the constraint that {*,} Ss, 0, where s,, the set of domain-labels associated with Va2 (2₂). The general form of this loss for a given sample can then be written as
(zgg A NN(2g: (fo(z) E E sg),
(2)
1
Cunsup
Σ (v.)
2n)
(3)
K
where d' Rd x Rd IR is some distance function. Here, we follow 30 and define d to be the squared Fuclidean distance between normalised encodings for our experiments. Allowing the NN algorithm to select pairs in an unconstrained manner, given the pool of queries and keys, however, can lead to poorly-matched pairs that are detrimental the optimisation process. To address this, replace the standard NN algorithm with propensity-score -coee-bb-e variant, inspired by the statistical matching framework (58).
3.2 Cross-domain matching
For the matching component of our algorithm, we propose to use variant of k-NN which, in addition to incorporating the above cross-domain constraint, filters the queries and keys that represent probable

---

CaliperNN
Su
Matched Pairs
Cross domaan constrant
Foed Calpper
Distance
Sid. Caliper
Illustration of our proposed statistical matching algorithm, CaliperNN. Given the anchor image encoding the corresponding domain label (we consider the binary case of labelled vs. unlabelled for simplicity), and propensity score e, CaliperNN outputs the closest samples subject to their being from different domains, following filtering by the fixed and std. calipers.
Figure 1:
outliers, according to their learned propensity scores. The initial stage filtering employs fixed caliper, where samples with propensity scores surpassing fixed confidence threshold are removed; this is followed by a second stage of filtering wherein any two samples (from different domains) can only be matched if the Euclidean distance between their respective propensity scores is below a pre-defined threshold (std-caliper). See Fig 3.2 for pictorial representation of these steps and Appendix x G for reference pseudocode.
The propensity score, for given sample is estimated p(s]=) using linear classifier fo. hy Rd A where AS is the probability simplex over possible domain labels, S, induced by the softmax function. KG is trained via maximum (weighted) likelihood to predict the domain label of given sample for all samples within the aggregate dataset D, (typically) a subset of it, encoded by fa. Since we apply both calipers to the learned propensity score, the shape of this distribution can have significant effect on the outcome of matching. Accordingly, we apply temperature-scaliin distribution. We denote the
with scalar 7 E I to sharpen or flatten the learned propensity-scne set of associated parameters tt T) the threshold for the fixed-caliper, the threshold for the std-caliper, and the temperature, respectively) as & and discuss in Appendix D how one can determine suitable values for these in practice.
For convenience we define the set of all encodings, given by fo. as z (fo(z) x D), the set of (he (a:) E z). and the set of associated domain labels
all associated propensity scores the offline case, the matches for D are then computed as
MatchedSamples é {(2, CaliperNN e, k)) z, e s)
with CaliperNN returning the set of k--earest neighbours according to d, subject to the aforementioned
cross-domain and caliper-based constraints. We allow for the fact that there may be valid matches for some samples due to these constraints; in such cases we have as the second element of their tuples, indicating that Cunsup should be set to O.
3.3 Scaling up with Online Learning
Re-encoding the dataset following each update of the feature-extractor in order to recompute MatchedSanples, is prohibitively expensive, with cost scaling linearly with N & N; + N,. More- over, CaliperNN requires explicit computation of the pairwise distance matrices, which can be prohibitive memory-wiss for large values of N. We address these problems using a fixed-size memory bank, MM storing only the last NM (where NM & N) encodings from a mlving momentum encoder [30, 33), for, which we refer to as the targer encoder, in line with [30] and accordingly refer to So as the online encoder. Unlike (30, however, we make use of neither projector nor a predictor head (in the case of the target encoder) in order to compute the inputs to the consistency loss and simply use the output of the backbone as is this is possible our setting due to Loup preventing representational collapse. More specifically, the target encoder's parameters,
5

---

8', are computed as a moving average of the online encoder's, 8, with decay rate Ç E (0, 1), per the recurrence relation
8-= +0
(5)
As the associated domain labels are also needed both for matching and to compute the loss for the propensity scorer, we also store the labels associated with M in companion memory bank M, We initialise M. and MM 0, resulting in fewer than NM samples being used during the initial stages of training when the memory banks are yet to be populated.
Each iteration of training, we sample batch of size B from D consisting of inputs x and s. During the matching phase, the inputs are passed through the targel encoder to obtain z _ (a)) x). serving as the queries for CaliperNN. We also experiment with simpler variant where the online encoder is instead used for this query-generation step, such that we instead have 2 _ {[f (z)) E x}. and find this can work equally well if Ç sufficiently high. The keys are then formed by combining the current queries with the past queries contained in the memory bank: 25 _ z U.M.MM The domain labels associated with 2₃ are likewise formed by concatenating the domain labels in the current batch with those stored in M Sy _ Sn U UMM. Once the matches for the current samples have been computed, the oldest B samples in M. and M are overwritten with Z₁ and Sk. respectively. The consistency loss is then enforced between each query 2g & (fe(a)l E x), according to the differentiable online encoder, and each of its matches, (2)) e CaliperNNg Zg. (Zz), providing that VA (=1) A ( (that is, under the condition that the estimated propensity score for 2n does not violate the caliper(s) and there are at least valid matches whose estimated propensity scores also do not), with the loss simply 0 otherwise. Since for is frozen, 2₁ carries an implicit soop-gradient and gradients are computed only w.r.t. 0. These steps are illustrated pictorially in Fig 2 and as pseudocode in Appendix G
Similarly, rather than solving for the optimal parameters, vi for the propensity scorer given the current values of Zą. which is infeasible for the large values of NM needed to wpll-approx the full dataset, we resort to a biased estimate of v. Namely, we train by in an online fashion to minimise the per-batch loss
Pps
Σ =EK
uea (2))))
(6)
Z
where H is the standard cross-entropy loss between the predictive distribution and the (degenerate) ground-truth distribution, given by the one-hot encoded domain labels, and Way S + R is a function assigning to each importance weight [66] based on the inverse of its frequency in sy to counteract label imbalance. In the special case in which the Dj and D. are known to have disjoint substitute their domain labels with 1 and O respectively (such that we have D, A {zi, Vi, 1),² and D., A {2;, 0) ) thus reducing the propensity scorer and CaliperNN to their binary forms. Knowing whether this condition is satisfied a priori (and thus whether the use of domain labels can be forgone completely from our pipeline) is not unrealistic: one may, for example know that two sets of satellite imagery cover two different parts of the world (e.g. Africa and Asia) yet not know the exact coordinates underlying their respective coverage.
support over S (that is, S n.S. 0),
Related Work
Domain Generalisation The goal of domain generalisation (DG) is produce models that robust to a wide range of distribution shifts (including those outside the training distribution), given a training set consisting of samples sourced from multiple domains. Despite the various techniques (many well theoretialllymotivaeee current methods continue to fall short in the face of natural distribution shifts AM] Indeed, ERM has repeatedly shown to be a strong baseline frequently outperforming dedicated methods that leverage domain information or additional unlabelled data for DG [3I] 63), despite the theoretical problems associated with using it when the training and test sets are misaligned. Until now, only pre-training on larger, more diverse datasets (with harder examples), has consistently proven to improve OOD generalisation, yet allowing pre-trained models to fit the ID data too closely can undo any such benefit conferred by the pre-training 3 39 69 74.. Similar to Okapi, MatchDG (51) draws upon causal matching to tackle DG. Despite the surface-level similarity, there are number of
designed to improve the generalisation of deep neural networks

---

Past + Current Encodings//omain Labels
Propensity Scorer
Location (s):
19
0
19
185
91
Zn
Target Encoder
CaliperNN
Location [s]:
EMA Updato
Lunoup I!
19
6:): Bephant
Online Encoder
Predictor
fe
Zq
Suup 30(3,))
oa + frunsup?
Figure 2: Overview of Okapi's learnig pipeline based using the iWildCam dataset for the sake of illustration. For simplicity, limit & to I so that the output of matching rather than a set of vectors; for the same reason we illustrate the process for only a single sample taken from the labelled data set D, annotated with both domain (s; in this case, camera location) and class (y) information. Inspired by recent advances self-supervised learning, we maintain a copy (the target encoder) of the online encoder, fe. whose parameters, O', average (EMA) of 0. This EMA update is performed at the beginning of each training set at a rate governed by the decay coefficient, Ç. For a given sample, we first compute its embedding using the target encoder to produce the query vector, z and by the online encoder to produce 2g. which will serve as the 'anchor' in the consistency loss. This query vector is then used alongside the output of the propensity scorer by CaliperNN to compute its cross-domain nearest neighbour, 2n, where the keys are taken be the current and past (stored in the Memory Bank) NM encodings of the data. The cross-domain constraint, prohibiting matching of samples belonging to the same domain, is denoted through a red coloring the location identifiers, the nearest sample obeying this constraint and the constraints of the calipers with purple highlighting. The consistency loss is the distance between by and zn- defined by function some distance function d. Finally, the supervised loss, Leup (here instantiated as the standard cross-entropy loss, H), is computed using the output of the predictor acting on bg and the ground-truth given by y.
single vector
exponential moving
significant differences, principally the respects that we consider semi-supervised DG (whereas MatchDG requires full-labeling w.r.t. y) and employ an augmented form of k-NN for bias-reduction in the absence of y.
Self-Supervised Learning In self-supervised learning (SeifSL), models trained solve pretext tasks constructed from the input data. This learning paradigm has led to significant breakthroughs in unsupervised learning in recent years, with performance now approaching (or even surpassing, along some axes such as adversarial robustness) that of supervised methods for many tasks while requiring significantly less labelled data. Due its generality, SelfSL has seen use across the complete spectrum of applications and modalities and underlies many of the foundation models [1] that have emerged in NLP 13 16 21) Computer Vision 22). and their intersection 2 78) Common pretext tasks include those based on the masked- masked-languuage waanuagge-moelllln
approach originally popularised by contrastive captioning (56 78), and instance discrimination and self-distillation [14 30 which rely on transformations
BERT (21 and recently generalised to other modalities [6] 7
of the data to generate multi-view inputs. Approaches belonging the latter two categories were originally limited by the fact that the transforms had to be tailored for a particular modality and for

---

Table 1: comparison between Okapi and different baselines on two benchmark image datasets. We include both the results of our re-run of the baselines and those of [63]. Both ID and OOD performances are reported. For iWildCam we average over results from 3 different seeds, for Pav we do so over the 5 pre-defined CV folds. Standard deviations are shown in parentheses.
Method
iWildCam
Mav
macro F1 f
worst U/R corr. f
worst U/R MSE 1
ID
OOD
ID
OOD
ID
00D
ERM 65
(1.4)
32.2(1..)
0.66 (0.04)
0.49 (0.06)
Fix.Match 66
46.3 (0.5)
(1.3)
0.54 (0.10)
0.30 (0.11)
ERM
48.6
33.3 (0.3)
0.72 (0.03)
0.53 (0.09)
0.23 (0.03)
0.35 (0.12)
FixMatch
51.1 (1.0)
(0.7)
0.50 (0.13)
0.34 (0.12)
0.59 (0.42)
0.88 (0.61)
Okapi (ours)
50.6 (0.7)
36.1 (0.9)
0.72 (0.02)
0.55 (0.10)
0.22 (0.02)
0.33 10)
Okapi (no calipers)
0.72 (0.02)
0.54 (0.12) 0.22 (0.02)
0.36 (0.14)
some modalities, such as tabular data, there is no obvious way to define them. A number of recent works have sought obviate this problem through the of MixUp [72], masking 6 (34, and k-NN 124 42 71], the latter of which is directly relevant to our work. Okapi bears closest resemblance to 42 in combining momentum-encooing with nearest-neighhourss lookup generate the views for BOOL-style (30] consistency loss. However, a key distinction lies in the use of an augmented form of neirhst-neighhburs, CaliperNN, which both constrains pairs of samples to being from different domains and filters out any queries or keys deemed outliers according to a learned propensity score.
Semi-Supervised Learning Semi-supervised learning (SemiSL) encompasses broad class of algorithms that combine unsupervised learning with supervised learning in order to improve the performance of the latter, especially when labelled data is limited. Many SemiSL methods are based on the self-training paradigm which can trace its roots back decades to the early work in pattern recognition by [65] and continues to be relevant in the modern era due its generality. both within SemiSL itself and in related fields such as domain adaptation (27), and fledgling field of SelfSL [4] discussed above. Self-training applies to any framework predicated on using a model's own predictions to produce pseudo-labels for the unlabelled data which can either be used as targets for self-distillation (55 or enforeing consistency between predictions that themselves have been perturbed [ 75 or that have been generated from perturbed/multi-vier inputs [67]. FixMatch method which has proven effective for semi-supervised classification, despite its simplicity, and various works (28) 46 have since built on the its framework prescribing the use weakly- and strongly-augment
667 is one example of a consistency-basee
inputs to generate the targets and predictions, respectively. Like these methods, Okapi also makes use of a cross-view consistency loss, however, the alternative views for a given sample are generated not through data-augmentationn but through statistical matching [58), with the aim being to achieve invariance to the domain rather than particular series of perturbations. Another example of particular relevance to our work is [70], which uses copy of the model with exponentially-average weights to generate the targets for the unlabelled data. Okapi also uses such a model to produce the targets for its consistency loss, but is more akin momentum-encoding
3 in the respect that the loss is imposed the latent space.
5 Experiments
5.1 Datasets
We evaluate Okapi three datasets taken from the WILDS 2.0 benchmark (63. These span variety of modalities and tasks, allowing us to showcase the generality of our proposed method (Okapi): iWildCam (images, multiclass classification), PovertyMap (multispectral images, regression), and CivilComments (text, binary classification). Details of each dataset can be found in Appendix
5.2 Image experiments
Results of our image-data experiments are summarised in Table [] Due to spacial constraints, we defer the full set of results, including those for the "offline" (w.rt. the matching) version of Okapi to Appendix. C For both datasets in question, the same metrics (63): macro-Fl for iWildCam
8

---

and worst-group (with the group defined as urban (U) vs. rural (R)) Pearson correlation for Poverty Map. For completeness, we include mean squared error (MSE) as a secondary metric for the latter dataset. Following 663, we compute the mean and standard deviation (shown in parentheses) over multiple runs for both and OOD test sets, with these runs conducted with 3 different random seeds and 5 pre-defined tross-validation folds for iWildCam and Pav, respectively.
We compare Okapi against two baselines, ERM and FixMatch 667 both according to our re- implementation and according to the original implementation given in [63 We note that since FixMatch, in its original form, is only applicable to classification problems due to its use of confidence- based thresholding, for the PovertyMap dataset, FixMatch represents a simplified variant (following [63]) without such thresholding, that is trained to simply minimise the MSE between all regressed values for the weakly- and strongly-augmente difference between the baselines run included in (63] and our re-runs is in the backbone architecture, with us opting for ConvNeXt 47 architecture over a ResNet one. For both datasets, and for both baselines we observe significant improvements stemming the change of backbone. Moreover, utilising ConvNeXt seems to be crucial in enabling FixMatch to surpass the ERM baseline in the classification task with 32.2 (ERM) vs 31.0 (FixMatch) and 33.3 (ERM) vs. 35.2 (FixMatch), with ResNet and ConvNeXt architecture respectively.
images. As described in Appendix D the main
Okapi, convincingly outperforms the baselines, w.r.t the OOD metric of interest, on both datasets. We observe an improvement of +0.9 macro F1, i.e. 36.1 1 vs 35.2 of Okapi and FixMatch (the best baseline for iWildCam) respectively. For the regression task in Pave Okapi achieves 0.55 and 0.33 on the OOD test set in terms of Pearson correlation and MSE, respectively, in contrast to the 0.53 and 0.33 of ERM. At the same time, we note that FixMatch fails to generalise well this task, yielding by far the worst results amongst the evaluated methods.
5.3 Text classification
Method
Civil Comments worst-group ace
Table 2: Comparison between Okapi and the baselines methods the Civil Comments dataset. We include both the original results of (63 well those of our reproduction of their ERM baseline. Performance is measured in terms of worst-group accuracy and averaged over seeds; standard deviations are shown in parentheses.
OOD
ERM 63
66.6 (1.6)
ERM (fully-aabelled)
63
69.4 (0.6)
ERM (reproduction)
68.5 (2.2)
Okapi (ours)
69.7 (2.0)
In Table 2 we summarise the numerical results for the CivilComments dataset. Remaining consistent with [63, we evaluate models according to the worst- group accuracy the minimum of the conditional accuracies obtained by conditioning on each of the 8 dimensions of s averaged over 5 replicates. Since there is no canonical ID test split available for this dataset, we report only the results only for the OOD split that is, rather than doing so for custom split avoid misrepresentation, We compare Okapi against both ERM variants featured in [X one trained on only the official labelled data and one trained with annotated unlabelled data (fully-labelled)- as well as our re-implementation of the ERM variant trained on only the labelled data with an identical hyperparameter configuration to the former. In contrast to the image datasets, we do not diverge in our choice of architecture, with all models trained with a pre-trained DistilBERT 664] backbone.
We observe marked improvement in the worst- worst-group accuracy of this baseline compared with that reported therein. We attribute this partly to the high variance of the model-selection procedure (inherited from (63] based on tidermtteetlyyccomputee consistently align with test performance) to determine the final model. This aside, we observe that Okapi outperforms the ERM baseline by a significant margin, to the point of parity with the fully-labelled baseline.
validation performance (which does not
5.4 Ablations and qualitatitive analysis
In order to evaluate the importance of the caliper-based filtering to the performance of Okapi, we perform an ablation experiment on PovertyMap dataset (Okapi (no calipers)) with said filtering disabled (and all else constant), such that instead of CaliperNN we have standard k-NN, albeit with

---

Imput
Matchh
Figure 3: Examples of input (labelled) images and their 1-NN matched (unlabelled) images retrieved using CaliperNN on iWildCam dataset. Here, we match images from the labelled-train set to images from the unlabelled-extra set, taking advantage the fact that their domains are disjoint.
the cross-group constraint still in place (per Eq. 2 We see that performance degrades according to both metrics of interest, and, crucially, that the standard deviation of the runs is significantly higher, in line with our expectation that filtering out poor matches should stabilise optimisation. We provide additional ablation experiments in Appendix FF) exploring the relative importance of the two (fixed and std-) calipers, the optimal number of neighbours to use for computing Cunsup, and the feasibility of using the online encoder to generate the queries for CaliperNN.
Finally, in Fig. 3 we show samples of matched pairs retrieved by CaliperNN from the encodings of the learned encoder for the iWildCam dataset. Here, we see that semantic information (encoding the
species of animal) is preserved across pairs, while nuisance factors such as illumination, background and contrast vary. Further examples from PovertyMap are shown in Appendix D In Appendix Hi
we include matching results for the PACS (photo (P), art painting (A), cartoon (C), and sketch (S))
dataset [45] demonstrating how temperature scaling, in conjunction with the fixed caliper, can be used control the filtering rate.
6 Conclusion
In this work, we introduced, Okapi, a semi-supervised method for training distributionally-roounn models that is intuitive, effective, and is applicable any modality task. Okapi is based on the simple idea of supplementing the supervised loss with cross-domain consistency loss that encourages the outputs of an encoder network to be similar for neighbouring (within the latent space of the encoder itself) samples belonging to different domains, which is made efficient using an online-learning framework. Rather than simply using k-NN with a cross-domain constraint, however, we propose an augmented form based on statistical matching (CaliperNN) that combines propensity scores with calipers winnow out low-quality matches; we find this to be important for both the end-performance and consistency of Okapi. Our work serves as response to 63], in that we find that is in fact possible to effectively incorporate unlabelled data and domain information into a training algorithm in order to improve upon ERM with respect an OOD test set, assuming an appropriate choice of architecture. Namely, on three datasets from the WILDS 2.0 benchmark, representing two different tasks (classification and regression) and modalities (image and text), show that Okapi outperforms both the ERM and FixMatch baselines according to the relevant OOD metrics.
Buoyed by these promising results, we intend to apply Okapi to other tasks (e.g. object detection and image segmentation) and other modalities (e.g. audio) to further establish its generality. Furthermore, one limitation of the current incarnation of the method is that the thresholds for the calipers are fixed over the course of training whereas may be beneficial to set these adaptively with the view to optimise such measures of inter-domain balance as Variance Ratio and Standard Mean Differences that are commonly used to evaluate the the goodness statistical matching procedures.
Acknowledgments and Disclosure of Funding
This research was supported by a European Research Council (ERC) Starting Grant for the project "Bayesian Models and Algorithms for Fairness and Transparency". funded under the European Union's 's Horizon 2020 Framework Programme (grant agreement no. 851538). NQ is also supported by the Basque Government through the BERC 2018-2021 program and by Spanish Ministry of Sciences, Innovation and Universities: BCAM Severo Ochoa accreditation SEV-2017-0718.
10

---

References
[1] A. Agarwal, A. Beygelzimer, M. Dudík, J. Langford, and H. Wallach. A reductions approach to fair classification. In International Conference Machine Learning, pages 60-69. PMLR, 2018.
[2] J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lene, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: a visual language model for few-shot learning. arXiv preprint arXiv:2204. 14198, 2022.
[3] A. Andreassen, Y. Bahri, B. Neyshabur, and R. Roelofs. The evolution of out-of-distribution robustness throughout fine-tuning, arXiv preprint (iXi::2006.183111 2021.
[4] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: 1907.02899 2019.
[5] Bachman, O. Alsharif, and D. Precup. Learning with pseudo-ensembles. Advances in neural information processing systems, 27, 2014.
[6] A. Baevski, W.-N. Hsu, Q. Xu, A. Babu, J. Gu, and M. Auli. Data2vec: A general framework for supersupe learning in speech, vision and language. arXiv preprint arXiv:2202.03555, 2022.
[7] H. Bao, I. Dong, and F. Wei. Beit: Bert pre-training of image transformers. arXiv preprint arXiv:2106.08254, 2021.
[8] S. Beery, E. Cole, and A. Gjoka. The iwildcam 2020 competition dataset. arXiv preprint arXiv:2004. 10340, 2020.
[9] Q Berthet, M. Blondel, O. Teboul, M. Cuturi, J.-P. Vert, and F. Bach. Learning with differen- tiable pertubed optimizers. Advances in neural information processing systems, 33:9508-9553 2020.
[10] A. Biglan, D. Ary, and A. C. Wagenaar. The value of interrupted time-series experiments for community intervention research. Prevention Science, 1(1):31-49, 2000.
[11] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint arXiv:2108.07258 2021.
[12] D. Borkan, L. Dixon, J. Sorensen, N. Thain, and L. Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. Compunion proceedings of the 2019 world wide web conference, pages 491-500, 2019.
[13] Brown, B. Mann, Ryder, M. Subbiah, J. D. Kaplan, Dhariwal, A. Neelakantan, Shyam. G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 13:1877-1901, 2020.
[14] M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 9650-9660, 2021.
[15] T. Chen, S. Komblith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[16] A. Chowdhery, S. Narang. J. Devlin, M. Bosma, G. Mishra, A. Roberts, Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint xrXiw:2204.02311, 2022.
[17] P. Christian, L. F. Murray-Kolb, S. K. Khatry, J. Katz, B. A. Schaefer, P. M. Cole, S. C. LeClerg, and J. M. Tielsch. Prenatal micronutrient supplementation and intellectual and motor function in carly school-aged children in nepal. Jama, 304(2))2716-272.) 2010.
11

---

[18] E. Creager, D. Madras, J.-H. Jacobsen, M. Weis, K. Swersky, T. Pitassi, and R. Zemel. Flexibly fair representation learning by disentanglement. In International conference on machine learning, pages 1436-1445. PMLR, 2019.
[19] E. Creager, J.-H. Jacobsen, and R. Zemel. Environment inference for invariant learning. In International Conference on Machine Learning. pages 2189 2990. PMLR, 2021.
[20] R. K. Crump, V. J. Hotz, G. W. Imbens, and O. A. Mitnik. Dealing with limited overlap in estimation of average treatment effects. Biometrika, 96(1):187-19, 2009.
[21] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprini 1040v088100048000 2018.
[22] M. Donini, L. Oneto, S. Ben-David, J. S. Shawe-Taylor and M. Pontil. Empirical risk mini- mization under fairness constraints. Advances in Neural Information Processing Systems, 31, 2018.
[23] J. A. Dunnmon, D. Yi, C. P. Langlotz, C. Ré, D. L. Rubin, and M. P. Lungren. Assessment of convolutional neural networks for automated classification of chest radiographs. Radiology, 290(2):537-544 2019.
[24] D. Dwibedi, Y. Aytar, J. Tompson, P. Sermanet, and A. Zisserman. With little help from my friends: Nearest-neighbor contrastive learning of visual representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 9588-9597, 2021.
[25] C. Dwork, M. Hardt, T. Pitassi, O. Reingold, and R. Zemel. Fairness through awareness. In Proceedings of the 3rd innovations in theoreticul computer science conference, pages 214-226, 2012.
[26] M. Feldman, S. A. Friedler, J. Moeller, C. Scheidegger, and S. Venkatasuuramanian, Certifying and removing disparate impact. In proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pages 259-268, 2015.
[27] Y. Ganin, E. Ustinova, H. Ajakan, P. Germain, H. Larochelle, F. Laviolette, M. Marchand, and Lempitsky. Comain-adversariaa training of neural networks. The journal of machine learning research, 17(1):2096-2033 2016.
[28] C. Gong, D. Wang, and Q. Liu. Alphamatch: Improving consistency for semi-supervised learning with alpha-divergence. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 13683-13692, 2021.
[29] P. Goyal, Q. Duval, I. Seessel, M. Caron, M. Singh, I. Misra, L. Sagun, A. Joulin, and P. Bo- janowski. Vision models are more robust and fair when pretrained on uncurated images without supervision. arXiv preprint 222.:22200.0866000 2022.
[30] J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. Richemond, E. Buchatskaya, C. Doersch, B. Avila Pires, Z. Guo, M. Gheshlaghi Azar, et al. Bootstrap your own latent-a new ap- proach to self-supervised learning. Advances in Neural Information Processing Systems, 33: 21271-21284, 2020.
[31] Guirajani and D. Lopez-Paz. In search of lost domain generalization. In International Conference on Learning Representations, 2020.
[32] M. Hardt, F. Price, and N. Srebro. Equality of opportunity in supervised learning. Advances in neural information processing systems, 29, 2016.
[33] K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattem recognition, pages
2020.
[34] K. He, X. Chen, S. Xie, Y. Li. P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. 11X.v.2111.06377, 2021.
[35] M. Hurley and J. Adebayo. Credit scoring the era of big data. Yale Journal of Law and Technology, 18:148-2111 1164111,16220 2017.
12

---

[36] B. Y. Idrissi, M. Arjovsky, M. Pezeshki, and D. Lopez-Paz. Simple data balancing achieves competitive panst-rroupacceurccy In Conference on Causal Learning and Reasoning, pages 336-351, PMLR, 2022.
[37] Kamiran and T. Calders. Data preprocessing techniques for classification without discrimina- tion. Knowledge and information systems, 33(1):1-33, 2012.
[38] Kehrenberg, M. Bartlett, O. Thomas, and N. Quadrianto. Null-sampling for interpretable and fair representations. In European Conference on Computer Vision, pages 565-580. Springer, 2020..
[39] D. Kim, K. Wang, S. Sclaroff, and K. Saenko. A broad study of pre-training for domain generalization and adaptation. arXiv preprint arXiv:2203.11819 2022.
[40] G. King and R. Nielsen. Why propensity scores should not be used for matching. Political Analysis, 27(4):435-444 2019.
[41] P. W. Koh, S. Sagawa, H. Marklund, S. M. Xie, M. Zhang, A. Balsubramani, W. Hu, M. Ya- sunaga, R. L. Phillips, I. Gao, et al. Wilds: A benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pages 5637-5664. PMLR, 2021.
[42] S. A. Koohpayegani, A. Tejankar, and H. Pirsiavash. Mean shift for self-supervised learning. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 10326- 10335, 2021.
[43] D. Krueger, E. Caballero, J.-H. Jacobsen, A. Zhang, J. Binas, D. Zhang, R. Le Priol, and A. Courville. distif--ist generalization via risk extrapolation (rex). In International Conference Machine Learning, pages 5815-5826. PMLR, 2021.
[44] P. Lahoti, K. Gummadi, and G. Weikum. Operationalizing individual faimess with pairwise fair representations. Proceedings of the VLDB Endowment, 13(4):506 518, 2019.
[45] D. Li, Y. Yang, Y.Z. Song, and T. M. Hospedales. Deeper, broader and artier domain gen- eralization. In Proceedings of the IEEE international conference computer vision, pages 5542-5550, 2017.
[46] J. Lienen and E. Hüllermeier. Credal self-supervised learning. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 14370-14382. Curran As- sociates, Inc., 2021. URL https::/ neurips wwwprripscccpppprr//aeaa pttpppaa///////2///ieeeee 7766:11:SFFOFF12227777700055555.
pdf
[47] Z. Liu, H. Mao, C.-Y. Wu, C. Feichtenhofer, T. Darrell, and S. Xie. A convnet for the 2020s. arXiv preprint MxXiv:2001.05555 2022.
[48] I. Loshchilov and F. Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th
International Conference on Learning Representations. ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://opeereviiee net/ htt/forummiddSSk@@8000
[49] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018
[50] D. Madras, F. Creager, T. Pitassi, and R. Zemel. Learning adversarially fair and transferable representations. In International Conference on Machine Learning, pages 3384-3393. PMLR, 2018.
International Conference on Machine Learning, pages 7313-7324. PMLR, 2021. [51] D. Mahajan, Tople, and A. Sharma. Domain generalization using causal matching. In
[52] K. Muandet, D. Balduzzi, and B. Schölkopf. Domain generalization via invariant feature representation. In International Conference on Machine Learning, pages 10-18. PMLR, 2013.
13

---

[53] L. Oneto, M. Donini, G. Luise, C. Ciliberto, A. Maurer, and M. Pontil. Exploiting mmd and sinkhom divergences for fair and transferable representation learning. Advances in Neural Information Processing Systems, 33: 33:15360-1533 15370, 2020.
[54] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, nigh-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[55] N. Quadrianto, V. Sharmanska, and O. Thomas. Discovering fair representations in the data domain. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages $227-8236, 2019.
[56] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pages 8748-8763. PMLR, 2021.
[57] S. Romiti, C. Inskip, V. Sharmanska, and N. Quadrianto. Realpatch: A statistical matching framework for model patching with real samples. CoRR, abs/2208.02192, 2022.
[58] P. R. Rosenbaum and D. B. Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41-55 1983.
[59] P. R. Rosenbaum and D. B. Rubin. Constructing a control group using multivariate matched sampling methods that incorporate the propensity score. The American Statistician, 39(1): 33-38, 1985.
[60] D. B. Rubin. Matching to remove bias observational studies. Biometrics, pages 159-183, 1973.
[61] D. B. Rubin. Using propensity scores to help design observational studies: application to the tobacco litigation. Health Services and Outcomes Research Methodology, 2(3):169-188, 2001.
[62] S. Sagawa, W. Koh, T. B. Hashimoto, and P. Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: 191 1. 1.08711, 2019.
[63] S. Sagawa, P. W. Koh, T. Lee, I. Gao, M. Xie, K. Shen, A. Kumar, W. Hu, M. Yasunaga, H. Marklund, S. Beery, E. David, I. Stavness, W. Guo, J. Leskovec, K. Saenko. T. Hashimoto, Levine, C. Finn, and P. Liang. Extending the WILDS benchmark for unsupervised adaptation. In International Conference on Learning Representations, 2022. URL http:t////peeereeieeee net//oorumidd=77p2VG6H3
[64] V. Sanh, L. Debut, Chaumond, and Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint 1rXiv:1910.01108,, 2019,
[65] H. Scudder. Probability of error of some adaptive pattern-recognition machines. IEEE Transac- tions on Information Theory, 11()):363-37,, 1965.
[66] H. Shimodaira. Improving predictive inference under covariate shift by weighting the log- likelihood function. Journal of statistical planning and inference, 90(2):227-244, 2000.
[67] K. Sohn, D. Berthelot, N. Carlini, Z. Zhang, H. Zhang, C. A. Raffel, E. D. Cubuk, A. Kurakin, and C.-L. Li. Fixmatch: Simplifying sumi-super learning with consistency and confidence. Advances in Neural Information Processing Systems, 33:596-608, 2020.
[68] E. A. Stuart. Matching methods for causal inference: A review and a look forward. Stat Sci, 25 (1):1-21, 2010. ISSN 0883-4237.
[69] R. Taori, A. Dave, V. Shankar, N. Carlini, B. Recht, and L. Schmidt. Measuring robustness to natural distribution shifts in image classification. Advances in Neural Information Processing Systems, 33:18583-18599 18599, 2020.
[70] A. Tarvainen and H. Valpola. Mean teachers are better role models: Weight-averaged consis- tency targets improve eupisuupee deep learning results. Advances in neural information processing systems, 30, 2017.
14

---

[71] W. Van Gansbeke, S. Vandenhende, S. Georgoulis, and L. V. Gool. Revisiting contrastive methods for unsupervised learning of visual representations. Advances in Neural Information Processing Systems, 34, 2021.
[72] V. Verma, T. Luong, K. Kawaguchi, H. Pham, and Q. Le. Towards domain-agnostic contrastive learning. In International Conference on Machine Learning, pages 10530-10541. PMLR, 2021.
[73] D. S. Watson, J. Krutzinna, I. N. Bruce, C. E. Griffiths, I. B. McInnes, M. R. Barnes, and L. Floridi. Clinical applications of machine learning algorithms: beyond the black box. Bmj, 364, 2019.
[74] O. Wiles, S. Gowal, F. Stimberg, S.-A. Rebuffi, I. Ktena, K. D. Dvijotham, and A. T. Cemgil. A fine-grained analysis on distribution shift. International Conference on Learning Representa- tions, 2022. URL https://openreevior htt.://oooooomommdddoaootttuut
[75] Q. Xie, M.-T. Luong, E. Hovy, and Q. V. Self-training with noisy student improves imagenet classification. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10687-10698, 2020.
[76] C. Ych, A. Perez, A. Driscoll, G. Azzari, Z. Tang, D. Lobell, S. Ermon, and M. Burke. Using publicly available satellite imagery and deep learning to understand economic well-being in africa. Nature communications, 11(1):1-1,, 2020.
[77] Yu, H. Chen, X. Wang. W. Xian, Chen, F. Liu, V. Madhavan, and T. Darrell. Bdd100k: A diverse driving dataset for heterogeneous multitask learning. Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 2636-2643 2020.
[78] Yu, Z. Wang, Vasudevan, L. Yeung, M. Seyedhosseini, and Y. Wu. Coca: Contrastive captioners are image-text foundation models. arXiv preprint wwxiv:2205.01917 2022.
15

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Yes; we show through our experiments detailed in Sec. 5 that our proposed method is able to outperform our baseline methods, and the original implementations of them given in [63], according three recent OOD benchmark datasets. (b) Did you describe the limitations of your work?
Yes; we describe the foremost limitation of our work in Sec. 6 Namely, we point out that using fixed thresholds for the caliper the entire duration of training is likely suboptimal and these hyperparameters should instead be set adaptively.
(c) Did you discuss any potential negative societal impacts of your work?
No, we do not anticipate there being any negative societal impacts specific to our work.
On the contrary, we propose a method for training more distributionalll--roustt models with the view to develop safer, less biased, machine-learning systems.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Yes, have read the review guidelines and have made sure that paper complies with them.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
N/A
(b) Did you include complete proofs of all theoretical results?
N/A
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)?
Yes; the code containing with the requisite instructions and configuration files needed to run all of the experiments detailed in Sec is publicly available and linked to in the Abstract.
(b) Did you specify all the training details (e.g.. data splits, hyperparameters, how they were chosen)?
Yes; all implementation details, including those related to optimisation and hyperparameterr-sllccion are given in Appendix [
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)?
Yes, we ran all methods with multiple random seeds/folds; we report the standard deviation over the runs for each method-dataset combination in parentheses in Table and Table (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)?
No, however do provide estimates of the carbon footprint for a single run of our method
and of the ERM and FixMatch baselines for the iWildCam dataset.
4. If you using existing assets (e.g., code, data, models) or curating/releusing
assets...
(a) If your work uses existing assets, did you cite the creators? Yes, we use datasets from the WILDS 2.0 benchmark developed by (63 for evaluating models and cite the associated paper in both Sec. and Sec.
(b) Did you mention the license of the assets?
No.
(c) Did you include any new assets either in the supplemental material or as a URL?
Yes, we include a URL link to the code for the paper at the end of the Abstract.
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating?
N/A
16

---

(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
N/A
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable?
N/A
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
N/A
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
N/A
17


Reviews:
[{"review_id": "xKkbOOFmhUj", "reviewer": "Reviewer_mbZt", "summary": "The paper proposes a general and robust semi-supervised learning algorithm called Okapi that works for data from various domains. It is based on the simple idea of supplementing the supervised loss with a cross-domain consistency loss that encourages the outputs of an encoder network to be similar for k-nearest neighbor samples belonging to different domains. To make the online matching in a runtime- and memory-efficient way, they incorporates a statistical matching framework and fixed-size memory bank. The results show that Okapi outperforms the baseline methods in terms of out-of-distribution (OOD) generalization on iWildCam and PovertyMap datasets. The work suggests that it is possible to effectively incorporate unlabelled data and domain information into a training algorithm to improve OOD generalization. ", "questions": "1. What is the intuition behind the CaliperNN algorithm to mimic the standard nearest neighbor algorithm? \n\nThe first stage is to remove the samples with a propensity score exceeding a predefined threshold. If I understand correctly, that means we want to remove those samples with high predictive confidence in predicting domains. In other words, we want to keep those samples that lie on the domain decision boundary. How could those samples be more likely the nearest neighbors?\n\nThe second stage of filtering is to keep paired samples from different domains only if their Euclidean distance between their respective propensity scores is below a threshold. Do you assume that the small distance in propensity score will reflect small distance in their representations?\n\n2. There are many moving parts and hyperparameters in the algorithm. Appendix C studies the effect of choosing the number of nearest neighbours k, the use of the fixed- and std-calipers in CaliperNN, and which encoder (online or target) is used to generate the queries for statistical matching. How did you choose the threshold t_f and t_theta, and are the results sensitive to the threshold values? Moreover, since the ablation is done in PovertyMap dataset, is the conclusion transfer to other datasets like iWildCam?\n\n", "limitations": "yes", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality: Making use of the domain information and using propensity score to predict domain for efficient matching is novel. \n\nQuality: The experiments are done for comparison with the existing methods, and for ablation study are solidly done. \n\nClarity: The paper is well written.\n\nSignificance: Most of the results in Table 1 are statistically significant.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "rd8Mbqhv1NR", "reviewer": "Reviewer_SyKB", "summary": "The authors studied semi-supervised learning across multiple domains. For each input example, they search for its k-nearest neighbors from other domains and enforce their feature representations to be similar. To filter outliers, the authors develop cross-domain matching based on the propensity score.\n\nThe authors applied the proposed algorithm to two datasets: iWildCam and PovertyMap. Empirical results show that their proposed method outperform previous baselines.", "questions": "In the algorithm, you tried to push the representations of the query example to be close to its unalbeled neighbors. Why would it help generalization? Aren't their representations already very close (since they are selected by k-NN)?", "limitations": "No. The authors claimed that there is no negative societal impacts.\n\nAll papers have limitations. I think it is important for researchers to discuss the weakness of their proposed method so that people can better position the paper.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The idea of enforcing feature consistency on the unlabeled data is interesting.\n\n\nWeaknesses:\n1. The empirical results (OOD 36.1 vs. 35.2) are far from satisfying. Moreover, the in-domain performance of the proposed method drops from 51.1 to 50.6. With 819K unlabeled examples (almost 8 times the size of the labeled data), I don't feel that this matching algorithm will be practically useful.\n2. The presentation quality needs to be improved.\n    + The introduction focused too much on those common sense knowledge and only talked about their proposed method in the last paragraph.\n    + Confusing notation. For example in Eq (2), what are the arguments of $NN$? In Eq (4), what is $\\text{CaliperNN}_\\xi(z, \\mathbf{z}, \\mathbf{e}, s, k))$?\n    + Given the complexity of the algorithm, I think the authors should include a pseudo-code in the main text.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "pmkEBWVg0Hg", "reviewer": "Reviewer_xMLR", "summary": "The paper uses a modified nearest-neighbors method to match unlabeled data with the labeled data in the latent feature space for semi-supervised learning. The authors conduct experiments on iWildCam and PovertyMap datasets. They compare the performance of their proposed method to two current semi-supervised learning methods (i.e., FixMatch and ERM). ", "questions": "Did you record the memory, computation costs, and carbon footprints of your experiments? If so, will you be able to add the analysis to the paper for the camera-ready version?", "limitations": "The computational and memory costs in addition to the carbon footprint of the proposed method is not discussed. As the proposed method, theoretically, seems to be computationally intensive, analysis of these costs become crucial for the assessment of the method. ", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Originality\nThe general idea of using some similarity measure to pair labeled and unlabeled images for semi-supervised learning is not new (e.g., Chen et al. (2022)). In this paper, the authors use nearest-neighbors as a measure of similarity.\n# Quality\n## Strengths\nThe paper is well written. Apart from some grammatical errors, the quality of writing is good.\n\n## Points of concern\nThe proposed nearest-neighbors matching algorithm uses the feature vectors stored in a memory bank for all the previous unlabeled input images. I have concerns regarding the memory requirement for such a memory bank. Also, given the sheer size of modern unlabeled datasets (over 100k samples in the datasets studied here), using KNN to find a match for each input image among all the previous unlabeled images during training seems like a very computationally expensive approach. To address these concerns, I suggest that the authors report the memory and computation requirements for their method. Additionally, I suggest that they compare their method to a supervised learning approach not just in terms of OOD performance, but also in terms of memory and computation costs. Then plot a trade-off curve between the added costs and gain in performance for their proposed method compared to the supervised method. Then we can see whether Okapi has practical significance.\n\nThe concern about the practical significance of the proposed method is further exacerbated by the marginal improvements in the benchmarked datasets (0.9 improvement in macro F1 score for the iWildCam and 0.05 improvement in Pearson correlation for the PovertyMap). Again, extensive analysis of the added memory and computational costs of the proposed method is required to see whether it is justified to burden such additional costs in exchange for the minor improvements in performance.\n\n\n# Clarity\nThe overall clarity of the paper is good. There are some grammatical errors that need to be fixed (e.g., line 333: “each which…”)\n# Significance\nProposing new algorithms without proper cost analysis proves little practical significance.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "bueIMREcVS9", "reviewer": "Reviewer_88f2", "summary": "The paper proposes Okapi, a method for domain adaptation / domain generalization based on the idea of matching unlabeled examples to labeled in-distribution examples and enforcing consistency in predictions. The authors propose an elaborate strategy for matching examples, and show promising performance on two datasets from the Wilds benchmark.", "questions": "Question 1. Please see my questions in the previous section. \n\nQuestion 2. Is there a reason to train the propensity classifier online, along with the encoder / main classifier? It appears that you can train it once before training the main classifier, and compute all matchings, avoiding the complications with the memory bank and online training? Or is the fact that the propensity classifier using the same features as the main classifier important to the method?", "limitations": "No issues.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength 1. The paper tackles a challenging problem of using unlabeled data for improving out-of-distribution performance. Prior work suggested that it is not trivial to achieve performance improvements over standard ERM by using unlabeled data. The authors show promising results.\n\nStrength 2. On a high level, the method makes intuitive sense. The authors also perform an ablation evaluating the importance of some parts of the method.\n\nWeakness 1. The method is quite complex, with a lot of moving parts. It involves two versions of an encoder (regular and with smoothed weights), propensity score classifier based on the embeddings, regular classifier, memory banks, two different criteria for discarding datapoints from matching (referred to as calipers in the paper). The method is not easy to understand, and it is not presented very clearly. The two calipers (methods for discarding outlier datapoints), std and fixed, are not described in sufficient detail in the main text. The matching strategy is also only described somewhat vaguely. In particular, what is $\\xi$ in CaliperNN$_{\\xi}$? How exactly are propensity scores used? \n\nGiven the complexity of the approach, it would be useful to have an algorithm box explaining precisely what each stage of the procedure is doing.\n\nWeakness 2. The empirical evaluation is somewhat limited. The authors focus on only two (challenging) datasets from the Wilds benchmark, comparing to two baselines: ERM and FixMatch. FixMatch was shown by prior work to perform poorly on these datasets, as mentioned by the authors. It is hard to know how significant the improvement is relative to the baselines: on poverty map ERM performs comparably to Okapi in terms of MSE. On iWildCam FixMatch (which was reported to work quite poorly on this dataset according to the authors) performs comparably to Okapi.\n\nOn both datasets, Okapi performs slightly better, but it is hard to tell if the improvement is significant, especially given that the method is more or less designed to work well on these two datasets, while FixMatch was shown to perform particularly poorly on iWildsCam. \n\nAnother complication with the evaluation is that the authors chose to use a different backbone compared to prior work, ConvNext. Consequently, we cannot directly compare the numbers reported by the authors to the numbers reported in the prior work. Why was this choice made?\n\nThe Calipers in table 1 appear to provide a fairly minor improvement compared to Okapi no caliper.\n\nWeakness 3. It seems like the method may not work well in the cases when the domains are sufficiently different, such as in the common domain generalization benchmarks. The idea of the method is to reject datapoints which can be differentiated from the in-distribution datapoints (based on propensity scores). On datasets such as PACS the domains can be clearly differentiated, so my understanding is that the method would reject all the images from other domains. Is this interpretation correct?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept