Provably Learning from Modern Language Models via Low Logit Rank

TL;DR

This paper demonstrates that models with low logit rank, a structure observed in modern language models, can be efficiently learned with provable guarantees using a query-based algorithm, bridging empirical observations and theoretical understanding.

Contribution

It introduces an efficient algorithm for learning approximately low logit rank models, providing the first end-to-end provable learning guarantee for a model resembling modern language models.

Findings

Efficient algorithm for learning low logit rank models from queries.

Models with low logit rank can encode complex distributions like noisy parities.

Provides the first theoretical learning guarantee for a model similar to modern language models.

Abstract

While modern language models and their inner workings are incredibly complex, recent work (Golowich, Liu & Shetty; 2025) has proposed a simple and potentially tractable abstraction for them through the observation that empirically, these language models all seem to have approximately low logit rank. Roughly, this means that a matrix formed by the model's log probabilities of various tokens conditioned on certain sequences of tokens is well approximated by a low rank matrix. In this paper, our focus is on understanding how this structure can be exploited algorithmically for obtaining provable learning guarantees. Since low logit rank models can encode hard-to-learn distributions such as noisy parities, we study a query learning model with logit queries that reflects the access model for common APIs. Our main result is an efficient algorithm for learning any approximately low logit rank model from queries. We emphasize that our structural assumption closely reflects the behavior that is empirically observed in modern language models. Thus, our result gives what we believe is the first end-to-end learning guarantee for a generative model that plausibly captures modern language models.