SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Nora Schneider<sup>1</sup> Shirin Goshtaspbour<sup>1,2</sup> Fernando Perez-Cruz<sup>1,2</sup>
<sup>1</sup>Computer Science Department, ETH Zurich, Zurich, Switzerland  
<sup>2</sup>Swiss Data Science Center, Zurich, Switzerland
nschneide@student.ethz.ch shirin.goshtasbpour@inf.ethz.ch fernando.perezcruz@sdsc.ethz.ch
We propose a novel algorithm for data augmentation in nonlinear over-parametrized regression. Our data augmentation algorithm borrows from the literature on causality and extends the recently proposed Anchor regression (AR) method for data augmentation, which is in contrast to the current state-of-the-art domain-agnostic solutions that rely on the Mixup literature. Our Anchor Data Augmentation (ADA) uses several replicas of the modified samples in AR to provide more training examples, leading to more robust regression predictions. We apply ADA to linear and nonlinear regression problems using neural networks. ADA is competitive with state-of-the-art C-Mixup solutions.
Data augmentation is one of the key ingredients of any successful application of a machine learning classifier. The first example that typically comes to mind is the in-depth description of the data augmentation in the now-famous Alexnet paper [26]. Data augmentation algorithms come in different flavors, and they mostly rely on the expectation that small perturbations, invariances, or symmetries applied to the input will not change the class label. That way, we can present 'fresh new' samples as alterations of the available examples for training. These transformations modify the input distribution to make the algorithm more robust for cases where the distribution of the test set may differ from that of the training set. We refer the reader to the related work section (Section 2.1) for an overview and description of different data augmentation strategies.
The literature for data augmentation in regression is slim. The paper on Mixup augmentation [51] proposes a simple and general scheme for data augmentation using convex combinations of samples. The authors only apply their data augmentation proposal to classification problems. They conjecture in the discussion that the application to regression is straightforward, however, this is not the case in practice. Mixup is theoretically analyzed in [5, 52] as a regularization technique for classification and regression problems. However, it is only illustrated in classification problems.
The Mixup algorithm has been extended to regression problems in [18, 49], in which the authors explain that Mixup cannot be blindly applied to regression problems. To our knowledge, these are the only two papers in which data augmentation for regression is proposed. RegMix [18] relies on a hard-to-train prior neural network controller before augmenting the data using a Mixup strategy. C-Mixup [49], a method proposed more recently, solves some of the issues limiting the standard Mixup algorithm for regression problems. The authors propose to mix only closeby samples in the
$^{1}$ Our Python implementation of ADA is available at: https://github.com/noraschneider/anchordataaugmentation/
output space (i.e., samples which have close enough labels). This strategy is only valid when the target variables are monotonic with the input and is applied in a transformed space. The authors present comprehensive results in data augmentation for in-distribution generalization, task generalization and out-of-distribution robustness.
In this paper, we rely on the causality literature to provide a different avenue for augmenting data in regression problems. Causal discovery finds the causes of a response variable among a given set of observations or helps to recognize the causal relations between a set of variables [39]. These causes allow us to understand how these relations will change if we were to intervene in a subset of the (input) variables or what would be the effect on the output. So, in general, the regression model will be robust to perturbations in the input variables making the prediction less sensitive to changes in the distribution of the test set. For example, the authors in [40] use the invariance property for prediction to perform causal inference. In turn, Anchor Regression (AR) builds upon the causality literature to obtain robust regression solutions when the input variables have been perturbed [42]. The procedure relies on anchor variables capturing the heterogeneity within a dataset and a parameter  $\gamma$  that measures the deviation with respect to the least square solution. Once the values of the anchors are known, AR modifies the data and obtains the least square solution, as detailed in Section 2.2
In this paper, we propose **Anchor Data Augmentation (ADA)** to augment the training dataset with several replicas of the available data. We use a simple clustering of the data to encode a homogeneous group of observations and use different values of  $\gamma$  to robustify the solution to different strengths of potential distribution shifts. In every minibatch, we sample  $\gamma$  from a predetermined range around  $\gamma = 1$ . As AR was developed for linear regression, the data augmentation strategy needs to be modified for nonlinear regression accordingly. We validate ADA for in-distribution generalization and out-of-distribution robustness under the same conditions proposed in C-Mixup [49], as well as some illustrative linear and nonlinear regression examples. In the replicated experiments, ADA is competitive or superior to other augmentation strategies such as C-Mixup, although on some datasets the performance gain is marginal.
The rest of the paper is organized as follows: First, we provide background information in Section 2. We give a brief overview of related work on data augmentation in Section 2.1 and summarize the key concepts on Anchor Regression in Section 2.2. Second, Section 3 shows how we extend Anchor Regression and introduces ADA. Section 4 reports empirical evidence that our approach can improve predictions, especially in over-parameterized settings. We conclude the paper in Section 5.
2.1 Data Augmentation
Many different data augmentation methods have been proposed in recent years with several applications in mind. Still most augmentations we mention here use human-designed transformations based on domain knowledge which leave the target variable invariant. For instance, Cutout [10] is an image-specific augmentation technique that is successfully used to train models on CIFAR10 and CIFAR100 [25], but was determined to be unsuitable for larger image datasets like ImageNet with higher resolution [9]. Other augmentation methods for images such as random crop, horizontal or vertical mirroring, random rotation, or translation [29, 43] may similarly apply to a certain group of image datasets while being inapplicable to others, e.g. datasets of digits and letters.
In an attempt to automate the augmentation process and reduce human involvement, policy or search-based automated augmentation methods were developed. In AutoAugment [7] a neural network is trained with Reinforcement Learning (RL) to combine an assortment of transformations in varying strengths to apply on samples of a given dataset and improve the model accuracy. Methods such as RandAugment [8], Fast AutoAugment [30], UniformAugment [32] and TrivialAugment [36] aim at reducing the cost of the pretraining search phase in automated augmentation with randomized transformations and reduced search space.
Alternatively, in order to adapt the augmentation policy to the model during training, Population-Based Augmentation [16] and Online Hyperparameter Learning [31] use multiple data augmentation workers that are updated using evolutionary strategies and RL, respectively. Adversarial AutoAugment [53] and AugMax [47] optimize for the augmentation policy that deteriorates the training accuracy
and improves its robustness. DivAug [34] finds the policy which maximizes the diversity of the augmented data.
Having a separate search phase for optimal augmentation policy is computationally expensive and may exceed the required computation to train the downstream model [8, 48]. In addition, these methods and their online counterparts need to be trained separately on every single dataset. While OnlineAugment [44] and DDAS exploit meta-learning to avoid this problem, they still rely on a set of predefined class invariant transformations that require domain-specific information.
Generic transformations such as Gaussian or adversarial noise [10, 28, 45] and dropout [3] are also effective in expanding the training dataset. Generative models such as Generative Adversarial Networks (GAN) [13] and Variational Auto-Encoders (VAE) [22] are trained in [16, 44] to synthesize samples close to the low dimensional manifold of the data for classification.
Mixup [51] is a popular data augmentation using a convex combination of pairs of samples from different classes and their softened labels for augmentation. Mixup is only evaluated on classification problems, even though it is claimed that the application to regression is straightforward. Various extensions of Mixup have been proposed to prevent data manifold intrusion [46], use more complex mixing strategies [33, 50] or account for saliency in augmented samples [20, 21]. These methods were predominantly designed to excel in classification tasks. In particular, Mixup for regression was studied in [5, 18, 49, 52] but it was reported to adversely impact the predictions in regression problems when misleading augmented samples are generated from a pair of faraway samples.
2.2 Anchor Regression
We summarize the key concepts of Anchor Regression (AR) as presented in [42]. Let  $X \in \mathcal{X}$  and  $y \in \mathcal{Y}$  be the predictors and target variables sampled from distribution  $(X, y) \sim P_{\mathrm{train}}$ ,  $\mathcal{X} \subseteq \mathbb{R}^d$  and  $\mathcal{Y} \subseteq \mathbb{R}$ . Traditionally, a causal framework models the relation of  $y$  and  $X$  to accurately predict the value of  $y$  under given interventions or arbitrary perturbations on  $X$ . A commonly held assumption is that the underlying causal relation among variables remains the same while the sampling distribution  $P_{\mathrm{train}}$  is altered by the intervention shift or the applied perturbation. For instance, if the distribution  $P_{\mathrm{train}}$  is induced by an unknown linear causal model, then the causally optimal parameters can be expressed as the solution to the optimization problem:
$$
b _ {\text {c a u s a l}} = \underset {b} {\arg \min } \underset {P \in \mathcal {P}} {\max } \mathbb {E} _ {P} \left[ (y - X ^ {T} b) ^ {2} \right], \tag {1}
$$
where  $\mathcal{P}$  is the class of distributions containing all interventions on components of  $X$  [41]. Therefore, causal parameters provide distributionally robust predictions that are optimal under the intervention in  $\mathcal{P}$ . In comparison, Ordinary Least Squares (OLS):
$$
b _ {\mathrm {O L S}} = \underset {b} {\arg \min } \mathbb {E} _ {P _ {\text {t r a i n}}} [ (y - X ^ {T} b) ^ {2} ], \tag {2}
$$
may lead to arbitrarily large predictive errors on distributions in  $\mathcal{P}$ . On the other hand, on  $P_{\mathrm{train}}$ , causal parameters  $b_{\mathrm{causal}}$  lead to conservative predictions, while  $b_{\mathrm{OLS}}$  presents optimal least squared performance.
To trade-off predictive accuracy on the training distribution with distribution robustness and to enforce stability over statistical parameters, AR [4, 42] proposes to relax the regularization in the optimization problem in (1) to a smaller class of distributions  $\mathcal{P}$ .
Assume that  $X$  and  $y$  are centered and have finite variance. We use  $A \in \mathbb{R}^q$  (called anchors) to denote the exogenous (random) variables in  $X$  which generate heterogeneity in  $y$ . We further denote the  $L_{2}$ -projection on the linear span of the components of  $A$  with  $\mathrm{P}_A$  and  $\operatorname{Id}(y) = y$ . Under linear assumption between  $A$  and  $(X, y)$ , we can write the relaxed optimization problem as:
$$
b _ {\gamma , A} = \underset {b} {\arg \min } \mathbb {E} _ {P _ {\text {t r a i n}}} \left[ \left(\left(\operatorname {I d} - \mathrm {P} _ {A}\right) \left(y - X ^ {T} b\right)\right) ^ {2} \right] + \gamma \mathbb {E} _ {P _ {\text {t r a i n}}} \left[ \left(\mathrm {P} _ {A} \left(y - X ^ {T} b\right)\right) ^ {2} \right], \tag {3}
$$
where  $\gamma > 0$  is a hyperparameter. The first term of the AR objective in Equation 3 is the loss after "partially out" the anchor variable, which refers to first linearly regressing out  $A$  from  $X$  and  $y$  and subsequently using OLS on the residuals. The second term is the well-known estimation objective used in the Instrumental Variable setting [11]. Therefore, for different values of  $\gamma$  AR interpolates between the partially out objective ( $\gamma = 0$ ) and the IV estimator ( $\gamma \to \infty$ ) and coincides with
OLS for  $\gamma = 1$ . The authors show that the solution of AR optimizes a worst-case risk under shift-interventions on anchors up to a given strength. This in turn increases the robustness of the predictions to distribution shifts at the cost of reducing the in-distribution generalization.
In the finite-sample case with  $n$  observations from  $P_{\mathrm{train}}$ , let matrix  $\mathbf{X} \in \mathbb{R}^{n \times d}$  contain the observations of  $X$  and let  $\mathbf{Y} \in \mathbb{R}^n$  be the vector of corresponding targets. Similarly, we denote the matrix containing the observations of  $A$  with  $\mathbf{A} \in \mathbb{R}^{n \times q}$  and we use  $\Pi_{\mathbf{A}} = \mathbf{A}\left(\mathbf{A}^T\mathbf{A}\right)^{\dagger}\mathbf{A}^T$  as the projection operator on the column space of the anchor matrix  $\mathbf{A}$  where  $\mathbf{A}^{\dagger}$  denotes the pseudo-inverse of matrix  $\mathbf{A}$ . Further,  $\mathbf{I}$  denotes the identity matrix. Then, the finite-sample optimization regression problem can be written as
$$
\hat {b} _ {\gamma , \mathbf {A}} = \underset {b} {\arg \min } \| (\mathbf {I} - \boldsymbol {\Pi} _ {\mathbf {A}}) (\mathbf {Y} - \mathbf {X} b) \| _ {2} ^ {2} + \gamma \| \boldsymbol {\Pi} _ {\mathbf {A}} (\mathbf {Y} - \mathbf {X} b) \| _ {2} ^ {2}. \tag {4}
$$
The AR regression estimate  $\hat{b}_{\gamma, \mathbf{A}}$  can be obtained by applying the OLS solution to a modified set of inputs and outputs:
$$
\tilde {\mathbf {X}} _ {\gamma , \mathbf {A}} = \mathbf {X} + (\sqrt {\gamma} - 1) \boldsymbol {\Pi} _ {\mathbf {A}} \mathbf {X} \tag {5}
$$
$$
\tilde {\mathbf {Y}} _ {\gamma , \mathbf {A}} = \mathbf {Y} + (\sqrt {\gamma} - 1) \boldsymbol {\Pi} _ {\mathbf {A}} \mathbf {Y} \tag {6}
$$
3 Anchor Data Augmentation
In this section, we introduce Anchor Data Augmentation (ADA), a domain-independent data augmentation method inspired by AR. ADA does not require previous knowledge about the data invariances nor manually engineered transformations. As opposed to existing domain-agnostic data augmentation methods [10, 45, 46], we do not require training of an expensive generative model, and the augmentation only adds marginally to the computation complexity of the training. In addition, since ADA originates from a causal regression problem, it can be readily applied to regression problems. Even when ADA does not improve performance, its effect on performance remains minimal.
Data augmentation aims to introduce informative data in addition to the original dataset during the training procedure of the model to improve its generalization. Similar to AR, ADA employs a linear projection, given by the anchor variables  $A$ , to determine the most relevant perturbation directions based on the similarity of the samples. ADA inherits the generalization properties from AR. In [42], the authors recommend that the anchor variable can be set as an indicator of the datasets, where each dataset is a homogeneous set of observations. A key insight of our work is that this can be achieved by clustering the data into  $q$  clusters. The matrix  $\mathbf{A} \in \mathbb{R}^{n \times q}$  is then constructed as an indicator matrix with a one-hot encoding of the assigned cluster index per row. For our experiments, we use  $k$ -means clustering [35] to construct  $\mathbf{A}$ . Further, in AR, only one value for  $\gamma$  is used, which should be chosen based on the desired strength of perturbations on test datasets, in comparison to the training dataset [42]. We suggest that the value of  $\gamma$  is sampled from a distribution with density  $p(\gamma)$ . In our experiments, we use a uniform distribution between  $1 / \alpha$  and  $\alpha$ , where  $\alpha > 1$  is a hyperparameter to be tuned.
ADA augments a sample  $(\mathbf{X}^{(i)},\mathbf{Y}^{(i)})$  by normalizing the original AR modifications (5 and 6) by  $1 + (\sqrt{\gamma} -1)\sum_{j}(\Pi_{\mathbf{A}})^{(ij)}$  to unify the noise level across the augmentations independent of the value of  $\gamma$ , while approximately preserving the potentially nonlinear relation between  $X$  and  $y$  (see also section 3.2):
$$
\tilde {\mathbf {X}} _ {\gamma , \mathbf {A}} ^ {(i)} = \frac {\mathbf {X} ^ {(i)} + (\sqrt {\gamma} - 1) \left(\boldsymbol {\Pi} _ {\mathbf {A}}\right) ^ {(i)} \mathbf {X}}{1 + (\sqrt {\gamma} - 1) \sum_ {j} \left(\boldsymbol {\Pi} _ {\mathbf {A}}\right) ^ {(i j)}}, \tag {7}
$$
$$
\tilde {\mathbf {Y}} _ {\gamma , \mathbf {A}} ^ {(i)} = \frac {\mathbf {Y} ^ {(i)} + (\sqrt {\gamma} - 1) \left(\boldsymbol {\Pi} _ {\mathbf {A}}\right) ^ {(i)} \mathbf {Y}}{1 + (\sqrt {\gamma} - 1) \sum_ {j} \left(\boldsymbol {\Pi} _ {\mathbf {A}}\right) ^ {(i j)}}, \tag {8}
$$
where we denote  $(\mathbf{M})^{(i)}$ ,  $(\mathbf{M})^{(j)}$ , and  $(\mathbf{M})^{(ij)}$  denote respectively the  $i$ -th row, the  $j$ -th column and  $(i,j)$  component of some matrix  $\mathbf{M}$ . As is standard practice, we rely on stochastic gradient descent to optimize our (nonlinear) regressors and apply ADA on each minibatch rather than the entire dataset.
ADA combines samples from the same cluster and generates augmented samples along the data manifold. For a general  $\mathbf{A}$ ,  $\Pi_{\mathbf{A}}$  provides a "collective" mixing approach for the samples in a batch
Figure 1: Comparison of ADA augmentations on a nonlinear Cosine data model. For a larger partition size, ADA augmentations are more accurate due to the high local variability of the Cosine function. We used  $k$ -means clustering to construct A and  $\gamma \in \{1/2, 2/3, 1.03/2, 2.0\}$ .
by determining a center, while  $\gamma$  controls the extend of contraction or expansion of the augmented sample around this center. In particular, for a one-hot encoding matrix  $\mathbf{A}$ ,  $\Pi_{\mathbf{A}}^{(i)}\mathbf{X}$  defines the centroid of the cluster to which sample  $i$  belongs. Then, the modified samples are located on the ray that originates from the centroid and goes through the original data point  $(\mathbf{X}^{(i)},\mathbf{Y}^{(i)})$ . As  $\gamma$  increases, the augmented samples move towards their corresponding centroid and specifically, for  $\gamma = 1$  they coincide with the original samples. Furthermore, the cluster size, regulated by the number of clusters  $q$ , directly impacts the number of samples mixed together; with smaller clusters, fewer samples are combined. Applying ADA on each minibatch introduces further diversity and enhances robustness, because the composition of samples being mixed together and the value of  $\gamma$  changes in each minibatch. In Appendix A.2 we provide a detailed explanation and analysis of the impact of ADA hyperparameters,  $q$  controlling the number of clusters and  $\alpha$  controlling the range of values for  $\gamma$ . In Appendix B.4 we empirically show how regression performance varies with respect to these hyperparameters.
In Figure 1 we visually illustrate the augmentation effects of ADA. We uniformly sampled 30 data points between  $\pm 3$  (i.e.  $x_{i} \sim \mathcal{U}[-3,3]$ ) and set the corresponding target variable as  $y_{i} = \cos(\pi x_{i})$  without added noise. We then clustered this data in  $q = 5$  and  $q = 12$  groups using k-means and applied eq. (7) and eq. (8) to the 30 samples with  $\gamma \in \{1/2,2/3,1,3/2,2\}$  resulting in 150 augmented data points.
3.1 Comparison to C-Mixup
ADA can be interpreted as a generalized variant of C-Mixup [49]. In C-Mixup samples are mixed in pairs, and the combination probability of each sample pair is given by the similarity of their labels, measured by a Gaussian kernel. Augmented samples are then obtained as the convex combination between the pair. In contrast, ADA allows mixing multiple samples based on their cluster membership and the resulting augmentations that may reside in the convex hull of the original samples of a cluster if  $\gamma \geq 1$  or beyond it when  $\gamma < 1$ . In particular, ADA and C-Mixup augmentations would be similar if the anchor matrix  $\mathbf{A}$  indicates pairs of samples weighted by the similarity of their labels and  $\gamma > 1$ .
3.2 Preserving nonlinear data structure
In the following we show that the scaled transformations in eq. (5) and eq. (6) preserve the nonlinear relationship, so that we can use the modified pair  $(\tilde{\mathbf{X}}_{\gamma,\mathbf{A}}, \tilde{\mathbf{Y}}_{\gamma,\mathbf{A}})$  to augment the dataset  $(\mathbf{X}, \mathbf{Y})$ . Let  $(\mathbf{X}^{(\mathrm{i})}, \mathbf{Y}^{(i)})$  be the  $i$ th sample from  $P_{\mathrm{train}}$  corresponding to the  $i$ th row of  $\mathbf{X}$  and  $i$ th component of  $\mathbf{Y}$ . When the data has a nonlinear relation,
$$
\mathbf {Y} ^ {(i)} = f _ {b} \left(\mathbf {X} ^ {(\mathrm {i})}\right) + \epsilon^ {(i)} \tag {9}
$$
given the zero mean noise variable  $\epsilon^{(i)}$ , we can alter the anchor loss accordingly [4],
$$
b _ {\text {N O N L I N}, \gamma , A}, f _ {\gamma , A} = \underset {b, f} {\arg \min } \mathbb {E} _ {P _ {\text {t r a i n}}} [ ((\operatorname {I d} - \operatorname {P} _ {A}) (y - f _ {b} (X))) ^ {2} ] + \gamma \mathbb {E} _ {P _ {\text {t r a i n}}} [ (\operatorname {P} _ {A} (y - f _ {b} (X))) ^ {2} ],
$$
The AR modification Equations 5 and 6 do not preserve the nonlinear relation between the target and predictors,
$$
\tilde {\mathbf {Y}} ^ {(i)} \neq f _ {b} (\tilde {\mathbf {X}} ^ {(i)}) + \tilde {\epsilon} ^ {(i)}
$$
with another zero mean variable  $\tilde{\epsilon}^{(i)}$  operating as the observation noise in the augmented data. Therefore, we propose to further extend the original AR and perform the data augmentation with
scaled transformations to get the modified sample  $(\tilde{\mathbf{X}}_{\gamma, \mathbf{A}}^{(i)}, \tilde{\mathbf{Y}}_{\gamma, \mathbf{A}}^{(i)})$  which approximately preserves the nonlinear relationship of sample  $(\mathbf{X}^{(\mathbf{i})}, \mathbf{Y}^{(i)})$  as shown below.
We can rewrite  $\tilde{\mathbf{Y}}_{\gamma ,\mathbf{A}}^{(i)}$  in Equation (8) as
$$
\tilde {\mathbf {Y}} _ {\gamma , \mathbf {A}} ^ {(i)} = \frac {f _ {b} (\mathbf {X} ^ {(\mathbf {i})}) + (\sqrt {\gamma} - 1) (\boldsymbol {\Pi} _ {\mathbf {A}}) ^ {(i)} \mathbf {F} _ {b} (\mathbf {X})}{1 + (\sqrt {\gamma} - 1) \sum_ {j} (\boldsymbol {\Pi} _ {\mathbf {A}}) ^ {(i j)}} + \tilde {\epsilon} _ {\gamma , \mathbf {A}} ^ {(i)}
$$
where  $\tilde{\epsilon}_{\gamma,\mathbf{A}}^{(i)}$  is a zero mean noise variable and  $\mathbf{F}_b(\mathbf{X}) = [f_b(\mathbf{X}^{(1)}),\dots,f_b(\mathbf{X}^{(n)})]^T$ . In Appendix A.1, for continuously differentiable function  $f$ , we can use the first order Taylor expansion of  $\tilde{\mathbf{Y}}_{\gamma,\mathbf{A}}^{(i)}$  around  $\tilde{\mathbf{X}}_{\gamma,\mathbf{A}}^{(i)}$  to show that
$$
\tilde {\mathbf {Y}} _ {\gamma , \mathbf {A}} ^ {(i)} \approx f _ {b} \left(\tilde {\mathbf {X}} _ {\gamma , \mathbf {A}} ^ {(i)}\right) + \tilde {\epsilon} _ {\gamma , \mathbf {A}} ^ {(i)} \tag {10}
$$
which approximately has the same nonlinear relation as the original model for small  $\| \mathbf{X}^{(\mathbf{i})} - \tilde{\mathbf{X}}_{\gamma ,\mathbf{A}}^{(i)}\| _2$  or  $\| \sum_j(\Pi_{\mathbf{A}})^{(ij)}(\mathbf{X}^{(j)} - \tilde{\mathbf{X}}_{\gamma ,\mathbf{A}}^{(i)})\| _2$