Reconsidering the Past: Optimizing Hidden States in Language Models (2021)

Abstract: We present Hidden-State Optimization (HSO), a gradient-based method for improving the performance of transformer language models at inference time. Similar to dynamic evaluation (Krause et al., 2018), HSO computes the gradient of the log-probability the language model assigns to an evaluation text, but uses it to update the cached hidden states rather than the model parameters. We test HSO with pretrained Transformer-XL and GPT-2 language models, finding improvement on the WikiText-103 and PG-19 datasets in terms of perplexity, especially when evaluating a model outside of its training distribution. We also demonstrate downstream applicability by showing gains in the recently developed prompt-based few-shot evaluation setting, again with no extra parameters or training data.

## introduction finetuning a pretrained transformer language model (lm) @xcite @xcite is now the default method for attacking a task in modern nlp. due to the high cost of pretraining, much research has been focused on how better to apply the pretrained models, rather than just improving pretraining itself. however, even finetuning can be too costly, especially for models such as the 175 billion parameter @xcite . as such, researchers have sought low cost alternatives, such as finetuning a small set of auxiliary parameters @xcite , or more recently leaving the lm weights fixed and passing a textual context designed to elicit the desired behavior via token prediction, such as in @xcite . one direction for language modeling in particular is to leave the lm parameters fixed, but update its intermediate quantities (e.g., @xcite . in this paper, we introduce hidden-state optimization (hso), a method that contributes to this line of work. hso first computes the language modeling loss as usual, then modifies the lm hidden states using the gradient of the loss (but critically reports the original loss). this process is repeated for each window of 10-25 tokens, updating the cached hidden states each time. attending to these modified hidden states creates higher quality predictions for future tokens. as an example of how future information can help embed past tokens, consider the garden path sentence: "the old man the boat." the embedding for "man" will only depend on "the", "old", and "man", so it will not reflect that "man" is being used as a verb. hso can be seen as a method of incorporating future information into the representation of a context while still using a left-to-right lm. bert @xcite showed that bidirectional information passing improves embedding quality, which suggests that doing so should improve performance on downstream tasks. we demonstrate hso in the setting of language model evaluation on the wikitext-103 @xcite and pg-19 @xcite corpora, and find improvements in measured perplexity. in order to demonstrate that this translates into value for downstream applications we apply hso to fewshot classification with the 1.5b parameter gpt-2, and find improvement in that setting as well. ## method let f be a transformer language model computing the distribution for token x t given tokens x 1:t-1 : in practice, one may cache the hidden states, h @xmath0, where ℓ is the number of layers and d is the embedding size. we represent this by factoring f into f h which computes hidden states (possibly depending on past hidden states) and f p which computes output probabilities from the hidden states: given a loss function l which takes as arguments the ground truth next word and a distribution over word types, one can then compute its gradient with respect to both the present hidden states h t , and with respect to the cached hidden states h 1:t-1 : denoting the concatenation of these two quantities along the time axis as g @xmath1; g present , we can make a gradient update to the hidden states: in practice, we use standard cross entropy as our loss function l. so, intuitively, we are updating the hidden states to make the actual word at position t + 1 more likely under the language model's distribution p t by altering only the previously computed hidden states. note that when we update the hidden states with gradient-based updates, it will no longer be the case that the set of hidden states follow the feedforward procedure defined by the architecture of the transformer language model. while computing the hidden state for x t+1 , we then substitute h1:t into eq. 1 in place of h 1:t-1 : provided that the loss for timestep t is computed with the unmodified hidden state h t rather than ht , this may be done at test time without the loss being improved by "looking into the future." we continue to update all hidden states at each step. 1 in practice taking a gradient step after each token is too costly, so we can process blocks of k tokens (which we will refer to as a window size of k): this sequence of computations is done in a single forward pass, but we have broken it up by token to make clear how a mix of unmodified and modified hidden states is used to embed each token in the window. once the loss function, l, is applied to x t+2:t+k+1 and p t+1:t+k , a backwards pass is done to compute the gradient of the sum of the losses with respect to the hidden states, at which point the modified hidden states h1:t+k are computed. ## experiments we demonstrate hso with the transformer-xl (txl) @xcite and gpt-2 2 (radford et al., 2 for gpt-2, we backpropagate into the key and value vectors rather than the full embeddings at each layer for ease method wt-103 pg-19 baseline 21.3/22.4 166.4/164.2 hso 20.7/21.7 140.0/145.7 ## conclusion and future work we presented a method that optimizes transformer language model hidden states, which improves lm perplexity and prompt-based few-shot classification, without additional parameters or data. future work will explore improving the cost of hso by further investigation into updating only a subset of hidden weights, and approximation of the exact gradient update. other directions we will explore are its application to conditional generation by improving the representation of the context, and its interaction with other methods for improving prompt-based few-shot classification.