Anchor function: a type of benchmark functions for studying language models

TL;DR

This work introduces the anchor function as a lightweight, controllable benchmark for studying transformer-based language models, addressing challenges faced by resource-limited academic groups. By formalizing one-anchor, two-anchor, and general anchor functions and pairing them with modular data-separation schemes, the authors show how anchor-driven tasks can instantiate diverse language tasks, analyze generalization, and probe mechanisms such as shift and broadcast in attention. Through nine task experiments and mechanism analyses, they demonstrate that transformers can learn anchor operations with data-efficient training, reveal how data structure shapes attention, and relate findings to existing theories like the frequency principle and condensation. The study further confirms that attention-based operations observed in large models like Llama2-6B align with the anchor-function framework, suggesting practical avenues for theoretical exploration and more accessible investigations into LM behavior and generalization. Overall, the anchor-function framework provides a cost-effective, interpretable platform for dissecting the internal workings of transformers and guiding future theoretical and empirical research in LM learning and generalization.

Abstract

Understanding transformer-based language models is becoming increasingly crucial, particularly as they play pivotal roles in advancing towards artificial general intelligence. However, language model research faces significant challenges, especially for academic research groups with constrained resources. These challenges include complex data structures, unknown target functions, high computational costs and memory requirements, and a lack of interpretability in the inference process, etc. Drawing a parallel to the use of simple models in scientific research, we propose the concept of an anchor function. This is a type of benchmark function designed for studying language models in learning tasks that follow an "anchor-key" pattern. By utilizing the concept of an anchor function, we can construct a series of functions to simulate various language tasks. The anchor function plays a role analogous to that of mice in diabetes research, particularly suitable for academic research. We demonstrate the utility of the anchor function with an example, revealing two basic operations by attention structures in language models: shifting tokens and broadcasting one token from one position to many positions. These operations are also commonly observed in large language models. The anchor function framework, therefore, opens up a series of valuable and accessible research questions for further exploration, especially for theoretical study.

Figures (28)

• Figure 1: Schematic diagram of three ways to split the training set and the test set.
• Figure 2: Task examples of anchor functions.
• Figure 3: Loss and accuracy of each task.
• Figure 4: Performance comparison of FNN, LSTM, and transformer. For each experimental setting, we used 10 random seeds for the experiment. Each experiment was trained for 4000 epochs. The scatter points are the average of the ten experimental results, and the shaded parts are the maximum and minimum values.
• Figure 5: The t-SNE distribution of token embeddings after training for numbers from 20 to 99.

Theorems & Definitions (1)

• Definition 1: mapping sequence