summaries: This work proposes an approach to learn to share attention heads in multitask (multilingual, multi-domain) transformer models with multi-head attention layers. Given a multi-head attention layer with $H'$ heads being trained on a set of tasks, each task learns to use $H$ heads, and the choice of attention heads is dictated by learnt latent variables. 2 strategies are proposed for this attention head selection: 1. Subset selection, selects top-$H$ heads for each task. 2. Group selection, where $H'$ heads are divided into groups of $H'/H$ heads, and 1 head is selected per task from each group.  This approach is compared against baseline approaches with full parameter sharing (transformers for MT, speech-transformer for other tasks) and baselines enhanced with adapter modules [1] on multilingual Machine Translation, multilingual ASR, multi-domain ASR, multilingual speech translation and multi-domain speech translation. On all these tasks the proposed head-sharing approach shows modest improvements over the fully-shared baselines.  References: [1] Simple scalable adaptation for neural machine translation, Bapna et al.

summaries: This paper address the problem of negative interference which happens when a single transformer model is shared among the different tasks in a multi-task learning setup specifically those of multilingual translation, speech recognition, and multi-domain speech translation.  The proposed approach is to adaptively make use of specific attention heads for a particular task among a pool of attention heads. In this way, the common attention heads between tasks can reinforce the positive transfer while the attention heads unique to each task will avoid the negative transfer. Extensive experiments on multiple tasks of multilingual translation and speech translation demonstrate improvements in performance over strong baselines.

summaries: The paper proposes a transformer architecture for multi-task learning, based on sharing some attention heads in multi-head attention between different tasks. The model is supplied with more heads than necessary and a neural module using gumbel-softmax is used to select a subset of these heads for each task.  A variant of this is described, where heads are distributed between groups and exactly one head is selected from each group.

summaries: Multi-head attention is an essential component for popular Transformer models. This paper proposes to learn shared and specialized attention heads for different languages and domains. The authors formulate attention selection as latent variables and adopt Gumbel softmax to select attention heads.  Experiments on text-to-text and speech-to-text translation verify the effectiveness of the proposed approach.

summaries: This paper aims to mitigate the problem of negative interference when learning models across different languages and domains. In particular, it proposes an approach that learns to share or specialize attention heads based on two different strategies, namely subset, and group selection.  The evaluation of multi-lingual/domain settings shows that this approach outperforms standard transformers and adapters without increasing the number of parameters.

abstractText: Multi-head attention has each of the attention heads collect salient information from different parts of an input sequence, making it a powerful mechanism for sequence modeling. Multilingual and multi-domain learning are common scenarios for sequence modeling, where the key challenge is to maximize positive transfer and mitigate negative interference across languages and domains. In this paper, we find that non-selective attention sharing is sub-optimal for achieving good generalization across all languages and domains. We further propose attention sharing strategies to facilitate parameter sharing and specialization in multilingual and multi-domain sequence modeling. Our approach automatically learns shared and specialized attention heads for different languages and domains. Evaluated in various tasks including speech recognition, text-to-text and speech-to-text translation, the proposed attention sharing strategies consistently bring gains to sequence models built upon multi-head attention. For speech-to-text translation, our approach yields an average of +2.0 BLEU over 13 language directions in multilingual setting and +2.0 BLEU over 3 domains in multi-domain setting.

authors:
Ankur Bapna
Orhan Firat

title: Simple, scalable adaptation for neural machine translation

venue: In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),

authors:
Graeme W. Blackwood
Miguel Ballesteros
Todd Ward

title: Multilingual neural machine translation with task-specific attention

venue: In Proceedings of the 27th International Conference on Computational Linguistics,

authors:
Zhao Chen
Vijay Badrinarayanan
Chen-Yu Lee
Andrew Rabinovich

title: Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks

venue: Proceedings of the 35th International Conference on Machine Learning,

authors:
Raj Dabre
Chenhui Chu
Anoop Kunchukuttan

title: A survey of multilingual neural machine translation

venue: ACM Computing Surveys (CSUR),

authors:
Praveen Dakwale
Christof Monz

title: Finetuning for neural machine translation with limited degradation across in-and out-of-domain data

venue: Proceedings of the XVI Machine Translation Summit,

authors:
William Fedus
Barret Zoph
Noam Shazeer

title: Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity

venue: arXiv preprint arXiv:2101.03961,

authors:
Mattia Antonino Di Gangi
Matteo Negri
Marco Turchi

title: Adapting transformer to end-to-end spoken language translation

venue: Interspeech 2019,

authors:
Xinwei Geng
Longyue Wang
Xing Wang
Bing Qin
Ting Liu
Zhaopeng Tu

title: How does selective mechanism improve self-attention networks

venue: Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics,

authors:
Jiatao Gu
Hany Hassan
Jacob Devlin
Victor OK Li

title: Universal neural machine translation for extremely low resource languages

venue: In Proceedings of the 2018 Conference of the North American Chapter of the Association

authors:
Pengsheng Guo
Chen-Yu Lee
Daniel Ulbricht

title: Learning to branch for multi-task learning

venue: In International Conference on Machine Learning,

authors:
Georg Heigold
Vincent Vanhoucke
Andrew W. Senior
Patrick Nguyen
Marc’Aurelio Ranzato
Matthieu Devin
Jeffrey Dean

title: Multilingual acoustic models using distributed deep neural networks

venue: IEEE International Conference on Acoustics, Speech and Signal Processing,

authors:
H. Inaguma
Kevin Duh
Tatsuya Kawahara
Shinji Watanabe

title: Multilingual end-to-end speech translation

venue: IEEE Automatic Speech Recognition and Understanding Workshop (ASRU),

authors:
Hirofumi Inaguma
Kevin Duh
Tatsuya Kawahara
Shinji Watanabe

venue: In IEEE Automatic Speech Recognition and Understanding Workshop,

authors:
Javier Iranzo-Sánchez
Joan Albert Silvestre-Cerdà
Javier Jorge
Nahuel Roselló
Adrià Giménez
Albert Sanchis
Jorge Civera
Alfons Juan

title: Europarl-st: A multilingual corpus for speech translation of parliamentary debates

venue: IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP),

authors:
Eric Jang
Shixiang Gu
Ben Poole

title: Categorical reparameterization with gumbel-softmax

venue: In 5th International Conference on Learning Representations,

authors:
M. Johnson
M. Schuster
Quoc V. Le
M. Krikun
Y. Wu
Z. Chen
Nikhil Thorat
Fernanda B. Viégas
M. Wattenberg
G. Corrado
Macduff Hughes
J. Dean

title: Google’s multilingual neural machine translation system: Enabling zero-shot translation

venue: Transactions of the Association for Computational Linguistics,

authors:
Huda Khayrallah
Brian Thompson
Kevin Duh
Philipp Koehn

title: Regularized training objective for continued training for domain adaptation in neural machine translation

venue: Proceedings of the 2nd Workshop on Neural Machine Translation and Generation,

authors:
Dietrich Klakow
Jochen Peters

title: Testing the correlation of word error rate and perplexity

venue: Speech Communication,

authors:
Catherine Kobus
Josep Maria Crego
Jean Senellart

title: Domain control for neural machine translation

venue: In Proceedings of the International Conference Recent Advances in Natural Language Processing, RANLP 2017,

authors:
Dmitry Lepikhin
HyoukJoong Lee
Yuanzhong Xu
Dehao Chen
Orhan Firat
Yanping Huang
Maxim Krikun
Noam Shazeer
Zhifeng Chen

title: Gshard: Scaling giant models with conditional computation and automatic sharding

authors:
X. Li
Changhan Wang
Y. Tang
C. Tran
Yuqing Tang
J. Pino
Alexei Baevski
Alexis Conneau
Michael Auli

title: Multilingual speech translation with efficient finetuning of pretrained models

venue: arXiv: Computation and Language,

authors:
Pengbo Liu
Hailong Cao
Tiejun Zhao

title: Gumbel-attention for multi-modal machine translation

authors:
Yinhan Liu
Jiatao Gu
Naman Goyal
Xian Li
Sergey Edunov
Marjan Ghazvininejad
Mike Lewis
Luke Zettlemoyer

title: Multilingual denoising pre-training for neural machine translation

authors:
Niko Moritz
Takaaki Hori
Jonathan Le

title: Streaming automatic speech recognition with the transformer model

authors:
Myle Ott
Sergey Edunov
Alexei Baevski
Angela Fan
Sam Gross
Nathan Ng
David Grangier
Michael Auli

title: fairseq: A fast, extensible toolkit for sequence modeling

venue: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,

authors:
Kishore Papineni
Salim Roukos
Todd Ward
Wei-Jing Zhu

title: Bleu: a method for automatic evaluation of machine translation

venue: In Proceedings of the 40th annual meeting of the Association for Computational Linguistics,

authors:
MinhQuang Pham
Josep Maria Crego
François Yvon

title: Revisiting multi-domain machine translation

authors:
Vineel Pratap
Anuroop Sriram
Paden Tomasello
Awni Y. Hannun
Vitaliy Liptchinsky
Gabriel Synnaeve
Ronan Collobert

title: Massively multilingual asr: 50 languages, 1 model, 1 billion parameters

venue: In INTERSPEECH,

authors:
Devendra Singh Sachan
Graham Neubig

title: Parameter sharing methods for multilingual selfattentional translation models

venue: In Proceedings of the Third Conference on Machine Translation: Research Papers, WMT 2018,

authors:
Elizabeth Salesky
Matthew Wiesner
Jacob Bremerman
Roldano Cattoni
Matteo Negri
Marco Turchi
Douglas W. Oard
Matt Post

title: The multilingual tedx corpus for speech recognition and translation

title: Domain adaptation and multi-domain adaptation for neural machine translation: A survey

authors:
Sukanta Sen
Kamal Kumar Gupta
Asif Ekbal
Pushpak Bhattacharyya

title: Multilingual unsupervised nmt using shared encoder and language-specific decoders

venue: In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics,

authors:
Noam Shazeer
Azalia Mirhoseini
Krzysztof Maziarz
Andy Davis
Quoc V. Le
Geoffrey E. Hinton
Jeff Dean

title: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer

authors:
T. Standley
A. Zamir
Dawn Chen
L. Guibas
Jitendra Malik
S. Savarese

title: Which tasks should be learned together in multi-task learning

venue: In ICML,

authors:
Brian Thompson
Jeremy Gwinnup
Huda Khayrallah
Kevin Duh
Philipp Koehn

title: Overcoming catastrophic forgetting during domain adaptation of neural machine translation

authors:
Laura Cross Vila
Carlos Escolano
José A.R. Fonollosa
Marta R

title: Costa-jussà. End-to-end speech translation with the transformer

venue: Fourth International Conference,

authors:
Changhan Wang
Yun Tang
Xutai Ma
Anne Wu
Dmytro Okhonko
Juan Pino

title: Fairseq s2t: Fast speech-to-text modeling with fairseq

venue: In Proceedings of the 1st Conference of the AsiaPacific Chapter of the Association for Computational Linguistics and the 10th International Joint Conference on Natural Language Processing: System Demonstrations,

title: fairseq s2t: Fast speech-to-text modeling with fairseq

venue: In Proceedings of the 2020 Conference of the Asian Chapter of the Association for Computational Linguistics (AACL): System Demonstrations,

authors:
Changhan Wang
Anne Wu
Juan Pino

title: Covost 2: A massively multilingual speech-to-text translation

venue: corpus. CoRR,

authors:
Yong Wang
Longyue Wang
Shuming Shi
Victor OK Li
Zhaopeng Tu

title: Go from the general to the particular: Multi-domain translation with domain transformation networks

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Zirui Wang
Zachary C Lipton
Yulia Tsvetkov

title: On negative interference in multilingual language models

venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP),

authors:
Sen Wu
Hongyang R Zhang
Christopher Ré

title: Understanding and improving information transfer in multi-task learning

venue: In International Conference on Learning Representations,

authors:
Tianhe Yu
Saurabh Kumar
Abhishek Gupta
Sergey Levine
Karol Hausman
Chelsea Finn

title: Gradient surgery for multi-task learning

authors:
Jiali Zeng
Jinsong Su
Huating Wen
Yang Liu
Jun Xie
Yongjing Yin
Jianqiang Zhao

title: Multidomain neural machine translation with word-level domain context discrimination

venue: In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing,

heading: 1 Introduction

text:
Recent progress on deep learning models, in particular multi-head attention, has brought significant gains to sequence modeling tasks including speech recognition (Moritz et al., 2020), text-to-text translation (Vaswani et al., 2017), and speech-to-text translation (Vila et al., 2018; Gangi et al., 2019). Attention mechanism allows a model to focus on informative parts of the inputs, and multi-head attention computes attention over inputs by multiple heads independently. With each head attending to different information, multi-head attention potentially captures more complicated data patterns and extracts sophisticated knowledge.
Sequence modeling has attracted a lot of research interest in multilingual and multi-domain settings, where a model is trained on data in multiple language directions and data from different domains respectively. Key advantages of these settings are better data efficiency and the support of knowledge transfer among languages or domains. This is critical for resource-limited scenarios. For example, multilingual translation enhances the performance of low-resource languages via knowledge transfer from high-resource languages (Gu et al., 2018; Inaguma et al., 2019b). Given the data scarcity in individual domains, a common practice is to combine the data from various domains to augment the training set (Wang et al., 2020d). Another appealing aspect of multilingual or multi-domain models is their low deployment and maintenance costs compared with numerous models trained for individual language pairs or domains.
Despite the positive knowledge transfer, negative interference has also been observed in multilingual (or multi-domain) training especially when languages (or domains) are dissimilar. Recent studies reveal from the optimization perspective that conflicting gradients in shared parameters is one cause of interference between languages (or domains) (Yu et al., 2020). A promising direction for interference
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
mitigation is to design better strategies of parameter sharing. In some previous works, sharing is based on the similarity between languages (or domains), which require expert knowledge or pre-computed relatedness (Wu et al., 2019). Recent studies also propose branches and components specific to languages (or domains) in addition to shared modules (Bapna and Firat, 2019; Guo et al., 2020).
In this work, we bring the mitigation of language and domain interference under a common umbrella, and tackle it by improving parameter sharing within multi-head attention. We propose strategies to select attention heads for different languages or domains. Instead of sharing everything across languages or domains, our model automatically learns to share heads among a subset of languages or domains. It encourages positive transfer within the subset and preserves their specificity without interference from outside the subset. The major contributions of this work are summarized below:
1. We propose attention head selection to mitigate language or domain interference;
2. The parameter sharing strategies are lightweight and preserve inference efficiency;
3. We extensively evaluate attention sharing strategies on various sequence modeling tasks including speech recognition, text-to-text and speech-to-text translation. Consistent gains are achieved across multiple benchmark datasets.
The paper is structured as follows. Section 2 discusses related works on sequence modeling in multilingual and multi-domain setting. In Section 3, we introduce the proposed strategies of head selection in multi-head attention. Section 4 describes the empirical evaluation, followed by a discussion in Section 5. We conclude this paper in Section 6.

heading: 2 Related Work

text:
Multilingual learning. Multilingual modeling has the potential to improve low-resource language performance through knowledge transfer from high-resource languages, and it draws great interest from researchers in speech recognition and translation (Pratap et al., 2020; Heigold et al., 2013; Johnson et al., 2017; Dabre et al., 2020; Liu et al., 2020; Inaguma et al., 2019a; Li et al., 2020). Although impressive progress has been made for low-resource or zero-shot tasks, it is also found the multilingual model has inferior performance on high-resource tasks due to multilingual interference. In order to address this issue, some works focus on multilingual models with task-specific parameters. Different parameter sharing strategies are examined on Transformer (Sachan and Neubig, 2018). Attention dependent on target languages is proposed to enhance multilingual translation (Blackwood et al., 2018). Treating multilingual modeling as an adaptation problem, Bapna and Firat (2019) first build a universal multilingual model for all languages and then finetune newly added adapters for each language pair. Another thread of work is to increase the model capacity to compensate the performance loss in high-resource languages (Pratap et al., 2020). Shazeer et al. (2017) propose mixture-of-experts and select RNN cells based on input tokens. Lepikhin et al. (2020) integrate a mixture of FFN experts in the GShard model, and later Fedus et al. (2021) propose Switch Transformer to route tokens to different FFN sub-layers. Different from previous works, we propose strategies of attention sharing among languages in the level of attention heads for multilingual modeling.
Multi-domain learning. Similar to multilingual learning, multi-domain learning (MDL) can effectively utilize data from different domains but also suffers from interference due to inter-domain heterogeneity (Saunders, 2021; Pham et al., 2021). Previous works address this issue from two perspectives: optimization and model architecture. In the optimization aspect, attempts have been made to synchronize the learning speed of different tasks (Chen et al., 2018), adjust the gradients of individual tasks to alleviate gradient conflicts (Yu et al., 2020) and apply regularization to achieve better generalization in different domains (Dakwale and Monz, 2017; Khayrallah et al., 2018; Thompson et al., 2019). In terms of model architecture, domain-specific labels (Kobus et al., 2017), word embedding (Zeng et al., 2018a), sub-networks (Wang et al., 2020d) are adopted to address the issue of domain divergence. The architecture can be specified during the general training with the mixed data from multiple domains (Wang et al., 2020d) or during the finetuning in individual domains (Bapna and Firat, 2019). In this work, we deal with domain interference by leveraging domain-specific attention heads in multi-head attention.
Attention selection. Selective self-attention networks propose to apply masking to the inputs and pay more attention to content words (Geng et al., 2020). Liu et al. (2021) select text-related image regions
with attention in multi-modality translation. Compared to these methods, we conduct automatic attention head selection for different tasks and focus on mitigating task interference.

heading: 3 Model

text: In this section, we start with preliminaries of multi-head attention, and introduce our approach to attention interference mitigation. We put multilingual and multi-domain sequence modeling under the same umbrella in this study. For the simplicity of the following discussions, we refer to the two settings as multi-task modeling, where a task is one language or one domain. Different from the standard multi-head attention, our model provides more attention heads than those used in computation. Different subsets of heads are assigned to each task so that partial attention sharing enables knowledge transfer and meanwhile mitigates interference. We introduce latent variables to modulate head selection, and propose strategies to learn the head assignment to different tasks.

heading: 3.1 Preliminary

text:
Multi-head attention. As a core component of Transformer, multi-head attention parameterizes each head with key, query and value transformation matrices (Vaswani et al., 2017). The token representation is transformed into key, query and value vectors via these transformations. Each head assigns the attention of this token over the input sequence based on the matching between its query vector and key vectors of other tokens. The value vectors are weighted by the attention as the contextualized token representation. It is passed through linear projection as the output of the attention head. Suppose that head h has output x(h). Multi-head attention with H heads yields an output x for the given token, which is the concatenation of all head outputs.
x = x(1) ⊕ · · · ⊕ x(h) ⊕ · · · ⊕ x(H), (1) where ⊕ is the vector concatenation. Interference. Maximal parameter sharing aims to learn universal knowledge across languages (Wang et al., 2020e) and domains (Zeng et al., 2018b). To capture the task specificity, different languages or domains compete for model capacity, which is observed as the interference in previous studies. The interference results in degraded performance in jointly trained models. However, few works look into the improvement of parameter sharing within multi-head attention. This study explores head selection strategies to mitigate the interference in multilingual and multi-domain models.

heading: 3.2 Latent Variable for Head Selection

text:
First, we outline our approach to learn a more general-purpose multi-head attention in Transformer from the Bayesian neural network perspective. Suppose that the input sequence is x and the output sequence is y. For conditional sequence modeling tasks such as machine translation, the posterior of p(y | x) can be computed by marginalizing over the posterior of latent variable z, which modulates parameters Θ in the standard Transformer architecture:
p(y | x,Θ) = Ep(z|Θ)[p(y | x, z)] = ∫ p(y | x, z)p(z|Θ) dz (2)
Parameterization of zt. In this work, we define zt as modulating the selection of attention heads by task t. We have zt = {z(h)t }h where z (h) t is a discrete latent variable from Bernoulli distribution indicating whether task t selects attention head h. This modeling choice allows us to prune attention heads, which preserves computation efficiency as well as regularizes training.
Marginalizing over zt is intractable given numerous heads in neural models. Therefore, we use variational inference to derive an approximate solution. Suppose that (xt, yt) is from task t. Specifically, we learn an inference network qφ(zt), which is paramterized with φ, to approximate the true distribution p(zt) and optimize the evidence lower bound (ELBO) of p(y|x):
log p(y | x) ≥ ∑ t ( Eqφ(zt)[log pθ(yt | xt, zt)]− KL(qφ(zt) ‖ p(zt)) ) , (3)
where KL is the KL-divergence between two distributions. In our work, we assume identical probability of each head being selected. Therefore, we have p(zt = 1) = HH′ , where H and H
′ are numbers of selected attention heads and all head candidates.
Training and interference. We use the Gumbel-Softmax reparameterization (Jang et al., 2017) to draw samples of z(h)t from the posterior qφ(z (h) t ). It makes the model end-to-end differentiable, while learning discrete policies of head selection without resorting to policy gradients. We adopt a lightweight estimator of qφ(z (h) t ) by directly learning the logit parameters {φ (h) t }:
qφ(z (h) t ) = exp((φ (h) t (1) + (1))/τ)∑
j∈{0,1} exp((φ (h) t (j) + (j))/τ)
, ∼ G(0, 1) (4)
where G(0, 1) is the Gumbel distribution, and τ is a temperature hyperparameter which increases the discreteness of samples when τ → 0. We will discuss different head selection strategies in Section 3.3, which make binary selection decisions based on real-valued posterior qφ(z (h) t ).
3.3 Attention Selection Strategies
Suppose that the output dimension of multi-head attention is d, and the dimension of each attention head is dH . We provide a large pool of H
′ (H ′ > H) attention head candidates in every Transformer layer, and H ′ is a hyperparameter controlling the search space size of attention selection strategies. The model requires attention outputs to have a consistent dimension d, so each task needs to select exact H heads among H ′ candidates. We introduce two strategies for the attention head selection: subset strategy and group strategy.
Subset strategy. The subset strategy is straightforward, and we compare the posterior {qφ(z(h)t ) : h ∈ [1, H ′]} of all H ′ heads given a task t. A subset of H heads with the highest posterior are selected by the task, and there are CH ′
H subset choices. The subset strategy is described in Fig. 1(a). The binary mask s(h)t indicates whether an attention head h is assigned to task t.
s (h) t = { 1, h ∈ TopH({qφ(z(h)t )}), 0, otherwise,
(5)
where TopH(·) returns the top H heads with the highest values. The outputs of the selected heads are concatenated as the attention output. Note that the subset strategy does not consider the order of the attention heads. For example, when head 2 and 3 are selected, head 2 contributes to the beginning part of attention output. With head 1 and 2 selected, the output of head 2 goes to the last part of the attention output.
Group strategy. We further propose group strategy to preserve the order of attention heads during head selection. Different from the subset strategy, the group strategy first divides H ′ heads into H groups. As is shown in Fig. 1(b), each group contains r = H ′
H candidates. Each task could choose one attention head from each group, and has access to H heads per layer. There are rH possible combinations of heads. The group strategy keeps the head order in that heads from group g only contribute to g’s corresponding dimensions in the attention output. The head with the highest
posterior in its group would be selected by a given task t. We use the binary mask s(h)t to indicate the selection of head h in group g.
s (h) t =
{ 1, h = argmax({qφ(z(h ′) t ) : h
′ ∈ g}), 0, otherwise.
(6)
The output of group g is:
x(g) = ∑ h∈g s (h) t · x(h). (7)
The outputs of H groups are concatenated as the output of the attention module for task t.
With either subset or group strategy, the sequence model is trained to assign attention heads to different tasks to maximize the lower bound in inequality (3). The number of additional parameters {φ(h)t } introduced by our attention selection is only O(T × H ′ × L), where T is the number of tasks, H ′ is the number of head candidates per layer, and L is the number of layers. It is small compared with the total parameter size of the model, and head selection is thus lightweight and memory efficient. Moreover, the head selection is inherently a pruning process. Regardless of the size of head candidates, only a fixed number of attention heads are involved in computation for a given task. Hence our approach is also computationally efficient in model inference.

heading: 4 Experiments

text:
We evaluate sequence models in multilingual and multi-domain settings respectively. Various applications are considered including multilingual machine translation (MT), automatic speech recognition (ASR) and speech translation (ST) in both multilingual and multi-domain settings. We integrate attention selection strategies into the self-attention1 module. Our implementation is based on the FAIRSEQ toolkit (Ott et al., 2019; Wang et al., 2020b). We include widely used sequence models built on multi-head attention as strong baselines below.
1.Transformer (Vaswani et al., 2017). It is a state-of-the-art model in machine translation, which takes texts in source languages as inputs and generates texts in target languages.
2. S2T Transformer (Wang et al., 2020a). As a variant of Transformer for speech processing, S2T Transformer is a stack of a convolutional subsampler and Transformer, where the subsampler processes audio log mel-filter features and sends them to Transformer for text generation.
3. Adapter model (Bapna and Firat, 2019). Adapters have been shown as an effective approach to language and domain adaptation. Task-specific layers are added on top of each Transformer layer in a well-trained (S2T) Transformer. A typical adapter layer consists of two feed-forward sub-layers.
4. Static strategy of head selection. A static strategy assigns each task with a fixed subset of attention heads based on the task similarity (Standley et al., 2020; Sen et al., 2019). In the multilingual setting, we group languages into linguistic families, and each family is assigned with an exclusive set of heads. As for the multi-domain setting, each domain has its own set of attention heads.
We report parameter size and decoding speed as memory and computation efficiency metrics respectively. Decoding speed is measured by the number of tokens decoded per second by one GPU.

heading: 4.1 Machine Translation

text:
The task of machine translation is to translate a text from one language to another. The metric BLEU measures the overlap between model translations and the ground truth (Papineni et al., 2002).
Dataset. We experiment with public multilingual machine translation datasets collected by WMT shared tasks as used by (Liu et al., 2020). The dataset consists of parallel sentences between English and other 14 languages2. Its data statistics are summarized in Appendix A.1. We evaluate models on
1We also tried head selection in the encoder-decoder attention but did not observe big improvements when using it alone or in combination with self-attention head selection.
2The 14 languages are: Chinese (zh), Czech (cs), Estonian (et), Finnish (fi), French (fr), German (de), Gujarati (gu), Kazakh (kk), Latvian (lv), Lithuanian (lt), Romanian (ro), Russian (ru), Spanish (es), Turkish (tr).
both one-to-many (O2M) and many-to-one (M2O) translations, which are translation from English to 14 languages and from 14 languages to English respectively.
Model configurations. The attention selection is based on the source language on the encoder side for M2O translation, and is based on the target language in the decoder part for O2M translation. For both subset and group strategies, the number of attention head candidates is set as 8 in each layer (i.e., H’=8), and only 4 heads (i.e., H=4) are selected for computation. We will discuss how the hyperparameter H ′ affects model performance in Section 5. For static strategy, we group languages into 5 linguistic families3. Each family is assigned with 4 attention heads which are shared by all languages in this family. Therefore, a total of 20 attention heads are used in the static strategy.
Other baselines have 4 attention heads in each Transformer layer. We also include a Transformer baseline with 8 attention heads, which measures the effect of increased attention heads. All models have 6 encoder layers and 6 decoder layers, the embedding dimension is 512 and the feed-forward dimension is 1024. They are trained with a batch size of 131k tokens and a learning rate of 0.0007. For O2M translation, attention selection models and Transformer are trained for 140k steps. As for M2O translation, they are trained for 100k steps. The adapter model is initialized with parameters from the trained Transformer, and tunes adapter layer parameters for 40k steps with Transformer parameters frozen. Adapter layers are added to Transformer for each language direction, and they have an intermediate dimension of 256. The dimension is selected so that the number of parameters (460M) in the adapter model is close to the parameter size (420M) in attention selection models.
Results. We group 14 language directions based on their amount of training data. We have 6 high-resource languages with more than 10M parallel sentences, and 8 low-resource languages with fewer than 10M sentence pairs. Table 1 shows model performance on WMT datasets. More attention heads improve Transformer performance while hurting the decoding speed. In comparison with Transformer with 4 heads, group strategy achieves +0.9 and +0.7 BLEU on average of 14 language directions in O2M and M2O translations respectively at a comparable decoding speed. Transformer with 8 heads and adapter achieve BLEU scores comparable to both group and subset strategies but fall behind in inference efficiency. Static strategy demonstrates comparable performance to group and subset strategies in all metrics except the parameter size.

heading: 4.2 Speech Recognition

text:
The task of Automatic Speech Recognition (ASR) is to transcribe source audios in the same language. Word error rate (WER) is ASR evaluation metric, which measures the difference of model outputs from the ground truth (Klakow and Peters, 2002). Lower WER indicates better recognition.
Model configuration. Models included in the experiments of speech recognition are S2T Transformer, S2T Transformer with adapter layers, S2T Transformer with static, group and subset strategies. With static strategy, we group 8 languages into 2 families4, and each family has an exclusive set of 4 attention heads. Following the setup of (Salesky et al., 2021), all models have 1024 channels in the input convolutional subsampler, 12 encoder layers and 6 decoder layers with 4 attention heads per layer. Again we include the S2T Transformer baseline with 8 heads. The embedding dimension is 256 and the feed-forward dimension is 2048. We set a batch size of 320k tokens and a learning
3(1) Indo-European family: cs, de, es, fr, gu, lt, lv, ro and ru; (2) Estonian family: et; (3) Uralic family: fi; (4) Turkic family: kk and tr; (5) Sino-Tibetan family: zh.
4(1) Afro-Asiatic family: ar; (2) Indo-European family: de, el, es, fr, it, pt and ru.
rate of 0.0005 during training. Attention selection models and S2T Transformer are trained for 250 epochs. Adapter model is initialized with parameters of the trained S2T Transformer, and is then trained for another 200 epochs with only adapter layer parameters tuned. The intermediate dimension of adapter layers is again set as 256. To prevent over-fitting, we stop the model training when the model does not improve on the validation set for 10 epochs. To reduce the performance variance, we average checkpoints of the last 10 epochs, and use the averaged model for evaluation.

heading: 4.2.1 Multilingual Speech Recognition

text:
Dataset. We use the multilingual TEDx (mTEDx) dataset for speech recognition (Salesky et al., 2021). It collects audio recordings from TEDx talks. Eight languages are covered including Arabic (ar), German (de), Greek (el), Spanish (es), French (fr), Italian (it), Portuguese (pt) and Russian (ru).
Results. S2T transformer share all parameters among languages. Attention selection models select attention heads based on the source and target languages. Adapter adds adapter layers based on the language directions. We report the ASR results in Table 2. It brings down 6.1% WER for S2T Transformer to increase from 4 to 8 attention heads. Compared to S2T Transformer with 4 heads, adapter model reduces the WER by 16.1% and subset strategy by 8.8%, while static strategy does not change the performance too much. Group strategy achieves the largest drop of 18.4% in WER of S2T Transformer (H=4) with comparable decoding speed. Moreover, it outperforms S2T Transformer with 8 heads in both speed and WER.

heading: 4.2.2 Multi-Domain Speech Recognition

text:
Dataset. Besides mTEDx data, we include two other public datasets, CoVoST 2 and EuroParl, which are commonly used for speech translation. Since source audios are accompanied by transcripts, we could use their source audio-text data for speech recognition tasks. We investigate multi-domain modeling with these three datasets.
1. CoVoST 2 (Wang et al., 2020c). With Common Voice as the audio source, CoVoST 2 covers speech-to-text translations from 21 languages to English and from English to 15 languages. 2. EuroParl (Iranzo-Sánchez et al., 2020). It provides paired audio-text instances from and into 6 European languages, which are compiled from the debates in European Parliament.
Results. In the multi-domain setting, attention selection models assign different heads to each domain. The static selection strategy provides each domain with an exclusive subset of attention heads. Adapter model adds domain-specific adapter layers to S2T Transformer. Table 3 reports WER of models trained for 400 epochs in three domains: mTEDx, CoVoST 2 and EuroParl respectively. The S2T Transformer jointly trained on multi-domain data (in the row of “Joint S2T Transformer
(H=4)”) reduces WER by 12.9%, 8.6% and 77.7% in three domains respectively, when compared with the models separately trained in individual domains (in the row of “Separate S2T Transformer (H=4)”). This demonstrates the benefits of positive transfer between domains.
The performance of speech recognition could be further improved by the mitigation of the domain interference. Attention selection with group and subset strategies outperform that with the static strategy. Both attention selection and adapter model achieve lower WER than the joint S2T Transformer with both 4 and 8 attention heads. Attention selection with group strategy has the lowest WER on both mTEDx and CoVoST 2 datasets, decreasing WER by 4.0% and 5.0% respectively in comparison with joint S2T Transformer (H=4). The best system on EuroParl is adapter model, yielding a WER reduction by 6.3% than the joint S2T Transformer (H=4).

heading: 4.3 Speech Translation

text:
Now with a focus on the task of speech translation, we again design experiments in multilingual and multi-domain settings. In the multilingual setup, we train translation models with samples in multiple languages to investigate language interference. As for the multi-domain setup, the models are trained with data from multiple domains so that we could look into the domain interference. BLEU serves as the evaluation metric of speech translation systems.
Baselines. We use the same baselines as in speech recognition. As recommended by (Salesky et al., 2021), we initialize the encoders in speech translation with the encoders trained in the task of speech recognition in Section 4.2 for the purpose of improving training efficiency and performance.
Model configurations. All models are trained for up to 400 epochs. Other model configurations in ST are the same as those in ASR.

heading: 4.3.1 Multilingual Speech Translation

text:
To explore language interference, we perform experiments on multilingual speech translation.
Dataset. We again use mTEDx dataset for multilingual speech translation. Besides speech recognition data, mTEDx also collects speech translation data from TEDx talks. Its test set covers 13 language directions. The training data is provided in 10 of these directions, so there are 3 zero-shot directions.
Results. Here the static strategy of head selection groups source languages into two families as in multilingual ASR task, and all target languages fall into the same Indo-European family. Table 4 summarizes the multilingual speech translation results on mTEDx. S2T Transformer has +0.4 BLEU with heads increased to 8. Since adapter model brings in language-specific layers, it cannot deal with zero-shot translations. Group strategy, subset strategy and adapter model bring improvements over S2T Transformer (H=4) which are jointly trained in 13 language directions. It suggests that multiple languages interfere within S2T Transformer whose parameters are shared by all languages. Attention selection with group strategy achieves the best translation performance. In comparison with S2T Transformer (H=4), group strategy achieves an average of +2.1 and +1.9 BLEU in training and zero-shot directions respectively. It leads to +2.0 BLEU on average of all directions.

heading: 4.3.2 Multi-Domain Speech Translation

text:
In this experiment, we investigate interference across domains in the task of speech translation, and evaluate the effectiveness of different models in multi-domain training. The attention selection now is
based on the data domain instead of languages, i.e., samples in different domains would choose their own attention heads. Similarly for adapter model, its adapter layers are domain-specific in this setup.
We again use CoVoST 2 and EuroParl as additional domains. We focus on the 13 language directions in mTEDx test set, and use the subset of CoVoST 2 and EuroParl corpora in the same directions. CoVoST 2 has 5 common directions5 and EuroParl has 11 common directions6 as mTEDx. Details about these datasets are included in Appendix A.1.
Results. Table 5 shows the average BLEU of speech translation in mTEDx, CoVoST 2 and EuroParl. We report results in rows of “Joint S2T Transformer” when S2T Transformers are trained with the mixture of three datasets. The results are included in rows “Separate S2T Transformer” when S2T Transformers are trained on each dataset independently. Zero-shot translations in mTEDx benefit a lot from additional data of CoVoST 2 and EuroParl, as the joint S2T Transformer (H=4) shows an average of +5.4 BLEU over separate S2T Transformer (H=4). However, there is a drop of 1.0 BLEU in its training directions, brought by the interference from CoVoST 2 and EuroParl domains.
Again we observe that static strategy falls behind group and subset strategies. Attention selection with learned strategies and adapter model bring gains to the joint model in individual domains. Compared with the joint S2T Transformer (H=4), adapter model improves mTEDx translation by 0.3 BLEU, CoVoST 2 translation by 0.6 BLEU and EuroParl by 1.0 BLEU on average. The attention selection with group strategy outperforms all other models. Its average BLEU gain over adapter model is 1.6 BLEU in mTEDx, 1.7 BLEU in CoVoST 2 and 0.8 in EuroParl.
5 Discussion
Hyperparameter H ′. The attention selection models set a hyperparameter H ′ as the total number of attention head candidates in multi-head attention, which controls the search space of attention sharing strategies. We now explore how the performance varies with H ′ for group and subset strategies.
Evaluated on the task of multilingual speech recognition, models have the same hyperparameters as those in multilingual ASR experiments except for H ′. Attention selection models are configured with H ′ = 4, 8, 12, 16 respectively, and Figure 2 shows the change of WER with H ′.
When H ′ = 4, there is no attention selection and all attention heads are shared by different languages. We observe a large drop of error rate asH ′ increases from 4 to 8. For the subset strategy, WER keeps decreasing when the number of head candidates grows from 4 to 16. As for group strategy, H ′ = 8 is the optimal hyperparameter on the ASR task. As we continue
5{es, fr, it, pt, ru}-en 6es-{en, fr, it, pt}, fr-{en, es, pt}, it-{en, es}, pt-{en, es}
increasing H ′ to 12 and 16, the error rate increases a bit. The performances of subset and group strategies are close when H ′ = 16.
The search space of group strategy is a strict subset of the space of subset strategy. However, we observe that group strategy shows comparable or better performance than subset strategy across tasks, including MT, ASR and ST. One possible explanation is that group strategy keeps the head order information while subset strategy does not. With a larger pool of head candidates, there is less sharing among tasks. The performance of the group strategy degrades a bit due to less positive transfer dependent on attention sharing. As for the subset strategy, better head assignments are learned in the enlarged search space.
Attention Sharing among Languages. We now analyze the attention sharing pattern among languages. Take the multilingual model on mTEDx speech recognition as an example, whose head selection is learned with group strategy. We count the number of heads shared by each language pair in the model, and visualize it with a heatmap in Fig. 3, where the darkness reflects the amount of sharing. The diagonal cells in the heatmap corresponds to the number of attention heads used by each language, i.e., the total number of attention heads in all layers.
For European languages including Spanish (es), French (fr), Italian (it) and Portuguese (pt), their shared attention heads are fewer in decoder than in encoder. This seems contradicted with previous findings that parameter sharing is beneficial for languages with high linguistic proximity. We note that they are high-resource languages in mTEDx corpus, which is also justified by their relatively lower WER. Their data is sufficient to learn good speech recognition, and sharing parameters with other languages hurt the preservation of the language specificity. This explains why the high-resource European languages do not share too many heads in the learned group strategy.
Another pattern we observe from Fig. 3 is that low-resource languages tend to share more attention heads with high-resource languages. For example, Arabic (ar) and Russian (ru) have relatively more sharing with Italian (it) than other languages. Low-resource languages benefit from the knowledge transfer from high-resource ones. Due to page limit, we include more discussions in Appendix A.3.

heading: 6 Conclusion

text:
Research efforts in multilingual and multi-domain modeling have been driven by the increasing need to improve data efficiency and model performance. In this work, we propose head selection strategies to allow attention heads to be shared or specialized for different languages or domains. It effectively mitigates interference within multi-head attention which is a core part of strong sequence models, and demonstrates good empirical gains in various text generation tasks.
This work has several limitations left for future research. We did not explore head selection based on both language and data domain. We did not analyze model fairness and robustness. As a technology used for text generation, the model might have systemic bias or produce inappropriate outputs.

title: Pay Better Attention to Attention: Head Selection in Multilingual and Multi-Domain Sequence Modeling


summaries: The paper studies the test error, as well as the bias and variance, of random feature regression in the asymptotic regime, in a setting with covariate shift — i.e., when the covariance matrix of the features is different at testing time than at training time. The paper defines a partial ordering over covariance shifts, and it proves that this corresponds to increased or decreased performance compared to the case of no shift. It also obtains a few results that relate these notions to overparameterization.

summaries: This paper presents a theoretical study of covariate shift under the setting of kernel ridge regression in the asymptotic regime. They charaterize the limiting test error, bias, and variance in this setting. They relates these with a quantity named hardness. They claim that their findings agree with several interesting empirical phenomena.

summaries: This paper studies random feature regression in high dimensional setting with the presence of covariate shift between train and test data. The authors compute the high-dimensional asymptotic limit of the test error for Gaussian covariates with a shift in the covariance matrix. They also provide results characterizing the bias and variance decomposition.

summaries: The authors study covariate shift in over-parameterized random features models.  They provide an analysis of the test error, as well as an analysis of the decomposition of the test error into the bias and variance components.  This results in a partial ordering over test errors WRT the "hardness" of the covariate shift.  Further, the theoretical results provide a concrete explanation of the linear trend that has been observed in past work between the test error on clean and corrupted error.

summaries: This work provides theoretical insight for why overparameterized models perform well when there is some covariate shift between the training and test sets. The authors' characterize their results by introducing a notion of "relative hardness". They present empirical evidence to match their intuition, and corroborate the main results in the paper using several empirical benchmarks.

abstractText: A significant obstacle in the development of robust machine learning models is covariate shift, a form of distribution shift that occurs when the input distributions of the training and test sets differ while the conditional label distributions remain the same. Despite the prevalence of covariate shift in real-world applications, a theoretical understanding in the context of modern machine learning has remained lacking. In this work, we examine the exact high-dimensional asymptotics of random feature regression under covariate shift and present a precise characterization of the limiting test error, bias, and variance in this setting. Our results motivate a natural partial order over covariate shifts that provides a sufficient condition for determining when the shift will harm (or even help) test performance. We find that overparameterized models exhibit enhanced robustness to covariate shift, providing one of the first theoretical explanations for this ubiquitous empirical phenomenon. Additionally, our analysis reveals an exact linear relationship between the in-distribution and out-of-distribution generalization performance, offering an explanation for this surprising recent observation.

authors:
B. Adlam
J. Levinson
J. Pennington

title: A random matrix perspective on mixtures of nonlinearities for deep learning

venue: arXiv preprint arXiv:1912.00827

authors:
B. Adlam
J. Pennington

title: The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization

venue: Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 74–84. PMLR

title: Understanding double descent requires a fine-grained bias-variance decomposition

venue: H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 11022–11032. Curran Associates, Inc.

authors:
M.S. Advani
A.M. Saxe
H. Sompolinsky

title: High-dimensional dynamics of generalization error in neural networks

venue: Neural Networks, 132:428–446

authors:
M. Banna
F. Merlevède
M. Peligrad

title: On the limiting spectral distribution for a large class of symmetric random matrices with correlated entries

venue: Stochastic Processes and their Applications, 125(7):2700–2726

authors:
M. Banna
J. Najim
J. Yao

title: A CLT for linear spectral statistics of large random informationplus-noise matrices

venue: Stochastic Processes and their Applications, 130(4):2250–2281

authors:
P.L. Bartlett
A. Montanari

title: and A

venue: Rakhlin. Deep learning: a statistical viewpoint

authors:
C.J. Becker
C.M. Christoudias
P. Fua

title: Non-linear domain adaptation with boosting

venue: C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 26. Curran Associates, Inc.

authors:
M. Belkin
D. Hsu
S. Ma
S. Mandal

title: Reconciling modern machine-learning practice and the classical bias–variance trade-off

venue: Proceedings of the National Academy of Sciences, 116(32):15849–15854

authors:
S. Ben-David
J. Blitzer
K. Crammer
F. Pereira

title: Analysis of representations for domain adaptation

venue: B. Schölkopf, J. Platt, and T. Hoffman, editors, Advances in Neural Information Processing Systems, volume 19. MIT Press

authors:
C. Cortes
M. Mohri

title: Domain adaptation and sample bias correction theory and algorithm for regression

venue: Theor. Comput. Sci., 519:103–126

authors:
A. D’Amour
K. Heller
D. Moldovan
B. Adlam
B. Alipanahi
A. Beutel
C. Chen
J. Deaton
J. Eisenstein
M.D. Hoffman

title: Underspecification presents challenges for credibility in modern machine learning

venue: arXiv preprint arXiv:2011.03395,

authors:
Y. Deshpande
A. Montanari

title: Sparse pca via covariance thresholding

venue: Journal of Machine Learning Research, 17(141):1–41

authors:
E. Dobriban
S. Wager

title: High-dimensional asymptotics of prediction: Ridge regression and classification

venue: The Annals of Statistics, 46(1):247 – 279

authors:
J. Duchi
T. Hashimoto

title: and H

venue: Namkoong. Distributionally robust losses for latent covariate mixtures

authors:
J.C. Duchi
H. Namkoong

title: Learning models with uniform performance via distributionally robust optimization

venue: The Annals of Statistics, 49(3):1378 – 1406

title: The matrix dyson equation and its applications for random matrices

venue: arXiv preprint arXiv:1903.10060

authors:
L. Erdős
H.-T. Yau
J. Yin

title: Bulk universality for generalized wigner matrices

venue: Probability Theory and Related Fields, 154(1-2):341–407

authors:
R.R. Far
T. Oraby
W. Bryc
R. Speicher

title: Spectra of large block matrices

venue: arXiv preprint cs/0610045

authors:
X. Glorot
A. Bordes
Y. Bengio

title: Domain adaptation for large-scale sentiment classification: A deep learning approach

venue: Proceedings of the 28th International Conference on International Conference on Machine Learning, ICML’11, page 513–520, Madison, WI, USA

authors:
I. Goodfellow
Y. Bengio
A. Courville

title: Deep learning

venue: MIT press

authors:
I.J. Goodfellow
J. Shlens
C. Szegedy

title: Explaining and harnessing adversarial examples

venue: Y. Bengio and Y. LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings

title: Percolation

venue: Percolation, pages 1–31. Springer

authors:
T. Hastie
A. Montanari
S. Rosset
R.J. Tibshirani

title: Surprises in high-dimensional ridgeless least squares interpolation

venue: arXiv preprint arXiv:1903.08560

authors:
J.W. Helton
T. Mai
R. Speicher

title: Applications of realizations (aka linearizations) to free probability

venue: Journal of Functional Analysis, 274(1):1–79

authors:
D. Hendrycks
S. Basart
N. Mu
S. Kadavath
F. Wang
E. Dorundo
R. Desai
T. Zhu
S. Parajuli
M. Guo
D. Song
J. Steinhardt
J. Gilmer

title: The many faces of robustness: A critical analysis of out-of-distribution generalization

venue: In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV),

authors:
D. Hendrycks
T.G. Dietterich

title: Benchmarking neural network robustness to common corruptions and perturbations

venue: 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net

authors:
S. Ioffe
C. Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: Proceedings of the 32nd International Conference on International Conference on Machine Learning - Volume 37, ICML’15, page 448–456. JMLR.org

authors:
A. Jacot
F. Gabriel
C. Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

venue: Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS’18, page 8580–8589, Red Hook, NY, USA

authors:
D. Kobak
J. Lomond
B. Sanchez

title: The optimal ridge penalty for real-world highdimensional data can be zero or negative due to the implicit ridge regularization

venue: Journal of Machine Learning Research, 21(169):1–16

authors:
P.W. Koh
S. Sagawa
H. Marklund
S.M. Xie
M. Zhang
A. Balsubramani
W. Hu
M. Yasunaga
R.L. Phillips
I. Gao
T. Lee
E. David
I. Stavness
W. Guo
B. Earnshaw
I. Haque
S.M. Beery
J. Leskovec
A. Kundaje
E. Pierson
S. Levine
C. Finn
P. Liang

title: Wilds: A benchmark of in-the-wild distribution shifts

venue: Proceedings of the 38th International Conference on Machine Learning,

authors:
A. Kumar
T. Ma
P. Liang

title: Understanding self-training for gradual domain adaptation

venue: Proceedings of the 37th International Conference on Machine Learning,

authors:
J. Lee
Y. Bahri
R. Novak
S.S. Schoenholz
J. Pennington
J. Sohl-Dickstein

title: Deep neural networks as gaussian processes

venue: 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings. OpenReview.net

authors:
J. Lee
S. Schoenholz
J. Pennington
B. Adlam
L. Xiao
R. Novak
J. Sohl-Dickstein

title: Finite versus infinite neural networks: an empirical study

venue: H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 15156–15172. Curran Associates, Inc.

authors:
Z. Liao
R. Couillet
M.W. Mahoney

title: A random matrix analysis of random fourier features: beyond the gaussian kernel

venue: a precise phase transition, and the corresponding double descent. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 13939–13950. Curran Associates, Inc.

authors:
M. Long
H. Zhu
J. Wang
M.I. Jordan

title: Unsupervised domain adaptation with residual transfer networks

venue: Proceedings of the 30th International Conference on Neural Information Processing Systems, NIPS’16, page 136–144, Red Hook, NY, USA

title: Deep transfer learning with joint adaptation networks

venue: Proceedings of the 34th International Conference on Machine Learning,

authors:
C. Louart
Z. Liao
R. Couillet

title: et al

venue: A random matrix approach to neural networks. The Annals of Applied Probability, 28(2):1190–1248

authors:
A. Madry
A. Makelov
L. Schmidt
D. Tsipras
A. Vladu

title: Towards deep learning models resistant to adversarial attacks

authors:
H. Mania
S. Sra

title: Why do classifier accuracies show linear trends under distribution shift

authors:
S. Mei
A. Montanari

title: The generalization error of random features regression: Precise asymptotics and the double descent curve

venue: Communications on Pure and Applied Mathematics,

authors:
G. Mel
S. Ganguli

title: A theory of high dimensional regression with arbitrary correlations between input features and target functions: sample complexity, multiple descent curves and a hierarchy of phase transitions

authors:
J. Miller
K. Krauth
B. Recht
L. Schmidt

title: The effect of natural distribution shift on question answering models

authors:
J.P. Miller
R. Taori
A. Raghunathan
S. Sagawa
P.W. Koh
V. Shankar
P. Liang
Y. Carmon
L. Schmidt

title: Accuracy on the line: on the strong correlation between out-of-distribution and in-distribution generalization

authors:
J.A. Mingo
R. Speicher

title: Free probability and random matrices

venue: volume 35. Springer

authors:
Z. Nado
S. Padhy
D. Sculley

title: A

venue: D’Amour, B. Lakshminarayanan, and J. Snoek. Evaluating prediction-time batch normalization for robustness under covariate shift. arXiv preprint arXiv:2006.10963

authors:
V. Nagarajan
J.Z. Kolter

title: Uniform convergence may be unable to explain generalization in deep learning

venue: H. Wallach, H. Larochelle, A. Beygelzimer, F. d Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.

title: Priors for infinite networks

venue: Bayesian Learning for Neural Networks, pages 29–53. Springer

authors:
B. Neyshabur
S. Bhojanapalli
D. Mcallester
N. Srebro

title: Exploring generalization in deep learning

venue: I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.

authors:
Y. Ovadia
E. Fertig
J. Ren
Z. Nado
D. Sculley
S. Nowozin
J. Dillon
B. Lakshminarayanan
J. Snoek

title: Can you trust your models uncertainty? evaluating predictive uncertainty under dataset shift

title: A note on the pennington-worah distribution

venue: Electronic Communications in Probability,

authors:
J. Pennington
P. Worah

title: Nonlinear random matrix theory for deep learning

title: The spectrum of the fisher information matrix of a single-hiddenlayer neural network

venue: S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.

authors:
A. Rahimi
B. Recht

title: Random features for large-scale kernel machines

venue: J. Platt, D. Koller, Y. Singer, and S. Roweis, editors, Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc.

authors:
B. Recht
R. Roelofs
L. Schmidt

title: and V

venue: Shankar. Do cifar-10 classifiers generalize to cifar-10? arXiv preprint arXiv:1806.00451

authors:
B. Recht
R. Roelofs
L. Schmidt
V. Shankar

title: Do imagenet classifiers generalize to imagenet? In K

venue: Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 5389–5400. PMLR

authors:
D. Richards
J. Mourtada
L. Rosasco

title: Asymptotics of ridge(less) regression under general source condition

venue: Proceedings of The 24th International Conference on Artificial Intelligence and Statistics,

title: On moments of the inverted wishart distribution

authors:
S. Sagawa
P.W. Koh
T.B. Hashimoto
P. Liang

title: Distributionally robust neural networks

venue: 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net

authors:
L. Schmidt
S. Santurkar
D. Tsipras
K. Talwar
A. Madry

title: Adversarially robust generalization requires more data

authors:
V. Shankar
A. Fang
W. Guo
S. Fridovich-Keil
J. Ragan-Kelley
L. Schmidt
B. Recht

title: Neural kernels without tangents

authors:
M. Sugiyama
M. Krauledat
K.-R. Müller

title: Covariate shift adaptation by importance weighted cross validation

venue: Journal of Machine Learning Research, 8(35):985–1005

authors:
R. Taori
A. Dave
V. Shankar
N. Carlini
B. Recht
L. Schmidt

title: Measuring robustness to natural distribution shifts in image classification

venue: H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 18583–18599. Curran Associates, Inc.

authors:
N. Tripuraneni
B. Adlam
J. Pennington

title: Covariate shift in high-dimensional random feature regression

venue: arXiv preprint arXiv:2111.08234

title: High-Dimensional Probability: An Introduction with Applications in Data Science

venue: Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press

title: High-Dimensional Statistics: A Non-Asymptotic Viewpoint

authors:
D. Wu
J. Xu

title: On the optimal weighted `2 regularization in overparameterized linear regression

venue: H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 10112–10123. Curran Associates, Inc.

authors:
Z. Yang
Y. Bai
S. Mei

title: Exact gap between generalization error and uniform convergence in random feature models

venue: M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 11704–11715. PMLR

authors:
H. Zhao
R.T.D. Combes
K. Zhang
G. Gordon

title: On learning invariant representations for domain adaptation

venue: Proceedings of the 36th International Conference on Machine Learning,

authors:
H. Zhao
S. Zhang
G. Wu
J.M.F. Moura
J.P. Costeira
G.J. Gordon

title: Adversarial multiple source domain adaptation

heading: 1 Introduction

text:
Theoretical justification for almost all machine learning methods relies upon the equality of the distributions from which the training and test data are drawn. Nevertheless, in many real-world applications, this equality is violated—naturally-occurring distribution shift between the training data and the data encountered during deployment is the rule, not the exception [31]. Even non-adversarial changes in distributions can uncover the surprising fragility of modern machine learning models [55, 56, 43, 26, 12, 50]. Such shifts are distinct from adversarial examples, which require explicit poisoning attacks [22]; rather, they can result from mild corruptions, ranging from changes of camera angle or blur [26], to subtle, unintended changes in data acquisition procedures [56]. Moreover, this fragility limits the application of deep learning in certain safety-critical areas [31].
Empirical studies of distribution shift have observed several intriguing phenomena, including linear trends between model performance on shifted and unshifted test distributions [56, 26, 31], dramatic degradation in calibration [50], and surprising spurious inductive biases [12]. Theoretical understanding of why such patterns occur across a variety of real-world domains is scant. Even basic questions such as what makes a certain distribution shift likely to hurt (or help) a model’s performance,
⇤Equal contribution. †Work performed while the author was at Google.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
and by how much, are not understood. One reason that these phenomena have eluded theoretical understanding is that there is often a strong coupling between model and distribution, implying that the effect of a given shift cannot usually be understood in a model-agnostic way. Another reason that satisfactory explanations have remained lacking is that the go-to formalism for studying generalization in classical models, namely uniform convergence theory (see e.g. [66]), may be insufficient to explain the behavior of modern deep learning methods (even in the absence of distribution shift) [47, 68]. Indeed, classical measures of model complexity, such as various norms of the parameters, have been found to lead to ambiguous conclusions [49].
In this paper, we follow a different approach: instead of focusing on worst-case bounds for generic distributions, we study average-case behavior for narrowly specified distributions. While this change in perspective sacrifices generality, it allows us to derive more precise predictions, which we believe are necessary to fully capture the relevant phenomenology. We study a specific type of distribution shift called covariate shift, in which the distributions of the training and test covariates differ, while the conditional distribution of the labels given the covariates remains fixed. Using random matrix theory, we perform an asymptotically exact computation of the generalization error of random feature regression under covariate shift. The random feature model provides a useful testbed to (1) investigate the interplay between various factors such as model complexity, label noise, bias, variance, and covariate shift; (2) rigorously define a model-agnostic notion for the strength of covariate shift; and (3) provide a theoretical explanation for the linear relationships recently observed between in-distribution and out-of-distribution generalization performance [55, 56, 27].

heading: 1.1 Contributions

text:
Our primary contributions are to:
1. Provide a model-agnostic partial order over covariate shifts that is sufficient to determine when a shift will increase or decrease the test error in random feature regression (see Def. 4.1);
2. Compute the test error, bias, and variance of random feature regression for general multivariate Gaussian covariates under covariate shift in the high-dimensional limit (see Sec. 5.1);
3. Prove that overparameterization enhances robustness to covariate shift, and that the error, bias, and variance are nonincreasing functions of the number of excess parameters (see Sec. 5.3);
4. Deduce an exact linear relationship between in-distribution and out-of-distribution generalization performance, offering an explanation for this surprising recent empirical observation (see Sec. 5.4).

heading: 1.2 Related work

text:
There is extensive literature on the empirical analysis of distribution shift in all of its myriad forms, ranging from domain adaptation [62, 20, 8, 70, 69, 37, 36] to defenses against adversarial attacks [39, 60] to distributionally robust optimization [59, 15, 16], among many others. Interestingly, for naturally occurring distribution shifts [31, 26], standard robustness interventions provide little protection [56, 63]. Indeed, empirical risk minimization on clean, unshifted training data often performs better on out-of-distribution benchmarks than more sophisticated methods [31]. One of the most striking observations in the context of natural distribution shifts is that model robustness improves with the classifier’s accuracy [56, 63, 26, 43]. For example, if a classifier’s accuracy increases by 1.0% on the unshifted CIFAR-10 test set, this tends to increase its accuracy by 1.7% on the CIFAR-10.1 dataset (a dataset with natural distribution shift) [56]. Moreover, such linear trends between the unshifted and shifted measures of error have now been observed in several contexts [56, 63, 43, 40, 44].
The number of theoretical works studying the impact of distribution shift on generalization is far smaller. One pioneering work provides VC-dimension-based error bounds for classification that are augmented by a discrepancy measure between source and target domains [10], while another demonstrates a similar class of uniform convergence-based results in the setting of kernel regression [11]. Recent work shows that learning domain-invariant features is insufficient to guarantee generalization when the class-conditional distributions of features may shift [69]. When the source domain gradually shifts toward the target domain, non-vacuous margin-based bounds for self-training can be established [32]. In [40], assumptions based on model similarity are used to help explain why classifiers exhibit linear trends between their accuracies on shifted and unshifted test sets [56, 43].
Our technical tools build on a series of works that have studied the exact high-dimensional limit of the test error for a growing class of model families and data distributions. In the context of linear models, recent work analyzes ridge regression for general covariances and a general non-isotropic source condition on the parameters which generate the targets [57], extending earlier work studying minimum-norm interpolated least squares and ridge regression in the random design setting [9, 14, 24]. The non-isotropy of source parameter effectively induces a shift on the bias term of this model, but the phenomenon is distinct from the covariate shifts we study here. Beyond linear regression, random feature models provide a rich but tractable class of models to gain further insight into generalization phenomena [2, 3, 41, 35]. These methods are of particular interest because of their connection to neural networks, with the number of random features corresponding to the network width (or model complexity) [48, 33, 29], and because they serve as a practical method for data analysis in their own right [54, 61]. In this context, a precise characterization of the gaps between uniform convergence and the (asymptotic) exact test error as a function of the sample size and number of random features can be derived [68]. Since this paper’s publication, we released follow-up work considering unequal scales in the training and test distributions and optimal regularization [64]. From the technical perspective, our analytic techniques build upon these works and a series of recent results stemming from the literature on random matrix theory and free probability [53, 52, 1, 2, 38, 51, 19, 45].

heading: 2 Preliminaries

heading: 2.1 Problem setup and notation

text:
As in prior work studying random feature regression [24, 41, 2, 1], we compute the test error in the high-dimensional, proportional asymptotics where the dataset size m, input feature dimension n0, and hidden layer size n1 all tend to infinity at the same rate, with := n0/m and := n0/n1 held fixed. We refer to / as the overparameterization ratio, which is the limit of n1/m and characterizes the normalized complexity of the (random) feature model.
Interestingly, in this high-dimensional limit, the conditional distribution of a linear labeling function is asymptotically equivalent to a wide class of nonlinear teacher functions (see [41, 2] for more details). With this in mind, we consider the task of learning an unknown function from m i.i.d. samples (xi, yi) 2 Rn0 ⇥ R for i 2 {1, . . . ,m}, where the covariates are Gaussian, xi ⇠ N (0,⌃) with positive definite covariance matrix ⌃, and the labels are generated by a linear function parameterized by 2 Rn0 , drawn from N (0, In0). In particular
y(xi) = >xi/ p n0 + ✏i, (1)
where ✏i ⇠ N (0, 2✏ ) is additive label noise on the training points. We study the class of prediction models defined by kernel ridge regression using unstructured random feature maps [54]. The random features are given by a single-layer, fully-connected neural network with random weights. Given a set of training data X = [x1, . . . ,xm] and a prospective test point x, the random features embeddings of the training and test data are given by
F := (WX/ p n0) and f := (Wx/ p n0) , (2)
for a random weight matrix W 2 Rn1⇥n0 with i.i.d. standard Gaussian entries and an activation function : R ! R applied elementwise. The induced kernel is
K(x1,x2) := 1
n1 (Wx1/
p n0) > (Wx2/ p n0) , (3)
and the model’s predictions are given by ŷ(x) = Y K 1Kx, where Y := [y(x1), . . . , y(xm)], K := K(X,X)+ Im, Kx := K(X,x), and 0 is a ridge regularization constant1. Owing to the implicit regularization effect of the nonlinear feature maps [7], in low noise settings the optimal value of can sometimes be negative [30]. For simplicity, we nevertheless make the standard assumption that 0, though we emphasize our techniques readily accommodate negative values.
1We overload the definition of K to include the additive regularization whenever no arguments are present. Also, if = 0 and K is not full-rank, K 1 should be understood as the Moore–Penrose pseudoinverse.
Our central object of study is the expected test loss for a datapoint x ⇠ N (0,⌃⇤) where ⌃⇤ may be different from the training covariance ⌃. The test error (without label noise on the test point) is
E⌃⇤ = E[( >x/ p n0 Y K 1Kx)2] (4)
= Ex, [(E[ŷ(x)] y(x))2]| {z } B⌃⇤ +Ex, [V[ŷ(x)]]| {z } V⌃⇤ , (5)
where the inner expectations defining the bias and variance are computed over W , X , and Y . We decompose the training and test covariance matrices into eigenbases as ⌃ = P n0
i=1 ivivi > and ⌃⇤ =P
n0 i=1 ⇤ i v⇤ i v⇤ i >, where the eigenvalues are in nondecreasing magnitude, i.e. 1  2  . . .  n0 and ⇤1  ⇤2  . . .  ⇤n0 . We define the overlap coefficients
ri := v > i ⌃⇤vi =
n0X
j=1
(v⇤ j · vi)2 ⇤j (6)
to measure the alignment of ⌃⇤ with the ith eigendirection of ⌃. In particular, ri is the induced norm of vi with respect to ⌃⇤. We use t̄r to denote the dimension-normalized trace: for a matrix A 2 Rn⇥n, t̄r(A) = 1
n tr(A). We use kAk1 and kAkF to denote the operator norm and Frobenius
norm of matrix A respectively. Finally, we use x to denote the Dirac delta function centered at x.

heading: 2.2 Assumptions

text:
Regularity assumptions on the spectra of ⌃ and ⌃⇤ are necessary to state the limiting behavior of this system. As in [67], it is not sufficient to consider the spectra of these matrices individually; they must be considered jointly. We do this in an eigenbasis of ⌃. Assumption 1. We define the empirical joint spectral distribution (EJSD) as
µn0 := 1
n0
n0X
i=1
( i,ri) (7)
and assume it converges in distribution to some µ, a distribution on R2+ as n0 ! 1. We refer to µ as the limiting joint spectral distribution (LJSD), and emphasize that this defines the relevant limiting properties of the train and test distributions2. Additionally, we require that lim sup
n0 max(k⌃k1, k⌃⇤k1)  C for a constant C.
Often we use ( , r) for random variables sampled jointly from µ and denote the marginal of under µ with µtrain. The conditional expectation E[r| ] is an important object in our study. We frequently overload the notation E[r| ] to view it as a function of , and we assume the following for simplicity.
Assumption 2. µ is either absolutely continuous or a finite sum of delta masses. Moreover, the expectations of and r are finite.
When the eigenspaces of ⌃ and ⌃⇤ are aligned and ri = ⇤i = ( i) for some smooth function , the support of the LJSD degenerates. Here, Assump. 1 is essentially equivalent to assuming the empirical spectral distribution of ⌃ converges in distribution to some µtrain, which is a standard assumption in the regression literature [14, 41]. One special case of note is when there is no shift, i.e. is the identity, in which case the LJSD degenerates to µ; defined by
µ;( , r) := µtrain( ) (r) , i.e. ( , ) ⇠ µ; for ⇠ µtrain . (8) As our analysis will eventually take place in the high-dimensional limit, we further define the asymptotic scales of the training and test covariances as s := limn0!1 t̄r(⌃) = Eµ[ ] and s⇤ := limn0!1 t̄r(⌃
⇤) = Eµ[r] under the limiting behavior specified in Assump. 1. Throughout this paper, we also enforce the following standard regularity assumptions on the activation functions to ensure the existence of the moments and derivatives we compute. Assumption 3. The activation function : R ! R is assumed to be differentiable almost everywhere. We assume that, | (x)|, | 0(x)|  c0 exp(c1x) for constants c0, c1.
2Note that the EJSD depends not only on ⌃ and ⌃⇤ but also on a choice of eigendecomposition for ⌃ when it has repeated eigenvalues. However, all possible choices for the EJSD form an equivalence class, and the ambiguity does not affect later definitions and conclusions. See Sec. A3.

heading: 2.3 A simple family of diatomic distributions

text:
As the above assumptions allow such a general class of covariance structures, it is useful to consider our results in the context of a simple family of distributions that readily admits a simple interpretation. Definition 2.1. For ↵ 1 and ✓ 2 R, we define the family of (↵, ✓)-diatomic LJSDs with ✓-powerlaw shifts as
µdiatomic ↵,✓
:= 1
↵+ 1 (↵,C↵✓) +
↵
↵+ 1 (↵ 1,C↵ ✓), (9)
where C is a normalization constant chosen so that E µ
diatomic ↵,✓
[r] = 1. Note that µdiatomic ↵,✓ is the limit of
⌃ij :=
8 ><
>: ↵ if i = j and i  b n01+↵c ↵ 1 if i = j and i > b n01+↵c 0 if i 6= j
and ⌃⇤ := 1
t̄r(⌃✓) ⌃✓ . (10)
This simple two-parameter family of distributions captures the fast eigenvalue decay observed in many datasets in machine learning, for which the covariance spectra are often dominated by several large eigenvalues and exhibit a long tail of many small eigenvalues [34]. Note that the trivial case of ↵ = 1 yields an identity covariance with no shift. For the nontrivial setting ↵ > 1, the exponent ✓ parameterizes the strength of the shift in an intuitive way: when ✓ = 1, there is no shift; when ✓ < 1, ↵✓ < ↵, so the large eigendirections of the training distribution are suppressed in the test distribution, suggesting that the shift makes learning harder; when ✓ > 1, ↵✓ > ↵, so the large eigendirections of the training distribution are further emphasized in the test distribution, suggesting that the shift makes learning easier. We will return to the notion of shift strength in Secs. 3 and 4.

heading: 3 Motivating example: linear regression

text:
We first consider the relatively simple case of ridgeless linear regression (LR), which will help build some intuition for the more general analysis of random feature regression in Sec. 5.1. Assuming the labels are generated by the linear model defined above, i.e. yi = >xi/ p n0 + "i, the estimator is given by ̂ = (XX>) 1XY , and the test risk (see Eq. (4)) has the following simple form. Proposition 3.1. For fixed dimension n0 and sample size m>n0+1, the test error of LR is given by
ELR⌃⇤ = 2 ✏ n0 m n0 1 t̄r(⌃⇤⌃ 1) = 2 ✏ n0 m n0 1 1 n0
n0X
i=1
ri i . (11)
Under Assump. 1, as n0,m ! 1 with = n0/m fixed, ELR⌃⇤ ! ELRµ = 2✏ /(1 )Eµ[r/ ].
One immediate question is whether a given shift will increase or decrease the test error relative to ELR⌃ . While the precise answer is of course determined by the value of t̄r(⌃
⇤⌃ 1), it useful for the subsequent analysis to develop an understanding of the individual contributions to this term. Similar decompositions of the test error into eigenspaces have proved useful in a variety of other contexts, e.g. [42, 4]. We begin with a specific example in the setting of the finite-dimensional analog of Def. 2.1. Example 3.1. For the finite form of the (↵, ✓)-diatomic density defined in Eq. (10), the test error of LR is given by
ELR⌃⇤ = 2 ✏ n0 m n0 1 ↵+ w(↵2✓ 1 ↵) 1 + w(↵2✓ 1) for w = 1 n0
n0
1 + ↵
⌫ , (12)
and so @ @✓ ELR⌃⇤ = 2 ✏ n0 m n0 1 2(1 w)w↵2✓ 1(1 ↵2) log(↵) (1+w(↵2✓ 1))2  0, which implies E LR ⌃⇤1  ELR⌃⇤2 whenever ✓1 ✓2, in accordance with the discussion in Sec. 2.3. It follows from Eq. (10) that the condition ✓1 ✓2 not only implies t̄r(⌃⇤1⌃ 1)  t̄r(⌃⇤2⌃ 1), but also that the ratios of overlap coefficients ri,1/ri,2 form a nondecreasing sequence. It is this condition involving all the eigendirections that will generalize to the nonlinear random feature setting in Sec. 4.
The following proposition captures the essence of these considerations in the context of linear regression. See Sec. A4 for the proof.
Proposition 3.2. Let ri,1 and ri,2 denote the overlap coefficients3 of ⌃⇤1 and ⌃⇤2 relative to ⌃. If tr(⌃⇤2) tr(⌃⇤1) and the ratios ri,1/ri,2 form a nondecreasing sequence, then in the setting of Prop. 3.1, ELR⌃⇤2 E LR ⌃⇤1 .
Whereas the (↵, ✓)-diatomic LJSDs explicitly enforce the trace normalizations tr(⌃) = tr(⌃⇤1) = tr(⌃ ⇤ 2) = 1, Prop. 3.2 provides sufficient conditions for the ordering of test errors for non-unit traces. The fact that ELR⌃⇤ scales linearly with the overall scale of ⌃ ⇤ is a unique feature of linear regression and does not generalize to the nonlinear random feature setting. We return to this issue in Sec. 4.

heading: 4 Definition of shift strength

text:
Deriving conditions on whether a shift will hurt or help a model’s performance is crucial to building an understanding of covariate shift. Motivated in part by the above results for linear regression, and in part by the results for random feature regression that we present in Sec. 5.3, we introduce the following definition of shift strength, which is a direct generalization of the conditions of Prop. 3.2: Definition 4.1. Let µ1 and µ2 be LJSDs with the same marginal distribution of , denoted µtrain. If the asymptotic overlap coefficients are such that Eµ1 [r| ]/Eµ2 [r| ] is nondecreasing as a function of on the support of µtrain and Eµ1 [r]  Eµ2 [r], we say µ1 is easier than µ2 (or µ2 is harder than µ1), and write µ1  µ2. Comparing against the case of no shift µ;, we say µ1 is easy when µ1  µ; and hard when µ1 µ;.
A priori, there is little reason to hope that such a model-independent definition of shift strength would adequately characterize a shift’s impact on the total error, bias, or variance of a given model. Even for the relatively simple case of random feature kernel regression, the nonlinear feature maps of Eq. (2) would seem to inextricably couple the covariance distribution to the model.
Nevertheless, as we show in Sec. 5.1, the coupling between model and shift simplifies considerably in the high-dimensional proportional asymptotics. It is characterized by a handful of constants that depend solely on the overall covariance scale Eµ[r] and a collection of functionals of µ, whose magnitudes can be bounded in terms of the ratio Eµ[r| ]/Eµ; [r| ]. The conditions that Def. 4.1 places on Eµ[r] and Eµ[r| ]/Eµ; [r| ] can be augmented by various constraints on the model to derive bounds on how the total error and bias will respond to a shift of a given strength. This perspective introduces considerable complexity and we present the details of this analysis elsewhere.
In this work, we focus on a simpler, surprising result: by merely normalizing the scales of the covariate distributions (i.e. enforcing s = s⇤), Def. 4.1 provides a model-independent definition of shift strength that determines how random feature models respond to shifts of different strength. This observation motivates the following assumption. Assumption 4. The training and test covariance scales are equal, Eµ[ ] = Eµ[r], i.e. s = s⇤.
We emphasize that Assump. 4 reflects common practice for many models and data modalities, as preprocessing techniques such as standardization are ubiquitous and many architectural components such as layer- or batch-normalization achieve a similar effect [21, 28, 46].

heading: 5 Covariate shift in random feature kernel regression

heading: 5.1 Main results

text:
Our main results characterize the high-dimensional limits of the test error, bias, and variance of the nonlinear random feature model of Sec. 2. Before stating them, we first introduce some additional constants that capture the effect of the nonlinearity . For z ⇠ N (0, s), define
⌘ := V[ (z)] , ⇢:= ( 1 s E[z (z)])2 , ⇣ := s⇢ , and ! := s(⌘/⇣ 1) . (13)
Our results also depend on the covariance spectra through two sets of functionals of µ,
Ia,b(x) := Eµ ⇣ a ( + x ) b ⌘ and I⇤ a,b (x) := Eµ ⇣ r a 1 ( + x ) b ⌘ . (14)
3Recall the definition of the overlap coefficients in Eq. (6).
(d) (f)(e)
(c)(b)(a) Total error Bias Variance
Sh ift
p ow
er , θ
Overparameterization ratio, ϕ/ψ Overparameterization ratio, ϕ/ψ Overparameterization ratio, ϕ/ψ
0.01 0.10 1 10 100 -3
-2
-1
0
1
2
3
0.01 0.10 1 10 100 -3
-2
-1
0
1
2
3
0.01 0.10 1 10 100 -3
-2
-1
0
1
2
3
Sh ift
p ow
er , θ
Sh ift
p ow
er , θ
0.025
0.05
0.1
0.2
0.4
0.8
1.6
Theorem 5.1. Under Assumps. 1, 2, 3 and 4, as n0, n1,m ! 1 the test error E⌃⇤ converges to Eµ = Bµ + Vµ, with the bias Bµ and variance Vµ given by
Bµ = I⇤1,2 (15)
Vµ = ⇢
@x
@
✓ I1,1(! + I1,2)(! + I⇤1,1) + 2 ⌧̄I1,2I⇤2,2
+ ⌧I2,2(! + I⇤1,2) + 2" ⇣ (! + I1,2)(! + I⇤1,1) +
⌧̄I⇤2,2
⌘◆ , (16)
where x is the unique nonnegative real root of x = 1 ⌧ !+I1,1 , @x @ = x +⇢ (⌧ / +⌧̄)(!+ I1,2) , and
⌧ =
p ( )2 + 4x /⇢+
2 and ⌧̄ =
1
+
⌧ 1 . (17)
Numerical predictions from Thm. 5.1 can be obtained by first solving the self-consistent equation for x by fixed-point iteration, x 7! 1 ⌧
!+I1,1 , and then plugging the result into the remaining terms. Fig. 1 shows excellent agreement between these asymptotic predictions and finite-size simulations.
At times we will find it convenient to consider the ridgeless limit of Thm. 5.1. By carefully expanding x and ⌧ for small , it is straightforward to obtain the following corollary. Corollary 5.1. In the setting of Thm. 5.1, as the ridge regularization constant ! 0, Eµ = Bµ+Vµ with Bµ given in Eq. (15) and Vµ given by
Vµ = | |x( 2 " + I1,1)(! + I⇤1,1) +
8 < : x ⇣ 1 x(! 2 ") 1 x2I2,2 ⌘ I⇤2,2 x 2 I2,2
x2 I2,2
! + I⇤1,2
<
, (18)
where x is the unique positive real root of x = min(1, / ) !+I1,1 .
Taking (x) = x and ! 0 in Cor. 5.1 yields an expression for the test error of ridgeless linear regression that agrees with [24] and with the asymptotic form of Prop. 3.1 (see Sec. A4.2).

heading: 5.2 Harder shifts increase the bias and test error

text:
The bias and variance in Thm. 5.1 depend on the covariate shift exclusively through I⇤ a,b — all other terms such as x, ⌧ , ⌧̄ , and Ia,b only depend on the marginal of under µ. The functionals I⇤a,b generalize the simple ratio of overlap coefficients to eigenvalues, Eµ[r/ ], that characterizes the error for linear regression (indeed, I⇤0,0 = Eµ[r/ ]). In contrast, the error in the random feature setting is a combination of multiple such terms. Nevertheless, Def. 4.1 enables comparisons of the individual I⇤ a,b
functionals, which provide sufficient conditions to order the error and bias. Proposition 5.1. Consider two LJSDs such that µ1  µ2 (see Def. 4.1). Then, in the setting of Thm. 5.1, Bµ1  Bµ2 and, if 2"  !, Eµ1  Eµ2 .
Prop. 5.1 shows that Def. 4.1 provides an essentially model-independent condition to determine the impact of covariate shift on the test error4. Interestingly, both the bias (which arises from both regularization and model misspecification) and the total error (which has additional variance contributions from the randomness induced by W , X , and ✏) respond to shifts in tandem in the regime of small label noise. (For large label noise, the variance can dominate the error and cause violations of monotonicity; see Sec. 5.5.) Prop. 5.1 is illustrated for the (↵, ✓)-diatomic LJSD in Fig. 1: following the vertical lines upward in (a), (b), or the x-axis rightward in (e) yields easier shifts and a corresponding decrease in the bias and total error.

heading: 5.3 The benefit of overparameterization

text:
While Prop. 5.1 shows that harder shifts increase the error, it is natural to wonder whether this increase can be mitigated by judicious model selection. In practice, empirical investigations have shown that the performance of large, overparameterized models tends to deteriorate less under distribution shift than their smaller counterparts [26]. We obtain a number of theoretical results that formally prove the benefit of overparameterization in our random feature setting.
First, we show that the bias decreases (or stays constant) when additional random features are added, which increases the model capacity and accords with the intuition of the bias as a measure of the model’s ability to fit the data. Proposition 5.2. In the setting of Thm. 5.1, the bias Bµ is a nonincreasing function of the overparameterization ratio / .
In contrast to the bias, which is monotonic for all overparameterization ratios, the variance can exhibit nonmonotonic behavior in the underparameterized regime. On the other hand, the following proposition shows that in the overparameterized regime, the variance is also nonincreasing. Note that our proof requires the setting of ridgeless regression ( = 0), but numerical investigation suggests this condition may not be necessary (see Fig. 1). Proposition 5.3. In the setting of Cor. 5.1 and in the overparameterized regime (i.e. < ), the variance Vµ is a nonincreasing function of the overparameterization ratio / .
The explosion of variance at the interpolation threshold and then its subsequent decay have been demonstrated in previous exact asymptotic studies of random feature regression in the absence of covariate shift, in stark contrast to what classical theory would suggest [2, 41]. Prop. 5.3 confirms the existence of analogous behavior under covariate shift.
Taken together, Props. 5.2 and 5.3 imply that some of the benefits of overparameterization extend to models evaluated out-of-distribution. An additional benefit is that overparameterized models are more robust: the difference in error between unshifted and shifted test distributions is smaller for larger models. A formal statement of this enhanced robustness is given in the following result. Proposition 5.4. Consider two LJSDs such that µ1  µ2 (see Def. 4.1). Then, in the setting of Cor. 5.1 and in the overparameterized regime (i.e. < ), the generalization gap Eµ2 Eµ1 is a nonincreasing function of the overparameterization ratio / .
4The weak model-dependence arising from the condition 2" < ! can instead be regarded as a small label-noise condition that can always be satisfied for ! > 0, which is true for any nonlinear (see Sec. A1).
= -0.5 (hard) = 1.5 (easy)
Identity line
-2.0 -1.5 -1.0 -0.5 0.0 -2.0
-1.5
-1.0
-0.5
0.0
(S hi
fte d
te st
e rro r) −E μ
(Unshifted test error)−Eμ∅
ImgNetv2 Grey ImgNetv2 TopImgs Linear fit (slope = 1.37) Linear fit (slope = 0.95) Identity line
30 40 50 60 70 80 90 30
40
50
60
70
80
90
Sh ift
ed Im
ag eN
et A
cc
ImageNet Val Acc
(b)(a)
Figure 2: Linear relationship between in-distribution and out-of-distribution generalization error. (a) Asymptotic predictions for shifted versus unshifted error for models with varying degrees of overparameterization / > 1, obtained via Cor. 5.1 for the (3, ✓)-diatomic LJSD (Eq. (9)) with = n0/m = 0.5, = ReLU, 2" = 0.01 and two different values of the shift-power ✓. Markers represent simulations for n0 = 512. The negated errors are plotted so that performance improves left to right and bottom to top, in order to match the behavior of the accuracy metric. (b) Reproduction of the empirical results of [56, 63], showing the relationship between the classification accuracy of various models on the original ImageNet test set and two shifted ImageNet datasets: a “hard" dataset with greyscale corruptions, Grey, and an “easy" dataset with high inter-annotator agreement, TopImgs. In both (a) and (b), the slope is greater than one for the hard shift and less than one for the easy shift, in accordance with Prop. 5.5.
In Fig. 1, Props. 5.2, 5.3 and 5.4 are illustrated. Following the horizontal lines rightward in (a), (b), and (c) or the x-axis rightward in (d) and (f) leads to models with more parameters. The monotonicity of the bias across the whole range of parameterization is evident, as is the necessity of considering the monotonicity of the variance and generalization gap only when > .

heading: 5.4 Linear trends between in-distribution and out-of-distribution generalization

text:
We have discussed how overparameterization yields improvements on both unshifted and shifted test distributions, which hints that these two quantities are positively correlated. Indeed, recent work has suggested increasing model size as a path to increased robustness [26]. Additional empirical studies have further refined this observation by discovering a linear relationship between the performance of models of varying complexity on unshifted and shifted data [56, 63]. In the context of ridgeless random feature regression, we provide a formal proof of this linear relationship. Proposition 5.5. In the setting of Cor. 5.1 and in the overparameterized regime (i.e. < ),
Eµ = E0 + ⇣! + I⇤1,1 ! + I1,1 ⌘
| {z } SLOPE
Eµ; , (19)
parametrically in the overparameterization ratio / , where E0 and SLOPE are constants independent of / , and Eµ; is the error on the unshifted distribution. Moreover, SLOPE 1 when µ is hard and SLOPE  1 when µ is easy.
Eq. (19) implies a parametrically linear relationship between Eµ and Eµ; by varying / . Prop. 5.5 also makes the nontrivial prediction that an improvement on the unshifted distribution leads to a relatively greater improvement on the shifted distribution when the shift is hard, and to a relatively smaller improvement when the shift is easy. This prediction is corroborated qualitatively in the data from [55, 56, 43]. We plot this linear behavior in Fig. 2, where (a) shows the random feature model and (b) shows an example of data from [56, 63]. The striking similarity in these plots is evident.

heading: 5.5 Importance of assumptions

text:
Props. 5.1, 5.2, 5.3, 5.4 and 5.5 rely on a number of assumptions and conditions; here we show the necessity of some of these prerequisites for our results.
(a) (c)(b)
Prop. 5.1 relies on a small label-noise condition ( 2 "  !) to ensure that the total error is ordered with respect to shift strength. The reason this condition is necessary is that the variance can actually increase as shifts become easier. While Fig. 1 presented a configuration for which the bias, variance, and error all decrease for easier shifts, Fig. 3(a) shows that a decrease is not guaranteed for the variance, and that it can increase even under the small label-noise condition. Moreover, while the bias continues to decrease for large label noise ( 2
✏ > !), the variance can become so large that bias
can no longer offset it, causing the error itself to increase, as seen in Fig. 3(a).
The strict ordering of overlap coefficients in Def. 4.1 are also necessary to guarantee a complete decoupling of the model and the shift strength. In the absence of these conditions, Fig. 3(b) shows how even a single out-of-order overlap coefficient induces a violation of the monotonicity with respect to shift strength suggested by Prop. 5.1 (this example is detailed further in Sec. A8).
Finally, we note that the exact linear relationship between in-distribution and out-of-distribution generalization characterized by Prop. 5.5 in Eq. (19) relies crucially on the overparameterization condition, < , as evidenced in Fig. 3(c), which shows marked nonlinearity in the underparameterized regime. This observation is perhaps unsurprising, as severely underparameterized models tend towards chance predictions, which produce comparable errors on shifted and unshifted data. Indeed, similar nonlinear behavior is seen in the low-accuracy regime for realistic models [56, Figure 17].

heading: 6 Conclusion

text:
We have presented an exact, asymptotic calculation of the test error, bias, and variance for random feature kernel regression in the presence of covariate shift. After defining a partial order over covariate shifts (motivated by the setting of linear regression), we have proved that harder shifts imply increased error. Our results capture many empirical phenomena such as the fact that overparameterization is beneficial even under covariate shift and that a linear relationship exists between the generalization error on shifted and unshifted data. Future directions include extending our results to the nonasymptotic regime, accommodating feature learning and more general neural network models, and investigating the impact of covariate shift for other loss functions.
Acknowledgments and Disclosure of Funding
The authors would like to thank Rodolphe Jenatton, Horia Mania, Ludwig Schmidt, D. Sculley, Vaishaal Shankar, Lechao Xiao, and Steve Yadlowsky for valuable discussions.
This work was performed at and funded by Google. No third party funding was used.

title: Overparameterization Improves Robustness to Covariate Shift in High Dimensions