CUE Vectors: Modular Training of Language Models Conditioned on Diverse Contextual Signals

TL;DR

The paper addresses the challenge of incorporating diverse sentence-external context into language models without requiring joint training of sentence- and context-encoders. It introduces CUE, a modular framework with a sentence encoder, a DistilBERT-based context encoder, and a decoder that blends sentence-internal and contextual information into predictions. Empirically, conditioning on context reduces perplexity on NYTimes data from 36.6 to 27.4, and the framework supports strong adaptation: updating only the context encoder recovers about 85% of the full joint-training gain, while a proxy-embedding pretraining approach yields about 67% of that gain. The results demonstrate robust modularity, enabling context-aware LMs to be trained and deployed incrementally, with cross-architecture compatibility and the ability to swap sentence LMs without retraining context encoders.

Abstract

We propose a framework to modularize the training of neural language models that use diverse forms of sentence-external context (including metadata) by eliminating the need to jointly train sentence-external and within-sentence encoders. Our approach, contextual universal embeddings (CUE), trains LMs on one set of context, such as date and author, and adapts to novel metadata types, such as article title, or previous sentence. The model consists of a pretrained neural sentence LM, a BERT-based context encoder, and a masked transformer decoder that estimates LM probabilities using sentence-internal and sentence-external information. When context or metadata are unavailable, our model learns to combine contextual and sentence-internal information using noisy oracle unigram embeddings as a proxy. Real contextual information can be introduced later and used to adapt a small number of parameters that map contextual data into the decoder's embedding space. We validate the CUE framework on a NYTimes text corpus with multiple metadata types, for which the LM perplexity can be lowered from 36.6 to 27.4 by conditioning on context. Bootstrapping a contextual LM with only a subset of the context/metadata during training retains 85\% of the achievable gain. Training the model initially with proxy context retains 67% of the perplexity gain after adapting to real context. Furthermore, we can swap one type of pretrained sentence LM for another without retraining the context encoders, by only adapting the decoder model. Overall, we obtain a modular framework that allows incremental, scalable training of context-enhanced LMs.

Figures (6)

• Figure 1: Overview of CUE architecture. Pretrained sentence encoder and DistilBERT modules are frozen.
• Figure 2: Priming the decoder with proxy embeddings. We add noise to an embedding of the target sentence unigram distribution as a proxy for the decoder to learn to attend to context as yet unknown during training. Modules in gray are frozen during decoder training.
• Figure 3: Change in normalized attention weights over training iterations. The weights of the self-attention component of the context encoder converge to the relative importance of each contextual category over time, with previous sentence receiving the most weight.
• Figure 4: Varying proxy embedding strength. The baseline is no context (green line) versus the lower bound of knowing all context in training and test (red line). We sweep the amount of noise added to the oracle unigram vector on the x-axis. When training and testing on only the unigram vector (blue line) the unigram vector is a powerful oracle without noise, but then becomes random as the variance increases. During adaptation (orange line), we discard the unigram embeddings, freeze the decoder parameters, and retrain the context encoder (5% of parameters). The amount of embedding noise is optimal roughly when the proxy embedding is as informative as actual context (where blue and red lines intersect).
• Figure 5: Relative strength of each contextual type. We trained the CUE model with only one contextual signal at a time and measured perplexity on the same heldout test set. Textual context types (previous sentence, title, descriptor) are the most powerful.