summaries: The paper aims to remove potential examples with label noise by discarding the ones with large losses in the training procedure. The idea also applies to the setting where instances may contain large noise. The proposed method may have an implicit trade-off between the robust to label noise and feature noise, which explains why the proposed method also has good performances on instance-dependent label noise. The paper is well-written and has sufficient experiments.

summaries: This paper presents ODD, a method that rejects incorrectly labeled / noisy examples from training on the fly. The motivation is sound, that with the capacity of modern neural networks, it's easy to memorize the mislabeled data and thus hurt generalization. If we could reject such mislabeled data, we may be able to get a more generalizable model. The authors made an observation that when training with large learning rate, examples with correct labeling and incorrect labeling exhibits different loss distributions. The authors further noticed that the loss distribution of incorrectly labeled examples can be simulated using eq.(1). Therefore, by setting a threshold that corresponds to a percentile of the incorrectly labeled loss distribution, the authors are able to reject incorrect examples.

abstractText: Memorization in over-parameterized neural networks can severely hurt generalization in the presence of mislabeled examples. However, mislabeled examples are to hard avoid in extremely large datasets. We address this problem using the implicit regularization effect of stochastic gradient descent with large learning rates. By the loss statistics, we are able to identify mislabeled examples with remarkable success. Then, we discard the mislabeled examples on the fly and continue training with the rest of the examples. This leads to ON-THE-FLY DATA DENOISING (ODD), a simple yet effective algorithm that is robust to mislabeled examples, while introducing almost zero computational overhead. Empirical results demonstrate the effectiveness of ODD on several datasets containing artificial and real-world mislabeled examples.

authors:
Devansh Arpit
Stanisław Jastrzebski
Nicolas Ballas
David Krueger
Emmanuel Bengio
Maxinder S Kanwal
Tegan Maharaj
Asja Fischer
Aaron Courville
Yoshua Bengio
Simon Lacoste-Julien

title: A closer look at memorization in deep networks

venue: arXiv preprint arXiv:1706.05394,

authors:
Carla E Brodley
Mark A Friedl
Others

title: Identifying and eliminating mislabeled training instances

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

authors:
Pratik Chaudhari
Anna Choromanska
Stefano Soatto
Yann LeCun
Carlo Baldassi
Christian Borgs
Jennifer Chayes
Levent Sagun
Riccardo Zecchina

title: Entropy-SGD: Biasing gradient descent into wide valleys

venue: arXiv preprint arXiv:1611.01838,

authors:
L Dinh
R Pascanu
S Bengio
Y Bengio

title: Sharp minima can generalize for deep nets

venue: arXiv preprint arXiv:1703.04933,

authors:
Simon S Du
Wei Hu
Jason D Lee

title: Algorithmic regularization in learning deep homogeneous models: Layers are automatically balanced

venue: arXiv preprint arXiv:1806.00900,

authors:
Aritra Ghosh
Himanshu Kumar
P S Sastry

title: Robust loss functions under label noise for deep neural networks

venue: In AAAI,

authors:
Priya Goyal
Piotr Dollár
Ross Girshick
Pieter Noordhuis
Lukasz Wesolowski
Aapo Kyrola
Andrew Tulloch
Yangqing Jia
Kaiming He

title: Accurate, large minibatch SGD: Training ImageNet in 1 hour

venue: arXiv preprint arXiv:1706.02677,

venue: arXiv preprint arXiv:1512.03385,

title: Delving deep into rectifiers: Surpassing Human-Level performance on ImageNet classification

venue: arXiv preprint arXiv:1502.01852,

authors:
Dan Hendrycks
Mantas Mazeika
Duncan Wilson
Kevin Gimpel

title: Using trusted data to train deep networks on labels corrupted by severe noise

venue: arXiv preprint arXiv:1802.05300,

authors:
Sepp Hochreiter
Jürgen Schmidhuber

title: SIMPLIFYING NEURAL NETS BY DISCOVERING FLAT MINIMA

venue: Advances in Neural Information Processing Systems

authors:
Elad Hoffer
Itay Hubara
Daniel Soudry

title: Train longer, generalize better: closing the generalization gap in large batch training of neural networks

authors:
Gao Huang
Yixuan Li
Geoff Pleiss
Zhuang Liu
John E Hopcroft
Kilian Q Weinberger

title: Snapshot ensembles: Train 1, get M for free

venue: arXiv preprint arXiv:1704.00109,

authors:
Lu Jiang
Zhengyuan Zhou
Thomas Leung
Li-Jia Li
Li Fei-Fei

title: MentorNet: Learning Data-Driven curriculum for very deep neural networks on corrupted labels

venue: arXiv preprint arXiv:1712.05055,

authors:
Nitish Shirish Keskar
Dheevatsa Mudigere
Jorge Nocedal
Mikhail Smelyanskiy
Ping Tak Peter Tang

title: On Large-Batch training for deep learning: Generalization gap and sharp minima

venue: arXiv preprint arXiv:1609.04836,

authors:
Robert Kleinberg
Yuanzhi Li
Yang Yuan

title: An alternative view: When does SGD escape local minima

venue: arXiv preprint arXiv:1802.06175,

authors:
Ranjay A Krishna
Kenji Hata
Stephanie Chen
Joshua Kravitz
David A Shamma
Li Fei-Fei
Michael S Bernstein

title: Embracing error to enable rapid crowdsourcing

venue: In Proceedings of the 2016 CHI Conference on Human Factors in Computing Systems,

authors:
Wen Li
Limin Wang
Wei Li
Eirikur Agustsson
Luc Van Gool

title: WebVision database: Visual learning and understanding from web data

venue: arXiv preprint arXiv:1708.02862,

authors:
Yuanzhi Li
Tengyu Ma
Hongyang Zhang

title: Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations

venue: arXiv preprint arXiv:1712.09203,

authors:
Tongliang Liu
Dacheng Tao

title: Classification with noisy labels by importance reweighting

venue: arXiv preprint arXiv:1411.7718,

authors:
Ilya Loshchilov
Frank Hutter

title: SGDR: Stochastic gradient descent with warm restarts

venue: arXiv preprint arXiv:1608.03983,

authors:
Dhruv Mahajan
Ross Girshick
Vignesh Ramanathan
Kaiming He
Manohar Paluri
Yixuan Li
Ashwin Bharambe
Laurens van der Maaten

title: Exploring the limits of weakly supervised pretraining

venue: arXiv preprint arXiv:1805.00932,

authors:
Stephan Mandt
Matthew D Hoffman
David M Blei

title: Stochastic gradient descent as approximate bayesian inference

venue: arXiv preprint arXiv:1704.04289,

title: Implicit regularization in deep learning

venue: arXiv preprint arXiv:1709.01953,

authors:
Curtis G Northcutt
Tailin Wu
Isaac L Chuang

title: Learning with confident examples: Rank pruning for robust classification with noisy labels

venue: arXiv preprint arXiv:1705.01936,

authors:
Scott Reed
Honglak Lee
Dragomir Anguelov
Christian Szegedy
Dumitru Erhan
Andrew Rabinovich

title: Training deep neural networks on noisy labels with bootstrapping

venue: arXiv preprint arXiv:1412.6596,

authors:
Mengye Ren
Wenyuan Zeng
Bin Yang
Raquel Urtasun

title: Learning to reweight examples for robust deep learning

venue: arXiv preprint arXiv:1803.09050,

authors:
David Rolnick
Andreas Veit
Serge Belongie
Nir Shavit

title: Deep learning is robust to massive label noise

venue: arXiv preprint arXiv:1705.10694,

authors:
Olga Russakovsky
Jia Deng
Hao Su
Jonathan Krause
Sanjeev Satheesh
Sean Ma
Zhiheng Huang
Andrej Karpathy
Aditya Khosla
Michael Bernstein
Alexander C Berg
Li Fei-Fei

title: ImageNet large scale visual recognition challenge

venue: International journal of computer vision,

authors:
Nitish Srivastava
Geoffrey Hinton
Alex Krizhevsky
Ilya Sutskever
Ruslan Salakhutdinov

title: Dropout: a simple way to prevent neural networks from overfitting

venue: Journal of machine learning research: JMLR,

authors:
Sainbayar Sukhbaatar
Joan Bruna
Manohar Paluri
Lubomir Bourdev
Rob Fergus

title: Training convolutional networks with noisy labels

venue: arXiv preprint arXiv:1406.2080,

authors:
Chen Sun
Abhinav Shrivastava
Saurabh Singh
Abhinav Gupta

title: Revisiting unreasonable effectiveness of data in deep learning era

venue: arXiv preprint arXiv:1707.02968,

authors:
Christian Szegedy
Sergey Ioffe
Vincent Vanhoucke
Alex Alemi

title: Inception-v4, InceptionResNet and the impact of residual connections on learning

venue: arXiv preprint arXiv:1602.07261,

authors:
George E Uhlenbeck
Leonard S Ornstein

title: On the theory of the brownian motion

venue: Physical review,

authors:
Andreas Veit
Neil Alldrin
Gal Chechik
Ivan Krasin
Abhinav Gupta
Serge J Belongie

title: Learning from noisy Large-Scale datasets with minimal supervision

venue: In CVPR,

authors:
Tong Xiao
Tian Xia
Yi Yang
Chang Huang
Xiaogang Wang

title: Learning from massive noisy labeled data for image classification

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Sergey Zagoruyko
Nikos Komodakis

title: Wide residual networks

venue: arXiv preprint arXiv:1605.07146,

authors:
Chiyuan Zhang
Samy Bengio
Moritz Hardt
Benjamin Recht
Oriol Vinyals

title: Understanding deep learning requires rethinking generalization

venue: arXiv preprint arXiv:1611.03530,

authors:
Hongyi Zhang
Moustapha Cisse
Yann N Dauphin
David Lopez-Paz

authors:
Zhilu Zhang
Mert R Sabuncu

title: Generalized cross entropy loss for training deep neural networks with noisy labels

venue: arXiv preprint arXiv:1805.07836,

heading: 1 INTRODUCTION

text:
Over-parametrized deep neural networks have remarkable generalization properties while achieving near-zero training error (Zhang et al., 2016). However, the ability to fit the entire training set is highly undesirable, as a small portion of mislabeled examples in the dataset could severely hurt generalization (Zhang et al., 2016; Arpit et al., 2017). Meanwhile, an exponential growth in training data size is required to linearly improve generalization in vision tasks (Sun et al., 2017); this progress could be hindered if there are mislabeled examples within the dataset.
Mislabeled examples are to be expected in large datasets that contain millions of examples. Webbased supervision produces noisy labels (Li et al., 2017a; Mahajan et al., 2018); whereas human labeled datasets sacrifice accuracy for scalability (Krishna et al., 2016). Therefore, algorithms that are robust to various levels of mislabeled examples are warranted in order to further improve generalization for very large labeled datasets.
In this paper, we propose On-the-fly Data Denoising (ODD), a simple and robust method for training with noisy examples based on the implicit regularization effect of stochastic gradient descent. First, we train residual networks with large learning rate schedules and use the resulting losses to separate clean examples from mislabeled ones. This is done by identifying examples whose losses exceed a certain threshold. Reasonable thresholds can be derived from the loss distribution for uniform label noise which does not depend on the amount of mislabeled examples in the dataset. Finally, we remove these examples from the dataset and continue training until convergence.
Empirically, ODD performs significantly better than previous methods in datasets containing artificial noise (Sections 4.1 and 4.2) or real-world mislabeled examples (Section 4.3), while achieving equal or better accuracy than the state-of-the-art on clean datasets (Sections 4.1 and 4.2). We further conduct ablation studies to demonstrate that ODD is robust w.r.t hyperparameters and artificial
noise levels (Section 4.4). Our method is also able to detect mislabeled examples in the CIFAR-100 dataset without any additional supervision (Figure 1).

heading: 2 DENOISING DATASETS ON-THE-FLY

text:
The goal of supervised learning is to find a function f ∈ F that describes the probability of a random label vector Y ∈ Y given a random input vector X ∈ X , which has underlying joint distribution P (X,Y ). Given a loss function `(y, ŷ), one could minimize the average of ` over P :
R(f) = ∫ `(y, f(x)) dP (x,y)
The joint distribution P (X,Y ) is usually unknown, but we could gain access to its samples via a potentially noisy labeling process, such as crowdsourcing (Krishna et al., 2016) or web queries (Li et al., 2017a). We denote the training dataset with N examples as D = (xi,yi)i∈[N ] = G ∪ B. G represents correctly labeled (clean) examples sampled from P (X,Y ). B represents mislabeled examples that are not sampled from P (X,Y ), but from another distribution Q(X,Y ); G ∩ B = ∅. We aim to learn the function f from D without knowledge about B, G or their statistics (e.g. |B|). A typical approach is to assume that B = ∅, i.e. all examples are i.i.d. from P (X,Y ), and minimizing the following objective:
R̂(f) = − 1 N N∑ i=1 `(y, f(x))
If B = ∅ is indeed true, then R̂(f) → R(f) as N → ∞. However, this is not true if B 6= ∅ since D is no longer an unbiased population of P (X,Y ). Moreover, when F is the space of large neural networks with parameters exceeding N , f could fit the entire training dataset (Zhang et al., 2016), including the mislabeled examples. This results in undesired behavior of f on inputs of the mislabeled set, let alone outside the training data.
To illustrate the harm of mislabeled examples to generalization, we consider training on CIFAR-10 where some examples are mislabeled uniformly at random. Compared with training on D, training only on G could decrease validation error from 11.53 to 4.25 if there are 20% mislabeled examples, and from 15.57 to 5.06 if there are 40% mislabeled examples1. Therefore, if we are able to identify examples that belong to G, we could vastly improve generalization on P (X,Y ).

heading: 2.1 SEPARATING MISLABELED EXAMPLES VIA IMPLICIT REGULARIZATION

text:
Fortunately, in the case of classification with deep residual networks (He et al., 2015a), the implicit generalization of stochastic gradient descent (SGD) with large learning rates (e.g. 0.1) can separate examples from G and examples from B via the loss statistics. We demonstrate this in Figure 2, where we train deep residual networks on CIFAR-100 and ImageNet with different percentages of uniform label noise. In early stages of training, the loss distributions of clean examples and mislabeled ones have notable statistical distance. The network starts to fit mislabeled examples when learning rate starts to decrease, which is also crucial for achieving better generalization on clean datasets.
The working of the implicit regularization of gradient descent is by and large an open question that attracts much recent attentions (Neyshabur, 2017; Li et al., 2017b; Du et al., 2018). Empirically, it has been observed that large learning rates are beneficial for generalization (Kleinberg et al., 2018). Recent work has shown that the stationary distribution of SGD iterates corresponds to an Ornstein-Uhlenbeck process (Uhlenbeck & Ornstein, 1930) with noise proportional to the learning rate (Mandt et al., 2017). Training with large learning rates would then encourage solutions that are more robust to large random perturbations in the parameter space and less likely to overfit to mislabeled examples. Therefore, given these empirical and theoretical evidence on large learning rate helps generalization, we propose to classify correct and mislabeled examples through the loss statistics, and achieve better generalization by removing the examples that are potentially mislabeled and training on clean examples only.
1Validation error is 3.73 for 0% mislabeled case. Additional details for CIFAR-100 in Table 1.

heading: 2.2 THRESHOLDS THAT CLASSIFY MISLABELED EXAMPLES

text:
To improve generalization in practice, one critical problem is to select a reasonable threshold for classification. High thresholds could include too many examples from B, whereas low thresholds could prune too many examples from G; reasonable thresholds should also adapt to different unknown ratios of mislabeled examples. Let us first consider the case where Q(Y |X) has the highest entropy, which is the uniform distribution over labels. From Figure 2, the loss distribution for B is relatively stable with different ratios of |B|/|D|; examples in B are making little progress when learning rate is large. We propose to characterize the (negative log-likelihood) loss distribution of uniform label noise pn(l) via the following generative procedure:
l = −ỹk + log  ∑ i∈[N ] exp(ỹi)  ≡ − log softmax(ỹ)[k] (1) ỹ = fc(relu(x̃)), x̃ ∼ N (0, I), k ∼ Uniform{0, . . . ,K}
where fc(·) is the final (fully connected) layer of the network, relu(x̃) = max(x̃,0) is the Rectified Linear Unit, and k represents a random label from K classes. This represents the case where the model’s prediction is uncorrelated with the labels. The actual noise distribution could skew to the left if the model overfits to the noise, and skew to the right if the model predicts a label different from the noisy one. We find that an identity covariance matrix for x̃ is able to explain the noise distribution; this could result from well-conditioned objectives defined via deep residual networks (He et al., 2015a) and careful initialization (He et al., 2015b). We qualitatively demonstrate the validity of our characterization on CIFAR-100 and ImageNet datasets in Figure 2.
Therefore, we could define a threshold via the p-th percentile of pn(l); it relates to approximately how much examples in B we would retain if Q(Y |X) is uniform. In Section 4.4, we show that this method is able to identify different percentages of uniform label noise with high precision.

heading: 2.3 A PRACTICAL ALGORITHM FOR ROBUST TRAINING

text:
We can utilize this implicit regularization effect to remove examples that might harm generalization, leading to On-the-fly Data Denoising (ODD), a simple algorithm robust to mislabeled examples:
1. Train all examples with large learning rates for E epochs.
2. Compute the p-th percentile of the distribution in Eq. (1), denoted as Tp.
3. Remove examples whose average loss of the past h epochs exceeds T from the dataset.
4. Continue training the remaining examples from epoch E + 1.
ODD introduces three hyperparameters: E determines the amount of training that separates clean examples from noisy ones; p determines Tp that specifies the trade-off between less noisy examples and more clean examples; h determines the window of averaged loss statistics to reduce variance from data augmentation. We do not explicitly estimate the portion of noise in the dataset, nor do we assume any specific noise model. In fact, the threshold Tp could be used to accurately predict the portion of uniform noise in the dataset, and works quite well even on other types of label noise; we will demonstrate this in Section 4. Moreover, ODD is compatible with existing practices for learning rate schedules, such as stepwise (He et al., 2015a) or cosine (Loshchilov & Hutter, 2016).

heading: 3 RELATED WORK

text:
Implicit Regularization of SGD The generalization of neural networks trained with SGD depend heavily on learning rate schedules (Loshchilov & Hutter, 2016). It has been proposed that wide local minima2 could result in better generalization (Hochreiter & Schmidhuber, 1995; Chaudhari et al., 2016; Keskar et al., 2016). Several factors could contribute to wider local optima and better generalization, such as smaller minibatch sizes (Keskar et al., 2016), reasonable learning rates (Kleinberg et al., 2018), and longer training time (Hoffer et al., 2017). Moreover, solutions that are further away from the initialization may lead to wider local minima and better generalization (Hoffer et al., 2017). In the presence of mislabeled examples, changes in optimization landscape (Arpit et al., 2017) could result in bad local minima (Zhang et al., 2016), although it is argued that larger batch sizes could mitigate this effect (Rolnick et al., 2017).
Training with Mislabeled Examples One paradigm to robust training with noisy labels involves estimating the noise distribution (Liu & Tao, 2014) or confusion matrix (Sukhbaatar et al., 2014). Another line of methods propose to identify and clean the noisy examples through predictions of auxillary networks (Veit et al., 2017) or via binary predictions (Northcutt et al., 2017); the noisy labels are either pruned (Brodley et al., 1996) or replaced with model predictions (Reed et al., 2014). Our method is comparable to these approaches, but the key difference is that we leverage the implicit regularization effect of SGD to identify noisy examples. Other approaches propose to reweigh the examples via a pretrained network (Jiang et al., 2017), meta learning (Ren et al., 2018), or surrogate loss functions (Ghosh et al., 2017; Zhang & Sabuncu, 2018). Some methods require a set of trusted examples (Xiao et al., 2015; Hendrycks et al., 2018).
ODD has several appealing properties compared to existing methods. First, the thresholds for classifying mislabeled examples from ODD do not rely on estimations of the noise confusion matrix. Next, ODD does not require additional trusted examples. Finally, ODD removes potentially noisy examples on-the-fly; it has little computational overhead compared to standard SGD training.

heading: 4 EXPERIMENTS

text: We evaluate our method on clean and noisy versions of CIFAR-10, CIFAR-100, ImageNet (Russakovsky et al., 2015) and WebVision (Li et al., 2017a) datasets. We use stochastic gradient descent with momentum for training while following standard image preprocessing and data augmentation practices. We do not consider dropout (Srivastava et al., 2014) or model ensembles (Huang et al., 2017) in our experiments. We use h = 2 for all our ODD experiments; we observe that having h ∈ [2, 5] yields similar results.

heading: 4.1 CIFAR-10 AND CIFAR-100

text:
We first evaluate our method on the CIFAR-10 and CIFAR-100 datasets, which contain 50,000 training images and 10,000 validation images of size 32×32 with 10 and 100 labels respectively. During training, we follow the data augmentations in (Zagoruyko & Komodakis, 2016), which performs horizontal flips, takes random crops from 40× 40 images padded by 4 pixels on each side, and fills missing pixels with reflections of the original images.
In our experiments, we train the wide residual network architecture (WRN-28-10) in (Zagoruyko & Komodakis, 2016) for 200 epochs with a minibatch size of 128, momentum 0.9 and weight decay
2Although the notion of wideness is controversial (Dinh et al., 2017).
5×10−4. We consider a cosine annealing schedule as described in (Loshchilov & Hutter, 2016) with T0 = 200, Tmult = 1, ηmax = 0.1, ηmin = 10
−5 (no warm restarts), as we observe this schedule outperforms the traditional stepwise schedules on the clean dataset. We include results for two types of stepwise schedules in Appendix A.1.

heading: 4.1.1 GENERALIZATION ON INPUT-AGNOSTIC LABEL NOISE

text:
We first consider label noise that are agnostic to inputs. Following Zhang et al. (2016), We randomly replace a 0%/20%/40%) of the training labels to uniformly random ones, and evaluate generalization error on the clean validation set. We compare with the following baselines: ORACLE, where the model knows the true identity of clean examples and only trains on them; Empirical Risk Minimization (ERM, Goyal et al. (2017)) which assumes all examples are clean; MENTORNET (Jiang et al., 2017), which pretrains an auxiliary model that predicts weights for each example based on its input features; REN (Ren et al., 2018), which optimizes the weight of examples via meta-learning; mixup (Zhang et al., 2017), a data augmentation approach that trains neural networks on convex combinations of pairs of examples and their labels; and Generalized Cross Entropy (GCE, Zhang & Sabuncu (2018)) that includes cross-entropy loss and mean absolute error (Ghosh et al., 2017).
We report the top-1 validation error in Table 1, where ? denotes methods trained with knowledge of 1000 additional clean labels. Notably, ODD significantly outperforms all other algorithms (except for the oracle) when there is artificial noise, and is on-par with ERM even when there is no artificial noise. On the one hand, this suggests that ODD is able to distinguish the mislabeled examples and improve generalization; on the other hand, it would seem that removing certain examples even in the “clean” dataset does not seem to hinder generalization.
ODD prevents overfitting to noise We compare the learning curves of ERM and ODD in Figure 3 with a stepwise schedule under 40% label corruption. ERM easily overfits the random labels when learning rate decreases, whereas ODD manages to continue improving generalization.
Mislabeled examples in CIFAR-100 We run our methods on three random seeds, and find the examples that are considered mislabeled by all the three instances (598 in total); we demonstrate some examples3 in Figure 1, which contains ambiguous / wrong labels.
3We include more examples in Appendix A.4.

heading: 4.1.2 GENERALIZATION ON INPUT-DEPENDENT LABEL NOISE

text: Images from some classes could be harder to label correctly than that from other classes. To simulate this, we perform experiments on settings where the label noise only comes from certain types of input data. Specifically, we remove a portion of classes from the dataset (e.g. class 9 in CIFAR-10), and assign the labels of all its examples to the remaining classes randomly (e.g. a class 9 example has a class 0 – 8 random label). This reduces the total number of classes, so on the validation set we only consider the classes that are not removed (e.g. classes 0 – 8). We compare ERM and ODD on datasets with 10% or 20% of the examples mislabeled, and summarize the results in Table 2. ODD is still able to significantly outperform ERM under such input-dependent noise.

heading: 4.1.3 GENERALIZATION ON NON-HOMOGENEOUS LABELS

text: We evaluate ERM and ODD on a setting without mislabeled examples, but the ratio of classes could vary. To prevent the model from utilizing the number of examples in a class, we combine multiple classes of CIFAR-100 into a single class, creating the CIFAR-20 and CIFAR-50 tasks. In CIFAR-50, we combine an even class with an odd class while we remove c% of the examples in the odd class. In CIFAR-20, we combine 5 classes in CIFAR-100 that belong to the same super-class4 while we remove c% of the examples in 4 out of 5 classes. This is performed for both training and validation datasets. Results for ERM and ODD with p = 10 and E = 75 are shown in Table 3, where ODD is able to outperform ERM in these settings where the input examples are not uniformly distributed.

heading: 4.2 IMAGENET

text:
We conduct additional experiments on the ImageNet-2012 classification dataset (Russakovsky et al., 2015). The dataset contains 1.28 million training images and 50,000 validation images from 1,000 classes. Input-agnostic random noise of 0%, 20%, 40% are considered. We follow standard data augmentation practices during training, including scale and aspect ratio distortions, random crops, and horizontal flips. We only use the center 224× 224 crop for validation. We train ResNet-50 and ResNet-152 models (He et al., 2015a) with the cosine schedule with initial learning rate 0.1, momentum 0.9, weight decay 10−4, 90 training epochs, and report top-1 and top-5 validation errors in Table 4. ODD significantly outperforms ERM in terms of both top-1 and top-5 errors on the 20% and 40% mislabeled examples, while being competitive with the clean dataset.
4Super-class information is available at https://www.cs.toronto.edu/˜kriz/cifar.html.

heading: 4.3 WEBVISION – A REAL-WORLD NOISY DATASET

text:
We further verify the effectiveness of our method on a real-world noisy dataset. The WebVision2017 dataset (Li et al., 2017a) contains 2.4 million of real-world noisy labels, that are crawled from Google and Flickr using the 1,000 labels from the ImageNet-2012 dataset. We train two architectures, ResNet-50 and Inception ResNet-v2 (Szegedy et al., 2016) with the same procedure in the ImageNet experiments, except for Inception ResNet-v2 we train for 50 epochs and use input images of size 299×299. We use both WebVision and ImageNet validation sets for 1-crop validation, following the settings in Jiang et al. (2017). We do not use a pretrained model or additional labeled data from ImageNet during training.
Our ODD method with p = 30 removes in the training set around 9.0% of the total examples with ResNet-50 and 9.3% of the total examples with Inception ResNet-v2 (Szegedy et al., 2016). Table 5 suggests that our method is able to outperform both ERM and MENTORNET when the training dataset is noisy, even as we remove a notable portion of examples. We include more results in Appendix A.3. In comparison, we removed around 1.1% of examples in ImageNet (Table 10, Appendix A.2); this may suggest that WebVision labels are indeed much noisier.

heading: 4.4 ABLATION STUDIES

text:
Sensitivity to p We first evaluate noisy ImageNet classification with ResNet-50 where p ∈ {1, 10, 30, 50, 80} and E = 60 in Table 6. A higher p includes more clean examples at the cost of involving more noisy examples. In the 20% and 40% noisy cases, the optimal trade-off for generalization is at p = 10, yet even when p = 50, the validation errors are still significantly better than ERM. When there is no artificial noise, generalization of ODD starts to match that of ERM as p ≥ 10. Therefore, ODD is not very sensitive to p in these cases, and empirically p = 10 represents the best trade-off. We include results for ResNet-152 in Appendix A.2.
Sensitivity to E We evaluate the validation error of ODD on CIFAR with 20% and 40% inputagnoistic label noise where E ∈ {25, 50, 75, 100, 150, 200} (E = 200 is equivalent to ERM). The results in Figure 4 suggest that our method is able to separate noisy and clean examples if E is
relatively small where the learning rate is high, but is unable to perform well when the learning rate decreases at later stages of the training.
Sensitivity to the amount of noise Finally, we evaluate the training error of ODD on CIFAR under input-agnostic label noise of {1%, 5%, 10%, 20%, 30%, 40%} with p = 5, E = 50 or 75. This reflects how much examples exceed the threshold and are identified as noise at epoch E. From Figure 5, we observe that the training error is almost exactly the amount of noise in the dataset, which demonstrates that the loss distribution of noise can be characterized by our threshold regardless of the percentage of noise in the dataset.

heading: 5 DISCUSSION

text:
We have proposed ODD, a straightforward method for robust training with mislabeled examples. ODD utilizes the implicit regularization effect of stochastic gradient descent to prune examples that potentially harm generalization. Empirical results demonstrate that ODD is able to significantly outperform related methods on a wide range of datasets with artificial and real-world mislabeled examples, maintain competitiveness with ERM on clean datasets, as well as detecting mislabeled examples automatically in CIFAR-100.
The implicit regularization of stochastic gradient descent opens up other research directions for implementing robust algorithms. For example, we could consider using a smaller network to remove examples, removing examples not only once but multiple times, retraining from scratch with the denoised dataset, or other data-augmentation approaches such as mixup (Zhang et al., 2017). Moreover, it would be interesting to understand the implicit regularization over mislabeled examples from a theoretical viewpoint.

heading: A ADDITIONAL EXPERIMENTAL RESULTS

text:
A.1 CIFAR INPUT-AGNOSTIC NOISE
In addition to the existing experiments, we include results for ORACLE, ERM, ODD with two stepwise annealing schedules. In the stepwise schedules, the learning rate starts from 0.1 and is then divided by 5 after 60, 120, 160 epochs (stepwise-i, which is used in Zagoruyko & Komodakis (2016)) or after 100, 150, 175 epochs (stepwise-ii). We consider cosine schedule in Section 4 because it achieves better generalization performance on the clean dataset.
For the stepwise schedules, we set E to be the epoch at which learning rate begins to decay. For the cosine schedule, we set E = 100 for CIFAR-10 and E = 50 for CIFAR-100. We set p = 20 and p = 10 for CIFAR-10 and CIFAR-100 respectively for the clean datasets; p = 10 and p = 5 for noisy datasets. This is motivated by the fact that CIFAR-10 has less labels, so the threshold has to take into account random labels that happens to be correct.
Tables 7 and 8 contain summary of the results. The cosine learning rate schedule generally outperforms the stepwise schedules. We note that for the stepwise schedules in ERM, the optimal validation error is achieved when the learning rate just starts to decay to 0.2, after that the model starts to overfit to noise, as demonstrated in Figure 3.
We evaluate precision and recall for examples classified as noise on CIFAR10 and CIFAR100 for different noise levels (1, 5, 10, 20, 30, 40) in Figure 6. The recall values are around 0.84 to 0.88 where as the precision values range from 0.88 to 0.92. This demonstrates that ODD is able to achieve good precision/recall with default hyperparameters even at different noise levels.
A.2 IMAGENET ABLATION RESULTS ON RESNET-152
We include the ImageNet ablation experiments on the hyperparameter p on the ResNet-152 architecture in Table 9. Compared to the ResNet-50 experiments, we can draw similar conclusions here: p = 10 generally represents the best trade-off. We show the percentage of examples discarded by NOISE CLASSIFIER in Table 10; the percentage of discarded examples by p = 10 is very close to the actual noise level. Moreover, the percentage of discarded examples does not vary significantly when we change our architecture from ResNet-50 to ResNet-152.
A.3 WEBVISION ABLATION RESULTS
We include the WebVision ablation experiments on the hyperparameter p on the ResNet-50 and Inception ResNet-v2 architecture with E = 60 in Tables 11 and 12, where we report top-1 and top-5 validation errors on WebVision and ImageNet validation sets respectively, as well as how many examples are discarded by our method at epoch 60. Similar to the results in ImageNet, generalization performance is generally insensitive to the hyperparameter p, except for p = 1, which discarded 25.3% of the examples. We use 2 seeds for each experiment setting. Notice that at each p, WebVision has more examples discarded compared to ImageNet (with 0% artificial noise), which further suggests that it has more mislabeled examples than ImageNet. Again, the percentage of discarded examples does not vary significantly across different architectures.
A.4 IMAGES IN CIFAR-100 CLASSIFIED AS NOISE
We display the examples in CIFAR-100 training set for which our ODD methods identify as noise across 3 random seeds. One of the most common label such examples have is “leopard”; in fact, 21 of 50 “leopard” examples in the training set are perceived as hard, and we show some of them in Figure 7. It turns out that a lot of the “leopard” examples contains images that clearly contains tigers and black panthers (CIFAR-100 has a label corresponding to “tiger”).
We also demonstrate random examples from the CIFAR-100 that are identified as noise in Figure 8 and those that are not identified as noise in Figure 9. The examples identified as noise often contains multiple objects, and those not identified as noise often contains only one object that is less ambiguous in terms of identity.


summaries: [UPDATE: Thanks to the authors for engaging with the comments in the review.  I hope to see them fully addressed in the revision and look forward to seeing the next version of the paper!]  The paper investigates how pretrained representations might be expected to help prediction on a downstream task.  In some sense, they should "obviously" help: the downstream predictor can learn to exploit any mutual information between its target output and the pretrained representations.  That's the standard motivation for multitask learning.  However, the paper exhibits specific settings in which the target output can be *exactly* predicted from the pretrained representations, using only a *simple* probe such as a linear classifier.    The paper is primarily theoretical, although a few experiments are included by way of illustration.  It is not clear what the experiments are intended to verify.  To make the analysis possible, the pretrained representations in the paper are not underlying representations such as word embeddings.  Rather, they are identifiable given the pretraining data, which makes it clear what "successful pretraining" means.  For example, the representation of word x_i in a sentence is the conditional distribution of x_i given all other words in the sentence, as predicted by a pretrained cloze language model.  The analysis assumes (though the experiments do not check) that the predicted conditional distribution actually matches the true conditional distribution.  This should be the case in the limit of infinite training data, provided that the cloze language model has adequate capacity and can be trained to a global optimum of log-loss.  It is shown that if the vocabulary is large enough, or under other assumptions, then this representation of x_1 is enough to recover the posterior distributions over h_1 and h_0, which are assumed to be sufficient for the downstream task.  Caveat: In the analysis, the observed sentences are assumed to be generated by a latent-variable model.  One would then expect the downstream task target output to be predictable from the actual latent variables.  However, the paper assumes more strongly that the target is a function of the *posterior distribution* of the latent variables given the observed sentences.  In other words, the goal of the downstream task is not to recover what the speaker *actually* meant, but only to recover what a hearer who knows the true generating distribution *thinks* it meant (given only the sentence and no other context).  Two structures are considered for the latent-variable model and its relationship to the downstream task.  For the more complex version, the predictor is also more complex, requiring argmax attention over all positions i.  In both cases, the paper considers allowing the predictor to prepend an arbitrary word embedding x_1 to the input.  This soft prompt can serve to shift the posterior predictions of x_i so as to be more informative for the downstream task.

summaries: This work puts forward a theoretical analysis of pre-trained language models and head and prompt tuning on downstream tasks. In particular it studies the assumptions on the distribution of language and downstream tasks under which the downstream label can be recovered from a pretrained LM by using head or prompt tuning, by considering a generative model of language that can be represented as a HMM. The key components of the analysis include: 1. The data distribution is generated from a HMM (or a memory augmented HMM). 2. The initial state of the HMM, $H_0$, contains all the information necessary to recover the downstream label for the sequence.  Under this setup, for head tuning, the downstream labels are recoverable under the condition that the token emission probability matrix $W$ has linearly independent columns (implying that the dimensionality of the HMM hidden states is smaller than the vocabulary size of the data generating distribution). On the other hand, prompt tuning requires only a relaxed version of this condition, only the columns corresponding to the support of the downstream task need to be independent, eliminating the condition on the hidden state dimensionality.  Under memory augmented HMMs, for downstream tasks that can be represented as a linear function of the memory (with a fixed number of cells, $N$), the non-degeneracy condition can be relaxed even for head tuning. For prompt tuning, the independence assumption can be further relaxed to hold only for the set of memory cells that the downstream task depends on.  The authors further perform simulation experiments to validate the findings on different settings of the data generating distribution. Simulation results agree with the findings: 1. Prompt tuning outperforms head tuning in the setting where hidden state size is significantly larger than the token vocabulary size. 2. When the data is generated using memory-augmented HMMs, and the downstream task information can be recovered entirely from the memory, even head tuning gets close to perfect performance for all hidden state sizes.

summaries: The paper studies the question why head and prompt tuning work for pre-trained LM. Analyzing complicated neural networks is a daunting task. The paper offers an analysis framework based on the Hidden Markov Model (HMM) and memory augmented HMM. Under this framework, the paper makes a connection with masked LM (BERT) that at time step i, the output of BERT(`[mask]_i`) represents the marginal P(X_i) of the HMM. The framework further allows for analysis of a binary classification task using head tuning or prompt tuning via the posterior distribution of a HMM. To this end, the paper runs a simulation study with data generated from HMM and memory augmented HMM.

summaries: This paper proposes an analysis framework to link the pre-training and downstream tasks. In this framework, the downstream task needs to predict the properties of the posterior distribution over latent variables in an underlying generative model. And When the generative model is a standard HMM, the downstream recovery is possible with a simple classification head under strong non-degeneracy assumptions. It further shows that changing the generative model to a memory-augmented HMM or using prompt tuning can relax the non-degeneracy conditions.

abstractText: Pretrained language models have achieved state-of-the-art performance when adapted to a downstream NLP task. However, theoretical analysis of these models is scarce and challenging since the pretraining and downstream tasks can be very different. We propose an analysis framework that links the pretraining and downstream tasks with an underlying latent variable generative model of text — the downstream classifier must recover a function of the posterior distribution over the latent variables. We analyze head tuning (learning a classifier on top of the frozen pretrained model) and prompt tuning in this setting. The generative model in our analysis is either a Hidden Markov Model (HMM) or an HMM augmented with a latent memory component, motivated by long-term dependencies in natural language. We show that 1) under certain non-degeneracy conditions on the HMM, simple classification heads can solve the downstream task, 2) prompt tuning obtains downstream guarantees with weaker non-degeneracy conditions, and 3) our recovery guarantees for the memory-augmented HMM are stronger than for the vanilla HMM because task-relevant information is easier to recover from the long-term memory. Experiments on synthetically generated data from HMMs back our theoretical findings.

authors:
Sanjeev Arora
Hrishikesh Khandeparkar
Mikhail Khodak
Orestis Plevrakis
Nikunj

title: Saunshi. A theoretical analysis of contrastive unsupervised representation learning

venue: In International Conference on Machine Learning,

authors:
Xiang Chen
Xin Xie
Ningyu Zhang
Jiahuan Yan
Shumin Deng
Chuanqi Tan
Fei Huang
Luo Si
Huajun Chen

title: Adaprompt: Adaptive prompt-based finetuning for relation extraction

venue: arXiv preprint arXiv:2104.07650,

authors:
Justin T Chiu
Alexander M Rush

title: Scaling hidden markov language models

venue: arXiv preprint arXiv:2011.04640,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova. Bert

title: Pre-training of deep bidirectional transformers for language understanding

venue: arXiv preprint arXiv:1810.04805,

title: How contextual are contextualized word representations? comparing the geometry of bert, elmo, and gpt-2 embeddings

venue: arXiv preprint arXiv:1909.00512,

authors:
Tianyu Gao
Adam Fisch
Danqi Chen

title: Making pre-trained language models better few-shot learners

venue: arXiv preprint arXiv:2012.15723,

authors:
Mario Giulianelli
Jack Harding
Florian Mohnert
Dieuwke Hupkes
Willem Zuidema

title: Under the hood: Using diagnostic classifiers to investigate and improve how language models track agreement information

venue: arXiv preprint arXiv:1808.08079,

authors:
Amit Gruber
Yair Weiss
Michal Rosen-Zvi

title: Hidden topic markov models

venue: In Artificial intelligence and statistics,

authors:
Karen Hambardzumyan
Hrant Khachatrian
Jonathan May

title: Warp: Word-level adversarial reprogramming

venue: arXiv preprint arXiv:2101.00121,

authors:
Jeff Z. HaoChen
Colin Wei
Adrien Gaidon
Tengyu Ma

title: Provable guarantees for selfsupervised deep learning with spectral contrastive loss, 2021

authors:
John Hewitt
Christopher D Manning

title: A structural probe for finding syntax in word representations. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)

authors:
Ganesh Jawahar
Benoît Sagot
Djamé Seddah

title: What does bert learn about the structure of language? In ACL 2019-57th Annual Meeting of the Association for Computational Linguistics, 2019

authors:
Zhengbao Jiang
Frank F Xu
Jun Araki
Graham Neubig

title: How can we know what language models know

venue: Transactions of the Association for Computational Linguistics,

authors:
Mandar Joshi
Danqi Chen
Yinhan Liu
Daniel S Weld
Luke Zettlemoyer
Omer Levy

title: Spanbert: Improving pre-training by representing and predicting spans

authors:
Taeuk Kim
Jihun Choi
Daniel Edmiston
Sang-goo Lee

title: Are pre-trained language models aware of phrases? simple but strong baselines for grammar induction

venue: arXiv preprint arXiv:2002.00737,

authors:
Daphne Koller
Nir Friedman

title: Probabilistic graphical models: principles and techniques

venue: MIT press,

title: Robust part-of-speech tagging using a hidden markov model

venue: Computer speech & language,

authors:
Jason D Lee
Qi Lei
Nikunj Saunshi
Jiacheng Zhuo

title: Predicting what you already know helps: Provable self-supervised learning

venue: arXiv preprint arXiv:2008.01064,

authors:
Brian Lester
Rami Al-Rfou
Noah Constant

title: The power of scale for parameter-efficient prompt tuning

venue: arXiv preprint arXiv:2104.08691,

authors:
Yoav Levine
Barak Lenz
Opher Lieber
Omri Abend
Kevin Leyton-Brown
Moshe Tennenholtz
Yoav Shoham

title: Pmi-masking: Principled masking of correlated spans

venue: arXiv preprint arXiv:2010.01825,

authors:
Xiang Lisa Li
Percy Liang

title: Prefix-tuning: Optimizing continuous prompts for generation

authors:
Matthew E Peters
Mark Neumann
Mohit Iyyer
Matt Gardner
Christopher Clark
Kenton Lee
Luke Zettlemoyer

title: Deep contextualized word representations

venue: arXiv preprint arXiv:1802.05365,

authors:
Guanghui Qin
Jason Eisner

title: Learning how to ask: Querying lms with mixtures of soft prompts

venue: arXiv preprint arXiv:2104.06599,

authors:
Lawrence Rabiner
Biinghwang Juang

title: An introduction to hidden markov models

venue: ieee assp magazine,

authors:
Alec Radford
Karthik Narasimhan
Tim Salimans
Ilya Sutskever

title: Improving language understanding by generative pre-training

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: arXiv preprint arXiv:1910.10683,

authors:
Anna Rogers
Olga Kovaleva
Anna Rumshisky

title: A primer in bertology: What we know about how bert works

authors:
Nikunj Saunshi
Sadhika Malladi
Sanjeev Arora

title: A mathematical exploration of why language models help solve downstream tasks

venue: arXiv preprint arXiv:2010.03648,

authors:
Timo Schick
Hinrich Schütze

title: Exploiting cloze questions for few shot text classification and natural language inference

venue: arXiv preprint arXiv:2001.07676,

title: It’s not just size that matters: Small language models are also few-shot learners

venue: arXiv preprint arXiv:2009.07118,

authors:
Taylor Shin
Yasaman Razeghi
Robert L Logan IV
Eric Wallace
Sameer Singh

title: Autoprompt: Eliciting knowledge from language models with automatically generated prompts

venue: arXiv preprint arXiv:2010.15980,

authors:
Koustuv Sinha
Robin Jia
Dieuwke Hupkes
Joelle Pineau
Adina Williams
Douwe Kiela

title: Masked language modeling and the distributional hypothesis: Order word matters pre-training for little

venue: arXiv preprint arXiv:2104.06644,

authors:
Sainbayar Sukhbaatar
Arthur Szlam
Jason Weston
Rob Fergus

title: End-to-end memory networks

venue: arXiv preprint arXiv:1503.08895,

authors:
Ian Tenney
Dipanjan Das
Ellie Pavlick

title: Bert rediscovers the classical nlp pipeline

venue: arXiv preprint arXiv:1905.05950,

authors:
Ian Tenney
Patrick Xia
Berlin Chen
Alex Wang
Adam Poliak
R Thomas McCoy
Najoung Kim
Benjamin Van Durme
Samuel R Bowman
Dipanjan Das

title: What do you learn from context? probing for sentence structure in contextualized word representations

venue: arXiv preprint arXiv:1905.06316,

authors:
Christopher Tosh
Akshay Krishnamurthy
Daniel Hsu

title: Contrastive estimation reveals topic posterior information to linear models

title: Contrastive learning, multi-view redundancy, and linear models

venue: In Algorithmic Learning Theory,

authors:
Colin Wei
Kendrick Shen
Yining Chen
Tengyu Ma

title: Theoretical analysis of self-training with deep networks on unlabeled data, 2020

authors:
Tianyi Zhang
Tatsunori Hashimoto

title: On the inductive bias of masked language modeling: From statistical to syntactic dependencies

venue: arXiv preprint arXiv:2104.05694,

authors:
Zexuan Zhong
Dan Friedman
Danqi Chen

title: Factual probing is [mask]: Learning vs. learning to recall

venue: arXiv preprint arXiv:2104.05240,

heading: 1 Introduction

text:
Natural language processing (NLP) has been revolutionized by large-scale pretrained language models such as BERT [4] and GPT [25], which are adapted to a variety of downstream NLP tasks. Although a large body of empirical work seeks to understand the effectiveness of pretrained models [7, 5, 12, 35, 34, 11, 27, 15], theoretical understanding is scarce. Theoretically analyzing the relationship between the pretraining and downstream tasks is challenging because pretraining and downstream settings can greatly differ.
The key starting point for our analysis is to link the pretraining and downstream settings through an underlying generative model of the data. We model the data distribution as a latent variable model and the downstream task as a function of the latent variables. Assuming that pretraining on a large corpus allows us to learn the generative model, the conditional token probabilities predicted by the pretrained model carry information about the hidden variables. In downstream adaptation, we aim to recover this information to solve the downstream task.
Though full finetuning is the de facto empirical standard, analyzing it is challenging because it requires characterizing the weights of the pretrained model. In this paper, we focus on head tuning
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
and prompt tuning, which both freeze all pretrained parameters and allow us to treat the pretrained model as a black box. Head tuning [22] trains task-specific heads on top of the pretrained model outputs. Prompt tuning [31, 19, 9, 21] optimizes a task-specific “prompt” that is concatenated to the model input. Studying prompt tuning is particularly interesting since it can match the performance of full finetuning with less computation time [19, 9, 21].
Our work contrasts with prior theoretical work [28], which assumes that downstream labels are recoverable via a linear head applied to the conditional token probabilities, and analyze how errors in pretraining or model misspecification propagate downstream. We consider specific generative distributions for which we can prove these assumptions, showing that head and prompt tuning can recover the downstream labels.
Our analysis considers two data-generating distributions with increasing realism. First, we consider data generated from a Hidden Markov Model (HMM), where the downstream task is to learn a linear classifier on the posterior distribution over the hidden states (Section 3). We prove that, under strong non-degeneracy conditions on token emission probabilities, a linear head applied to a pretrained model G which outputs exact conditional token probabilities (Gipxq “ P rXi |x´is) can recover the downstream label (Theorem 3.3). Furthermore, we can prove better recovery guarantees with relaxed non-degeneracy assumptions (Assumption 3.1) by using continuous prompt tuning (Theorem 3.6), reflecting the strong empirical performance of prompt tuning [19, 9, 21]. Intuitively, prompt tuning conditions the latent variables so that nonessential information for the downstream task can be ignored during the tuning phase, making task-essential information easier to recover.
Second, we also strengthen our analysis by leveraging additional structure in the data. Motivated by long-range dependences in natural language, we analyze HMM variants with additional latent “memory” variables that can store long-term information more easily than vanilla HMMs (Section 4). Here, the downstream task is to learn a linear classifier on the posterior distribution of the memory variables. We show that, under weaker non-degeneracy conditions than the first setting, an attentionbased classification head can recover ground-truth downstream labels from pretrained model outputs (Theorem 4.3). Intuitively, our recovery guarantees improve because the classification head can focus on the persistent, task-essential information in the memory while ignoring other transient and nonessential aspects of the latent variables. As with the vanilla HMM, we analyze prompt tuning for relaxing the non-degeneracy conditions even further (Theorem 4.6).
In summary, we relate the pretraining and downstream tasks by assuming that the downstream task is to learn a classifier on the posterior distributions of the latent variables defined by an underlying generative model of text. Our theoretical contributions are: 1) in this setting we analyze an HMM generative model show that simple classification heads can recover the true downstream labels under certain non-degeneracy assumptions, 2) we prove that soft prompt tuning can relax the non-degeneracy assumptions needed for downstream recovery making it easier to extract task-specific information, and 3) our recovery guarantees are stronger for memory-augmented HMMs in comparison to the vanilla HMM when tuning an attention-based classfication head.
We empirically evaluate our theoretical results with language models pretrained on synthetically generated data from HMMs. We find that prompt tuning obtains good downstream performance when our non-degeneracy conditions are relaxed, whereas head tuning performs poorly. Furthermore, we show that head tuning obtains better downstream performance when data is generated from a memory-augmented HMM, compared to a vanilla HMM, as is predicted by our theory.

heading: 1.1 Related works

text:
The black box nature of BERT and related models has inspired a variety of empirical works which seek to understand them. Probing papers study whether a pretrained model computes various types of structured information (e.g., syntactic [35, 11]) by evaluating the performance of simple classifiers, or probes, on the representations [7, 12, 34, 27, 15]. Other papers ablate various aspects of pretraining, such as changing the masking scheme [14, 20, 40] or permuting the word order [32].
In comparison, theoretical analysis of pretrained language models is limited. Besides [28], which we discussed in Section 1, Zhang and Hashimoto [40] analyze using a linear classifier to approximately recover the latent variable in a Gaussian graphical model with sparse dependencies between observed variables. However, their analysis and setting are focused towards understanding syntactic dependencies between tokens, whereas we directly model and analyze downstream performance.
Prompt-based tuning [31, 19, 9, 21, 13, 6, 41, 2, 23], which has improved empirical downstream performance for lightweight adaptation methods beyond head tuning to approach full finetuning, is an important focus of our theoretical analysis. Shin et al. [31] employ task-specific prompts that are optimized over the discrete token space. Schick and Schütze [29, 30] reformulate natural language tasks as cloze-style phrases to enable few-shot learning. Subsequent methods [19, 9, 21] optimize “soft” prompts, or continuous embedding vectors. Lester et al. [19] employ soft prompts on pretrained large-scale T5 [26] models and show that as the model size increases, prompt tuning performance can eventually match finetuning. Hambardzumyan et al. [9] applies a variant of soft prompt tuning to MLM models. Li and Liang [21] propose prefix tuning, which prepends a trainable prefix embedding sequence to all layers of the transformer.
More broadly, Lee et al. [18] analyze reconstruction-based self-supervised learning methods in a general setting and show that under certain conditional independence assumptions, predicting one observed variable from another allows recovery of the latent with a linear head. Other theoretical works analyzing self-supervised or constrastive learning include [1, 10, 36, 38, 37], but they do not directly relate to our particular setting.

heading: 2 Formulations and notations

text:
We analyze models pretrained on masked language modeling (MLM) objectives. Let X denote a finite vocabulary of input tokens, X ˚ the set of variable-length sequences of tokens, and X “ pX1, . . . , XT q P X ˚ a random sequence of T tokens. Let ∆|X | denote the space of probability distributions over tokens.
Pretraining and downstream task. Let Gpxq “ pG1pxq, G2pxq, . . .q denote the masked language model which predicts a probability vector for each timestep in the input x. Our theoretical abstraction is that Gi perfectly computes the distribution of Xi, the i-th token, conditioned on all other tokens: Gipxq “ P rXi|X´i “ x´is. Here P rXi |X´i “ x´is P ∆|X | is a probability vector. In particular, Gipxq does not depend on xi. The downstream task involves labeled examples px, F ‹pxqq P X ˚ˆY , where F ‹ : X ˚ Ñ Y provides ground-truth downstream labels and Y is a discrete set of labels for classification.
Head and prompt tuning. Head tuning trains a classification head f on top of fixed model outputs, resulting in the classifier F pxq “ 1pfpGpxqq ě 0q. We expect f to be a simple function such as a linear or one layer attention model. We also analyze variants where f also takes the tokens x or embeddings of x as input, which provides additional information. Soft prompt tuning requires viewing the pretrained model G as a function of the token embeddings; we refer to this model by G. Letting epxq “ epx1q, . . . , epxtq denote the token embeddings, we have Gpepxqq “ Gpxq. Soft prompt tuning concatenates a trainable prompt u so that the model output isGppu, epxqq. We consider simultaneously training the prompt parameter u and a classification head to fit the downstream task.
Notations. Let ∆d denote the space of d-dimensional probability vectors. We work with discrete random variables V taking values in a finite set V . We use P rV s P ∆|V| to denote the distribution of V and P rU |V “ vs P R|U | the conditional distribution of U given V “ v. PrpV “ vq P r0, 1s will denote the probability that V takes values v. We also let P rU “ u |V s P R|V| denote the vector with entries PrpU “ u |V “ vq. P rU |V s P R|U |ˆ|V| will describe the matrix with entries P rU |V su,v “ PrpU “ u |V “ vq. For a sequence v “ pv1, . . . , vtq, we use the notation vi:j for i ď j to denote pvi, . . . , vjq, and v´i to denote pv1:i´1, vi`1:tq. We let 1 denote the indicator function. For set V , we let V˚ “ V1YV2Y¨ ¨ ¨ denote variable-length sequences of elements of V . Let d denote elementwise product. Let 1d,0d denote the d-dimensional all-1’s and all-0’s vector. We omit the subscript if the dimension is clear from context. For two vectors a, b P Rd, we let a{b denote their element-wise division. We use supppaq to denote the set of indices where vector a is non-zero.

heading: 3 Analysis for Hidden Markov Models

text:
Defining a relation between pretraining and downstream tasks is the foremost challenge for analysis. We propose to link the two via latent variable generative assumptions on the input distribution. We
model the downstream task as a function of the posterior distribution of the latent variables. Towards a first result, this section studies the case where inputs are generated by HMMs (see Figure 1 (left)), which have been well-studied in the context of language and speech processing (see e.g. [24, 17, 3]).
Data distribution. Let H denote the hidden state space of the HMM. We use H “ pH0, H1, . . . ,HT q P H˚ to denote the sequence of hidden states. For all timesteps i ą 0, the transition probabilities are time-invariant, i.e. P rHi |Hi´1s “ A for A P R|H|ˆ|H|. For each timestep i ě 1, tokens Xi are emitted following some time-invariant probability: P rXi |His “W for W P R|X |ˆ|H|. The joint probability of X,H is
PrpX,H “ x, h |T “ tq “ PrpH0 “ h0q t ź
i“1 PrpHi “ hi |Hi´1 “ hi´1qPrpXi “ xi |Hi “ hiq.
Downstream tasks. We assume that H0 has the meaningful information for the downstream task, which is a binary classification task where the ground-truth labeling F ‹ is assumed to be a linear classifier on the posterior P rH0 |X1:T “ xs:
F ‹pxq “ 1pµJP rH0 |X1:T “ xs ě 0q (3.1)
for µ P R|H|. Our results are easily extended to the multiclass setting. We consider tuning a linear head for the downstream classifier, which formally computes 1pbJG1pxq ě 0q for b P R|X |. The following non-degeneracy condition is crucial for our recovery result in this setting. Assumption 3.1 (Non-degeneracy, vanilla HMM). The token emission probability matrix W has linearly independent columns.
We also require the following regularity conditions on H0 and the state transitions. Assumption 3.2 (Regularity). The Markov chain H0, H1, . . . is ergodic, and P rH0s has full support.
We show that if W has linearly independent columns, a linear head fits downstream labels. Theorem 3.3. Assume that non-degeneracy (Assumption 3.1) and regularity (Assumption 3.2) hold. Then any downstream task F ‹pxq of the form (3.1) can be computed by a linear head onG applied to a shifted sequence. That is, there exists linear head weights b P R|X | such that for all x P supppP rXsq,
F ‹pxq “ 1pbJG1px1q ě 0q
where x1 “ p∅, x1:tq is the concatenation of a special token ∅ with x.1
The key for the proof is to leverage the following general statement about random variables U, V, Z such that U K V |Z, which decomposes the expression for P rU |V s.
1We note that G1px1q does not depend on x11 and therefore x11 can be any token.
Proposition 3.4. Let U, V, Z be random variables such that U K V |Z. Then for any v, P rU |V “ vs “ P rU |Zs ¨ P rZ |V “ vs. Thus, if P rU |Zs has a left inverse pP rU |Zsq:, then P rZ |V “ vs “ pP rU |Zsq:P rU |V “ vs.
By the conditional independence structure of the HMM, Proposition 3.4 immediately implies
G1px1q “WP rH1|X2:T`1 “ xs ùñ P rH1|X2:T`1 “ xs “W :G1px1q where W : is the left inverse for W , guaranteed to exist by Assumption 3.1. This lets us recover P rH1|X2:T`1 “ xs by applying a linear function to G1px1q. Additional linear functions will be sufficient to obtain µJP rH0|X1:T “ xs from P rH1|X2:T`1 “ xs. We provide the full proof in Section B.
Proposition 3.4 is reminiscent of the arguments of [18], which leverages the independence structure in the same way. Subsequent sections will require more complicated analyses and recovery procedures.
A drawback of Theorem 3.3 is that it relies heavily on assuming W has full column rank, which implies the necessary condition that |H| ď |X |. Without this assumption, it is unclear how to recover P rH0 |X1:T “ xs from Gpxq alone. However, in realistic settings we would expect |H| ą |X |, as increasing the size of the hidden state space improves language modeling capabilities of HMMs [3].

heading: 3.1 Relaxed non-degeneracy assumptions via prompt tuning

text:
In this section, we study applying soft, or continuous, prompt tuning [19, 9] to the setting above. We show that by using soft prompt tuning, we can recover F ‹ using a linear head on G for HMMs where the non-degeneracy assumptions on W are relaxed. Our analysis provides insight into the empirical successes of prompt-tuning: intuitively, prompt tuning enables better recovery of the downstream task by conditioning the output of G to only contain task-specific information.
Soft prompt tuning trains task-specific embedding vectors, but analyzing how the model processes embedding vectors is challenging because it requires opening up the black box of the pretrained model. Thus, we require additional abstractions about how the pretrained model processes the embedding vectors. We will extend the mask language model G to a model G that maps a sequence of embeddings e1, . . . , et to conditional probabilities G1pxq, . . . , Gtpxq as follows. We observe that each token z in the vocabulary X naturally corresponds to a |H|-dimensional vector: the z-th row of the emission probability matrix W , or equivalently, P rXi “ z |His. We denote this embedding by epzq and call the family of embeddings tepzq : z P X u proper embeddings. A fundamental property of HMMs is that the conditional probability P rXi |X´i “ x´is only depends on x1, . . . , xt through their embeddings epxq “ pepx1q, . . . , epxtqq. In other words, there exists a function Gi such that
Gipx1, . . . , xtq “ Gipepx1q, . . . , epxtqq In particular, we let Gi compute the standard message passing algorithm [16] that computes the conditional probability of HMMs. This ensures thatGi is well defined on all sequences of nonnegative vectors in r0, 1s|H|, beyond sequences of proper embeddings.We assume that pretraining produces this Gi, which we treat as a blackbox for prompt tuning.
In particular, for prompt tuning we can consider the case where we pass an arbitrary nonnegative vector u P r0, 1s|H| to G in the first argument and proper embeddings at positions i ą 1. We can interpret u as the embedding of a fake token rz. Concretely, consider adding a new token rz to the vocabulary X , and changing the emission probability at position 1 to satisfy P rX1 “ rz |H1s “ u and for all z ‰ rz, P rX1 “ z |H1s9p1 ´ uq d epzq. Then Gipu, epx1q, . . . , epxtqq precisely computes the conditional probability P rXi |X´i “ prz, x1, . . . , xtq´is under the modified HMM. We refer the readers to Section C for the formal definition of Gi and formal proofs of the interpretation above.
We consider a downstream training algorithm which trains the prompt tuning parameter u described above and a linear classification head. Letting u denote the trainable prompt parameter and b P R|X | the trainable linear head weights, the model uses the embedding sequence
pepxq fi pu, ep∅q, epx1q, . . . , epxtqq (3.2) and outputs the prediction F pxq “ 1pbJG2ppepxqq ě 0q. We can provide recovery guarantees for this model if the ground-truth classifier weights µ (defined in (3.1)) and columns of the HMM transition matrix A satisfy the following relaxation of the requirement in Theorem 3.3 that W is nondegenerate.
Assumption 3.5 (Relaxed non-degeneracy condition). There exists a set of essential hidden states H‹ Ď H, so that the columns of W corresponding to H‹, tW:,huhPH‹ , are linearly independent. Furthermore, H‹ covers all meaningful information for the downstream tasks: supppµq Ď H‹.

heading: In addition, a last technical requirement on H‹ is as follows: there exists a set B Ď H such that

text:
H‹ “ YhPBsupppA:,hq. In other words, H‹ must be the set of all states reachable by starting from some state in B and transitioning one step in the hidden Markov chain.
Compared to Assumption 3.1, which required that all columns of W are linearly independent, Assumption 3.5 only requires linear independence on a subset H‹ of essential states. In the setting where |H| ą |X |, the condition for Theorem 3.3 can never hold. On the other hand, Assumption 3.5 could still hold, for example, if |supppµq| ă |X | and the set of columns of W corresponding to hidden states in supppµq is linearly independent. The last technical requirement in Assumption 3.5 is also required, which could be satisfied if columns of A are sparse. The following theorem shows that when Assumption 3.5 holds, we can recover F ‹ using soft prompt tuning with a linear head.
Theorem 3.6. In the above setting, assume that Assumptions 3.2 and 3.5 hold. Then F ‹ can be computed using soft prompt tuning with a linear head on G. Concretely, there is a continuous prompt parameter u P R|H| and weight vector b P R|X |, such that for all x P supppP rXsq,
F ‹pxq “ 1pbJG2ppepxqq ě 0q
where pe prepends u to the input embedding sequence, as defined in (3.2).
Theorem 3.6 provides a stronger recovery result than Theorem 3.3, which only used a linear head. This is also reflected in our synthetic experiments (Section 5), and prior work which shows that variants of prompt tuning can perform much better than only training the last few layers of the model [21]. Our theory suggests that prompt tuning could help by conditioning the hidden variables to remove nonessential information for the task from the output of G. This makes task-essential information easier to recover.
The key proof intuition is that although recovering P rH0 |X1:T “ xs is impossible without strong non-degeneracy conditions (Assumption 3.1), we can aim to recover P rH0 |X1:T “ xs on the subset of essential states H‹ defined in Assumption 3.5, which suffices for computing µJP rH0 |X1:T “ xs, since H‹ Ě supppµq. To recover P rH0 |X1:T “ xs on H‹, we observe in Lemma C.2 that prepending the prompt u is equivalent to introducing a modified random sequence pX and fake token rz which influences the posterior of H2 as follows:
G2ppepxqq “ rxWDpP rH2 | pX1 “ rzs d P rH0 |X1:T “ xsq (3.3)
for invertible diagonal matrix D and positive scalar rx. We choose u so P rH2 | pX1 “ rzs d P rH0 |X1:T “ xs is supported only on H‹. As corresponding columns of W are linearly independent (Assumption 3.5), we recover PrpH0 “ h |X1:T “ xq for h P H‹ via a linear function of G2ppepxqq. This suffices for computing µJP rH0 |X1:T “ xs. For more details, see Section C.

heading: 4 Analysis for memory-augmented Hidden Markov Models

text:
We study a memory-augmented HMM which explicitly disentangles the evolution of hidden states from a persistent “memory” variable. Inspired by natural sentences, this model is intended to better capture the distinction between syntax, which constantly evolves, and semantics, which changes less. This additional structure in the generative model allows us to strengthen our results by relaxing the non-degeneracy conditions on W , the token emission probabilities. Thus, both head and prompt tuning are more powerful in this setting compared to Section 3 and can recover the downstream label with weaker non-degeneracy assumptions on W . In Section 4.2, we show that soft prompt tuning also provides an advantage over head tuning alone.
Data distribution. The memory-augmented HMM, depicted in Figure 2, can be viewed as a generative variant of memory networks [39, 33] and is closely related to Hidden Topic Markov Models [8]. There are two sets of latent variables in the memory-augmented HMM: a Markov chain on hidden states H0, H1, . . ., meant to model the evolution of syntax, and a persistent “memory” M “ pM1, . . . ,MN q with N total cells, where each Mi takes values in a finite set M. The full joint
probability is as follows: PrpX,H,M “ x, h,m|T “ tq “
PrpM “ mqPrpH0 “ h0q t ź
i“1 PrpHi “ hi|Hi´1 “ hi´1qPrpXi “ xi|M “ m,Hi “ hiq
The hidden state is modified to explicitly consist of a disentangled cell index J P rN s and syntax state S P S, such that Hi “ pJi, Siq and H “ rN s ˆ S. To sample the token at timestep i given the hidden state Hi “ pJi, Siq, we first use Ji to index the memory M , obtaining the random variable MJi . Xi is then sampled according to some time-invariant probability depending on MJi , Ji, Si: P rXi |M “ m,Hi “ pj, sqs “ P rXi |MJi “ mj , Hi “ pj, sqs “W:,pmj ,j,sq Here W P R|X |ˆ|M||H| stores the emission probabilities for each choice of memory cell value and hidden state. Note that in particular, the conditional probabilities for Xi only depend on a single memory cell for each timestep. We also note that memory-augmented HMMs can be viewed as vanilla HMMs with structured transitions because pH0,Mq, pH1,Mq, . . . can be viewed as a Markov chain where the memory component does not change. Example 4.1 (Generating natural sentence with memory-augmented HMM). We consider how this model may generate the sentence “The cow in the pasture rolled on the grass’ happily.” M1 could store the subject (“cow”), M2 the location (“pasture”), M3 the sentiment (“happily”), and Si could determine part-of-speech. For timesteps where “cow” and “rolled” are emitted Ji “ 1 because we emit information related to the sentence subject. Timesteps for “pasture” and “grass” have Ji “ 2.
Downstream tasks. We consider downstream tasks where ground-truth labels are obtained via a linear classifier on the posterior distribution of a particular memory cell j‹ P rN s: F ‹pxq “ 1pµJP rMj‹ |X1:T “ xs ě 0q, where µ P R|M|. Intuitively, this formulation models downstream tasks which depend on a particular aspect of the semantics but not on syntax (e.g. in the setting of Example 4.1, if j‹ “ 3, the task is sentiment analysis).

heading: 4.1 Tuning attention head for recovering ground-truth downstream labels

text:
To recover the downstream labeling, we require an attention-based classification head, which is a function of both the input embeddings and outputs of G. Formally, let q P R|H|`1 denote a query parameter and β1, . . . , βt P R|H|`1 denote trainable position embeddings. Given pretrained model outputs Gipxq and trainable token embeddings epxiq, the attention head Attnp¨q applies key and value functions K,V to compute the output as follows:
I fi arg max i
tqJpKpGipxqq ` βiqu (4.1)
AttnppGipxq, epxiqqti“1q fi 1 |I| ÿ
iPI V pGipxq, epxiqq (4.2)
where arg max refers to the set of indices achieving the maximum in (4.1). We note that standard attention heads in practice rely on the softmax function, but the expression based on arg max above captures the limiting behavior as }q}2 Ñ8. We consider linear key functions given by KpGipxqq “ ΘpKqGipxq. The value function V : R|X | ˆ R|M||H| Ñ R uses parameters ΘpV q P R|M||H|ˆ|X | and b P R|M||H| and computes V pGipxq, epxiqq “ bJppΘpV qGipxqq d epxiqq. Because our generative model disentangles H and M , we can relax the non-degeneracy assumption on the token emission probabilities W , compared to Theorem 3.3. The relaxed assumption only requires the columns tW:,pm,hqumPM,hPH‹ to be linearly independent in a subset H‹ of “recoverable” hidden states, whereas Assumption 3.1 required all columns to be linearly independent. Assumption 4.2 (Existence of “recoverable” hidden states). There exists a set of recoverable hidden states H‹ “ tj‹u ˆ S‹, such that the collection of token emission probabilities from M ˆ H‹, tW:,pm,hqumPM,hPH‹ , is a linearly independent set of vectors. Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from Mˆ pHzH‹q: spanptW:,pm,hqumPM,hPH‹q X spanptW:,pm,h1qumPM,hPHzH‹q “ t0|X |u.
Note that the non-degeneracy condition of Theorem 3.3 would require tW:,pm,hqumPM,hPH to be linearly independent, whereas Assumption 4.2 only requires linear independence for h P H‹. The second condition states that H‹ and HzH‹ are distinguishable by the token emission probabilities. We explain Assumption 4.2 in the setting of Example 4.1. For natural language, there might be choices of h “ pji, siq for which the set tW:,pm,hqumPM of token emission probabilities is fundamentally not very diverse, and therefore not linearly independent. For example, if the syntax si indicates “article”, i.e. words such as “a”, “an”, and “the”, the token emission probabilities would carry little information about Mji because the choice of article does not depend much on semantics, so columns corresponding to si “ “article” would not be linearly independent, violating Assumption 3.1. However, Assumption 4.2 allows us to avoid this issue by placing such h in HzH‹, a set of hidden states which we can ignore, and only including hidden states which carry a lot of information about M in H‹. In Example 4.1, when Ji “ 2 (location), Si “ “noun”, the position i should convey a lot about the location (in this case, “pasture”), so it is more reasonable to assume that tW:,m,humPM is linearly independent for this hidden state.
Thus, our aim is to focus on recovering information for the downstream task from positions i where Hi P H‹. Formally, we define the following set of input sequences containing positions i where the posterior of Hi given x´i concentrates on H‹:
R fi tpx1, . . . , xtq P supppP rXsq : Di with supppP rHi |X´i “ x´isq Ď H‹u (4.3)
The following theorem shows that under Assumption 4.2, we can recover F ‹ using the attention head described above, if x P R is nonempty. Note that R is nonempty if the posterior of Hi concentrates on H‹ for some i. For natural language, it is realistic to assume this can occur because syntactic aspects of a sentence are typically low-entropy when the full sentence is observed. Theorem 4.3. Assume that non-degeneracy (Assumption 4.2) and regularity (Assumption 3.2) hold. Define R as in (4.3). Then there exist an attention head on Gpxq and token embeddings epxiq such that the following holds for any x P R:
F ‹pxq “ 1pAttnppGipxq, epxiqqti“1q ě 0q where the function Attn is in the form described in (4.2).
The idea is to use the attention mechanism to attend to positions i where supppP rHi |X´i “ x´isq Ď H‹. The intuition of Assumption 4.2 is that such positions are more informative for recovering the latent posteriors; indeed, from the outputs Gipxq at such i, the value function in the attention will be able to recover P rMj‹ |X1:T “ xs. A full proof is provided in Section D.1.

heading: 4.2 Guarantees for prompt-tuning

text:
Though the generative modeling assumptions in this section already allowed relaxed non-degeneracy assumptions, applying soft prompt tuning allows us to relax them even further. For simplicity, we consider the setting where there is a single memory cell, so M P M, and the downstream task is a linear classifier on the posterior of the memory: F ‹pxq “ 1pµJP rM |X1:T “ xs ě 0q. This
simplified setting doesn’t require the explicit disentanglement between Ji and Si in Hi. We analyze continuous prompt-tuning in a setting where the pretrained model G follows the same abstraction as in Section 3.1. We modify the model to take |M||H|-dimensional vectors, so the proper embedding for token z is given by epzq “ P rXi “ z|M,His “ WJz,:. In Section D.3, we describe the formal construction and interpretation of G in the more general setting with more memories.
Letting u P R|M||H| denote the trainable prompt parameter, we define the input embeddings pepxq fi pu, epx1q, . . . , epxtqq (4.4)
The downstream model applies an attention head to the output of G: F pxq “ 1pAttnppGippepxqq, peipxqqt`1i“1q ě 0q, where Attn is defined in (4.2). An additional stationarity assumption on P rH0s will simplify the recovery procedure (though it can be removed). Assumption 4.4 (Stationarity). Assumption 3.2 holds on the Markov chainH0, H1, . . .. Furthermore, P rH0s is the stationary distribution: P rH0s “ AP rH0s, where A is the transition matrix.
As before, we assume sparsity of µ and some non-degeneracy of W , though the assumption is more relaxed and easier to state compared to the vanilla HMM setting. Assumption 4.5 (Relaxed version of Assumption 4.2). Let M‹ fi supppµq denote the set of non-zero coordinates in µ. There exists a set of recoverable hidden states H‹, such that the collection of token emission probabilities from M‹ ˆH‹, tW:,pm,hqumPM‹,hPH‹ , is linearly independent. Furthermore, the span of these vectors must be disjoint from the span of token emission probabilities from M‹ ˆ pHzH‹q: spanptW:,pm,hqumPM‹,hPH‹q X spanptW:,pm,h1qumPM‹,hPHzH‹q “ t0|X |u.
We note that Assumption 4.5, and Assumption D.5 for multiple memories, are relaxations of Assumption 4.2, as they only consider memory values in supppµq, whereas Assumption 4.2 considers all m PM. An additional advantage of the memory-augmented HMM is that Assumption 4.2 is simpler than Assumption 3.1 and does not require any conditions on the transition matrix A. We now state our result for recovering F ‹ with soft prompt tuning and an attention head. Theorem 4.6. In the setting above, suppose that non-degeneracy Assumption 4.5 and stationarity Assumption 4.4 hold. Then there exists a prompt u and attention head on Gppepxqq and the token embeddings which can compute the ground-truth F ‹pxq for any x P R, defined in (4.3):
F ‹pxq “ 1pAttnppGippepxqq, peipxqqt`1i“1q ě 0q where pe is the embedding in (4.4) and Attn is defined in (4.2).
The intuition for this proof is similar to Theorem 3.6: the soft prompt conditions the memory M to concentrate on supppµq. As a result, all irrelevant information to the task is removed from Gippepxqq, making it easier to recover the task-specific information about the posterior of M . A more general theorem statement for the multiple memories setting, and the full proof, is provided in Section D.3

heading: 5 Simulations

text:
We empirically evaluate our theoretical results by pretraining a BERT-like masked language model (MLM) [4] on synthetic data generated by an HMM. Our goal is to verify key implications of our theory in a more realistic setting where some assumptions, such as that G outputs exact conditional probabilities, may not hold. First, we compare head and prompt tuning and show that prompt tuning improves downstream performance, especially when the recovery problem is degenerate. Second, we compare the effect of changing the data distribution from vanilla HMMs to memory-augmented HMMs on head tuning with an attention layer. We find that the downstream performance improves when the data has a long-term memory component. These observations support our theory.
Pretraining data and downstream task. We generate pretraining data from an HMM with randomly generated transition matrix, emission probabilities, and start distributions. In all experiments, the HMMs have 10 vocabulary symbols, while the hidden state size varies. The downstream task uses input sequences X1:T of length 129, where the first token X1 “ [MASK]. We consider binary classifcation where labels are generated using linear functions of the analytically-computed posteriors in the HMMs. In all experiments, the ground truth linear weight is sparse with 6 nonzero entries at uniformly random locations with Gaussian values. More details are in Appendix E.
Head vs. prompt tuning. We compare head and prompt tuning as the hidden state size of the HMM varies. The downstream label is computed via µJP rH1 |X´1 “ x´1s, where µ is a random ground-truth linear weight. Head tuning learns a linear head on top of the softmax probabilities predicted by the pretrained model for filling in the first [MASK] token. Prompt tuning uses the same setup but also optimizes a length 20 continuous embedding prepended to the input sequence.
Figure 3 (left) shows that prompt tuning improves downstream performance substantially across all hidden state sizes ({4,8,10,15,25,30}). Prompt tuning improves especially when the hidden state size increases beyond the vocabulary size, which makes the recovery problem degenerate. Thus, as suggested by Theorem 3.6, prompt tuning helps relax the non-degeneracy conditions.
Memory-augmented HMMs. We investigate the effect of augmenting the data-generating HMM with a long-term memory. We consider the single memory case with |H| “ 4 and varying memory sizes |M| P t2, 3, 5, 7u. The downstream label is generated by computing µJP rM |X´1 “ x´1s, where µ denotes the ground-truth weights. Viewing the memory HMM as a HMM where the component on M never changes, we can compare against the vanilla HMMs from the previous setting. For the memory-augmented HMM, we use head tuning with a single-cell attention layer on the entire sequence of softmax probability outputs. For the vanilla HMM in the comparison, we use a linear head on the output at the first position, as an attention head would perform worse since the downstream task depends only on H1 and not any other timesteps.
Figure 3 (right) verifies that head tuning recovers the downstream task better when there is more structure in the data, as predicted by Theorem 4.3. Head tuning achieves near 100% downstream accuracy on all hidden state sizes.

heading: 6 Conclusion

text: We analyze how pretraining on generic language modeling tasks can improve performance on diverse downstream tasks. In our analysis framework, the downstream task requires predicting properties of the posterior distribution over latent variables in an underlying generative model. When the generative model is a standard HMM, downstream recovery is possible with a simple classification head under strong non-degeneracy assumptions. We also show that we can relax the non-degeneracy conditions by changing the generative model to a memory-augmented HMM or using prompt tuning. The distributions studied here are meant to provide a first-cut result – we also expect similar theorems to hold for other generative models, which we leave as an interesting direction for future work.

heading: Acknowledgements

text:
We thank Percy Liang, Tianyi Zhang, and Nelson Liu for helpful discussions.
Funding statement
CW was supported by a NSF Graduate Research Fellowship. SMX was supported by a NDSEG Fellowship. TM acknowledges support of Google Faculty Award, NSF IIS 2045685, and JD.com. Additional revenue: CW received an honorarium for a talk at G-Research.

title: Why Do Pretrained Language Models Help in Downstream Tasks? An Analysis of Head and Prompt Tuning