Linearity of Relation Decoding in Transformer Language Models (2308.09124v2)

Abstract: Much of the knowledge encoded in transformer LLMs (LMs) may be expressed in terms of relations: relations between words and their synonyms, entities and their attributes, etc. We show that, for a subset of relations, this computation is well-approximated by a single linear transformation on the subject representation. Linear relation representations may be obtained by constructing a first-order approximation to the LM from a single prompt, and they exist for a variety of factual, commonsense, and linguistic relations. However, we also identify many cases in which LM predictions capture relational knowledge accurately, but this knowledge is not linearly encoded in their representations. Our results thus reveal a simple, interpretable, but heterogeneously deployed knowledge representation strategy in transformer LMs.

• The paper's main contribution is revealing that nearly 48% of relations in transformer LMs are linearly decodable using Linear Relational Embeddings (LREs).
• It introduces LREs to approximate subject-object mappings via affine transformations across 47 diverse relations in models like GPT-J and LLaMA-13B.
• Experimental results indicate a mix of linear and non-linear decoding, opening avenues for targeted knowledge insertion and refined LM probing techniques.

Exploring Linear Relation Decoding in Transformer LLMs

Introduction

In the contemporary landscape of neural LLMs (LMs), understanding the underlying mechanisms of relation representation and decoding has emerged as a pivotal research focus. Transformer LMs, despite their complex and highly non-linear architectures, demonstrate an intriguing behavior: certain relational knowledge can be linearly decoded from their embeddings. This paper explores the characterization of these linearly decodable representations within transformer LMs with respect to a diverse spectrum of relations encompassing factual, commonsense knowledge, and linguistic nuances.

Background: Knowledge Representation in LMs

The transformation of raw data into meaningful representations is a cornerstone of machine learning. In the context of transformer LMs, this transformation involves encoding a vast array of factual and commonsense knowledge into the model's weights. These encoded representations facilitate the regeneration of factually correct statements, a critical benchmark for evaluating LM efficacy. Prior research has illuminated the significance of the multi-layer perceptron layers within transformers, portraying them as key-value stores that enrich entity representations with pertinent knowledge.

Investigating Linear Relation Embeddings (LRE)

The paper introduces the concept of Linear Relational Embeddings (LREs) to approximate the decoding of relations by transformer LMs. LREs provide a simplified, linear framework for interpreting how models predict object entities based on subject-entity representations. The research posits that for a select subset of relations, the decoding mechanism in transformer LMs can be approximated by a linear function, culminating in an affine transformation operationalized through LREs.

Experimental Evaluation

A dataset spanning 47 relations was meticulously curated to facilitate the evaluation of LREs across various transformer models including GPT-J, GPT2-XL, and LLaMA-13B. The findings underscore the presence of robust LREs capable of faithfully recovering subject-object mappings for an array of relations. Notably, approximately 48% of the tested relations exhibited strong linear decodability, implying a pervasive yet heterogeneously applied linear knowledge representation strategy within LMs.

The discovery of a linear underpinning in the decoding process of transformer LMs unveils a nuanced understanding of knowledge representation strategies. This has profound implications, not only enriching our comprehension of LM internals but also opening avenues for refining LMs through targeted knowledge insertion and extraction methodologies.

Application and Future Directions

The unveilment of LREs paves the way for novel probing methods like the attribute lens, enhancing our prowess in visualizing and understanding the internal knowledge representations of LMs. Nevertheless, acknowledging the partial deployment of LREs across relations beckons further exploration into the mechanisms governing non-linear decoding processes within transformer architectures.

In summation, this paper accentuates the existence of linearly decodable relational knowledge in transformer LMs, fostering a deeper insight into the intricate tapestry of encoded knowledge and its implications for future AI research endeavors.