Research map / Learning theory and algorithms / Sample complexity
Learning Partitions from Context
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies learning a hidden partition of N tokens into K classes from samples of a class-determined function f, giving complexity (NP-completeness), information-theoretic (N ln N samples), and gradient-flow analysis (N^2 ln^2 N samples) results.
Problem
Recovering the class memberships (partition) of discrete tokens from a finite number of sparse interaction samples of a function that depends only on the classes of its arguments, motivated by how LLM token embeddings may learn structure from next-token prediction.
Approach
Formalizes partition recovery as a logical/combinatorial inference problem with a class-determined function over I slots; proves information-theoretic sample bounds for random partitions, NP-completeness in general, and shows gradient flow on token embeddings recovers clusters for slot-wise linear functions given about N^2 ln^2 N samples.
Key results
- For random partitions, on the order of N ln(N) samples (times constants depending on I, K) suffice to identify the partition information-theoretically, and Theorem 2 shows this is essentially tight.
- Recovering the partition from samples is NP-complete in general (shown for I=3, K=2 with known g), ruling out general fast algorithms.
- Gradient flow on token embeddings keeps clusters stable and recovers cluster structure for slot-wise linear functions when S is on the order of N^2 ln^2(N) samples (with error scaling as sqrt(ln(S)N/S)).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use random token partitions and uniformly sampled interactions | Partition recovery · Random partitions | samples sufficient for partition recovery: order N ln(N) | Helped |
| Use slot-wise linear prediction functions | Gradient flow on token embeddings · Approximately balanced partitions | samples sufficient for cluster recovery: order N 2 ln(N) | Helped |
| Initialize embeddings sufficiently close to a global minimum and use sufficiently many samples | Gradient flow with general prediction functions · Approximately balanced partitions | cluster stability: S ≥ c1 max(K3IN 2 ln2(N), N ln(N)K9I/2) | Helped |
| Solve the general partition-learning decision problem | Partition learning with known interaction function · I = 3, K = 2 | computational complexity: NP-complete | Hurt |
Surprising findings
- Gradient descent on embeddings can fail to cluster even with arbitrarily small initial cluster spread when initialization is far from a minimum (Example 1 with I=K=D=1).
Limitations
- Gradient flow results need assumptions such as slot-wise linear functions, small embedding dimension, and initialization near minimum; analysis uses continuous-time MSE gradient flow instead of SGD with cross-entropy.
- No noise or errors considered; partitions are fixed and uniform, unlike real language structure.
- Gap between information-theoretic N ln N and gradient-based N^2 ln^2 N sample requirements remains open.
Experimental setup
Metrics: sample complexity
Builds on: matrix completion sample complexity
Closest papers by meaning
- Replicability in Learning: Geometric Partitions and KKM-Sperner Lemma (NeurIPS 2024 · similarity 0.85)
- Learning Representations for Hierarchies with Minimal Support (NeurIPS 2024 · similarity 0.84)
- Token Embeddings Violate the Manifold Hypothesis (NeurIPS 2025 · similarity 0.84)
- Learning In-context $n$-grams with Transformers: Sub-$n$-grams Are Near-Stationary Points (ICML 2025 · similarity 0.84)
- Active Learning of General Halfspaces: Label Queries vs Membership Queries (NeurIPS 2024 · similarity 0.84)
- Optimal Sample Complexity of Contrastive Learning (ICLR 2024 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.