Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Provably Learning from Modern Language Models via Low Logit Rank

Noah Golowich, Allen Liu, Abhishek Shetty

TL;DR

The paper addresses the theoretical understanding of modern language models by leveraging the empirical observation of approximately low logit rank in logit matrices. It introduces a formal framework for approximate low logit rank and proves a polynomial-time learning algorithm that uses logit queries to recover a model close in total variation to the target. The core innovations are a two-part technical approach: (i) adaptive selection of futures via the elliptical potential lemma to cope with high-dimensional logit rows, and (ii) sampling via a linear-programming-based representation that bounds coefficient growth and yields an efficient sampler. The results provide the first end-to-end provable learning guarantees for a generative model that plausibly captures key aspects of modern language models, and they open avenues for further integration with practical training, online data, and interpretability. The work also connects to broader themes in learning theory, such as Kushilevitz-Mansour-style results for Boolean functions, under the low logit-rank paradigm.

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.

Provably Learning from Modern Language Models via Low Logit Rank

TL;DR

The paper addresses the theoretical understanding of modern language models by leveraging the empirical observation of approximately low logit rank in logit matrices. It introduces a formal framework for approximate low logit rank and proves a polynomial-time learning algorithm that uses logit queries to recover a model close in total variation to the target. The core innovations are a two-part technical approach: (i) adaptive selection of futures via the elliptical potential lemma to cope with high-dimensional logit rows, and (ii) sampling via a linear-programming-based representation that bounds coefficient growth and yields an efficient sampler. The results provide the first end-to-end provable learning guarantees for a generative model that plausibly captures key aspects of modern language models, and they open avenues for further integration with practical training, online data, and interpretability. The work also connects to broader themes in learning theory, such as Kushilevitz-Mansour-style results for Boolean functions, under the low logit-rank paradigm.

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.

Paper Structure

This paper contains 35 sections, 15 theorems, 57 equations, 1 figure, 3 algorithms.

Table of Contents

  1. Introduction
  2. Low Logit Rank
  3. Approximate Low Logit Rank
  4. Empirical Validation.
  5. Our Contributions
  6. Going forward: beyond the i.i.d. model.
  7. Technical Overview
  8. Learning with an approximate logit oracle
  9. Approximating rows of the logit matrix via the elliptical potential lemma
  10. Our approach: using the elliptic potential lemma to adaptively choose futures.
  11. Controlling the growth of coefficients via linear programming
  12. Our approach: "sampling by linear programming".
  13. Putting it all together: summary of the learning algorithm
  14. Learning a language model with approximately low logit rank
  15. Related Work
  16. ...and 20 more sections

Key Result

Theorem 1.2

There is a polynomial-time algorithm that, given the ability to make logit queries to a distribution $\mathbb{M}$ with $\varepsilon_{\mathsf{avg}}$-approximate logit rank $d$, outputs an efficiently sampleable distribution $\hat{\mathbb{M}}$ satisfying $D_{\mathsf{TV}}({\hat{\mathbb{M}}}, {\mathbb{M

Figures (1)

  • Figure 1: Low-rank approximation error (measured by average $L^1$ error) of the extended logit matrix for OLMo2-1b. For fixed sets $\mathcal{H}, \mathcal{F}$, the approximation errors for the logit matrix $\mathcal{L}_\mathbb{M}(\mathcal{H}, \mathcal{F})$ behave according to a similar power law as to those of various sub-matrices with $2$ or $4$ times fewer entries.

Theorems & Definitions (49)

  • Definition 1.1: Language Model
  • Definition 1.3: Logit Matrix
  • Definition 1.4: (Exact) Low Logit Rank
  • Definition 1.5: (Time-varying) ISAN
  • Definition 1.6
  • Definition 1.7: Logit oracle
  • Theorem 1.2: Main theorem; Informal version of \ref{['thm:approx-main']}
  • Definition 3.1: Low Rank Language Model
  • Definition 4.1: $\epsilon$-approximate logit oracle
  • Definition 4.2
  • ...and 39 more