[Paper] Provably Learning from Modern Language Models via Low Logit Rank

Source: arXiv - 2512.09892v1

Overview

Modern language models (LLMs) are notoriously intricate, yet recent empirical work shows that the matrix of their token‑level log‑probabilities often has low rank—meaning it can be well approximated by a small number of underlying factors. This paper leverages that observation to design an algorithm that provably learns such “low‑logit‑rank” models using the same kind of query interface that developers already enjoy via LLM APIs (e.g., “give me the probability of token X after this prompt”). The result is the first end‑to‑end learning guarantee for a generative model that plausibly mirrors the behavior of today’s LLMs.

Key Contributions

• Low‑Logit‑Rank Formalization: Introduces a clean mathematical abstraction of LLMs as matrices of log‑probabilities with approximate low rank, matching empirical findings.
• Query Learning Model: Defines a realistic query access model (logit queries) that mirrors the typical API calls developers make to LLM services.
• Efficient Learning Algorithm: Provides a polynomial‑time algorithm that, given logit queries, recovers an approximation of any low‑logit‑rank model up to provable error bounds.
• Hardness Insight: Shows that low‑logit‑rank models can encode notoriously hard distributions (e.g., noisy parity), highlighting why a specialized algorithm is needed.
• First Provable Guarantee: Delivers the first theoretical learning guarantee for a generative model that aligns with the empirical structure of modern LLMs.

Methodology

1. Matrix View of LLMs:

   • Treat the conditional log‑probability of each token (v) given a context (c) as an entry in a matrix (L) (rows = contexts, columns = tokens).
   • Empirically, (L) can be approximated by a rank‑(r) matrix (U V^\top) with (r \ll \min(#\text{contexts}, #\text{tokens})).

2. Query Model:

   • The learner can submit any context (c) and receive the full logit vector (L_c) (the log‑probabilities for all tokens).
   • This mirrors typical API calls where a prompt returns a probability distribution over the vocabulary.

3. Algorithmic Core:

   • Randomized Subspace Sampling: Randomly sample a modest number of contexts to capture the column space of (U).
   • Low‑Rank Matrix Completion: Use the observed logit vectors to solve a constrained regression that recovers (U) and (V) up to a small error.
   • Noise Handling: Incorporate robust statistical techniques (e.g., spectral filtering) to tolerate the noisy parity‑type distributions that can be hidden inside the low‑rank structure.

4. Theoretical Analysis:

   • Prove that with (O(r \cdot \text{polylog}(V))) queries (where (V) is vocabulary size) the algorithm outputs a model whose total variation distance from the true distribution is bounded by (\epsilon).
   • Show that the runtime is polynomial in (r), (V), and (1/\epsilon).

Results & Findings

• Sample Complexity: The algorithm needs only a linear number of queries in the rank (r) (often < 50 for real LLMs) and logarithmic dependence on the vocabulary size, making it practical even for large models.
• Error Guarantees: Guarantees an (\epsilon)-accurate reconstruction of the conditional distribution for any context, with high probability.
• Empirical Validation (Synthetic): Experiments on synthetic low‑rank logit matrices (including noisy parity embeddings) confirm that the algorithm recovers the underlying factors far more efficiently than generic black‑box learning methods.
• Alignment with Real LLMs: Preliminary analysis of GPT‑2‑like checkpoints shows that their logit matrices have effective rank ≈ 30–40, suggesting the theoretical regime is relevant for today’s models.

Practical Implications

• Model Extraction & Auditing: Developers can now extract a compact surrogate of a proprietary LLM using only the standard API, enabling offline analysis, debugging, or compliance checks without full model access.
• Efficient Fine‑Tuning: The low‑rank representation can serve as a lightweight backbone for downstream tasks, reducing compute and memory requirements when fine‑tuning on domain‑specific data.
• Distillation & Compression: Knowing that an LLM’s logits lie near a low‑rank subspace opens new avenues for knowledge distillation, where a smaller student model learns the subspace directly rather than mimicking the full output distribution.
• Robustness & Security: Understanding that low‑rank structure can hide hard distributions (e.g., noisy parity) informs adversarial testing—security teams can probe LLM APIs for hidden, hard‑to‑learn behaviors.
• API Design: Service providers might expose a “logit‑vector” endpoint deliberately, enabling developers to leverage the low‑rank learning algorithm for custom tooling while still protecting the underlying weights.

Limitations & Future Work

• Assumption of Exact Low Rank: Real LLMs exhibit approximate low rank; the current theory requires a bounded approximation error, and tighter characterizations for real‑world noise remain open.
• Query Cost in Practice: While the algorithm is query‑efficient theoretically, each query returns a full vocabulary logit vector, which can be bandwidth‑heavy for very large vocabularies. Future work could explore compressed query schemes.
• Extension to Context‑Dependent Rank: The rank may vary with context length or domain; adapting the algorithm to dynamic rank settings is an interesting direction.
• Empirical Evaluation on Closed‑Source APIs: Applying the method to commercial APIs (e.g., OpenAI, Anthropic) would validate practical feasibility and uncover engineering challenges.
• Beyond Generative Models: Investigating whether similar low‑rank abstractions hold for multimodal models (e.g., vision‑language) could broaden the impact of this line of work.

Authors

• Noah Golowich
• Allen Liu
• Abhishek Shetty

Paper Information

• arXiv ID: 2512.09892v1
• Categories: cs.LG, cs.AI, cs.DS, stat.ML
• Published: December 10, 2025
• PDF: Download PDF