Latent Concept Disentanglement in Transformer-based Language Models

TL;DR

This paper investigates how transformer-based language models represent and disentangle latent concepts during in-context learning, revealing their ability to identify and utilize latent structures in reasoning tasks.

Contribution

It demonstrates that transformers can successfully identify, disentangle, and utilize latent concepts in various reasoning tasks, advancing understanding of their internal representations.

Findings

Models identify latent concepts in reasoning tasks.

Low-dimensional subspaces reflect underlying parameters.

Transformers effectively use in-context learned concepts.

Abstract

When large language models (LLMs) use in-context learning (ICL) to solve a new task, they must infer latent concepts from demonstration examples. This raises the question of whether and how transformers represent latent structures as part of their computation. Our work experiments with several controlled tasks, studying this question using mechanistic interpretability. First, we show that in transitive reasoning tasks with a latent, discrete concept, the model successfully identifies the latent concept and does step-by-step concept composition. This builds upon prior work that analyzes single-step reasoning. Then, we consider tasks parameterized by a latent numerical concept. We discover low-dimensional subspaces in the model's representation space, where the geometry cleanly reflects the underlying parameterization. Overall, we show that small and large models can indeed disentangle and utilize latent concepts that they learn in-context from a handful of abbreviated demonstrations.