Learning and Transferring Sparse Contextual Bigrams with Linear Transformers

TL;DR

This paper introduces the Sparse Contextual Bigram (SCB) model, analyzes its training dynamics with linear transformers, and demonstrates how pretraining and finetuning can improve learning efficiency in language models.

Contribution

The paper proposes the SCB model as a natural extension of bigram models, analyzes its training process with linear transformers, and shows how pretraining enhances sample efficiency.

Findings

Training has an initial sample-intensive phase boosting correlation.

Finetuning from pretrained models bypasses the initial phase.

The proposed algorithm can outperform SGD in learning SCB.

Abstract

Transformers have excelled in natural language modeling and one reason behind this success is their exceptional ability to combine contextual informal and global knowledge. However, the theoretical basis remains unclear. In this paper, first we introduce the Sparse Contextual Bigram (SCB), a natural extension of the classical bigram model, where the next token's generation depends on a sparse set of earlier positions determined by the last token. We then analyze the training dynamics and sample complexity of learning SCB using a one-layer linear transformer with a gradient-based algorithm. We show that when trained from scratch, the training process can be split into an initial sample-intensive stage where the correlation is boosted from zero to a nontrivial value, followed by a more sample-efficient stage of further improvement. Additionally, we prove that, provided a nontrivial correlation between the downstream and pretraining tasks, finetuning from a pretrained model allows us to bypass the initial sample-intensive stage. We also empirically demonstrate that our algorithm can outperform SGD in this setting and discuss its relationship with the usual softmax-based transformers.