Next Concept Prediction in Discrete Latent Space Leads to Stronger Language Models

Paper: 2602.08984 Authors: Yuliang Liu, Yunchong Song, Yixuan Wang, Kewen Ge, Alex Lamb, Qipeng Guo, Kai Chen, Bowen Zhou, Zhouhan Lin Categories: cs.CL, cs.LG

Abstract

This paper proposes Next Concept Prediction (NCP), a generative pretraining paradigm built on top of Next Token Prediction (NTP). NCP predicts discrete concepts that span multiple tokens using Vector Quantization (VQ). The resulting model, ConceptLM, leverages both NCP and NTP objectives jointly. Trained from scratch at scales from 70M to 1.5B parameters with up to 300B tokens of data, results across 13 benchmarks show consistent gains. Continual pretraining on an 8B Llama model demonstrates that NCP can further improve models already trained with NTP.

Key Contributions

• Next Concept Prediction: A novel pretraining objective that operates on multi-token concepts in discrete latent space
• ConceptLM: A model architecture combining NCP and NTP for stronger language modeling
• Vector Quantization for concepts: Using VQ to discretize multi-token spans into learnable concept codes
• Scalability: Demonstrated improvements from 70M to 8B parameters

Motivation

Standard next token prediction treats language as a flat sequence of tokens, ignoring the hierarchical nature of language where meaning is often carried by multi-token spans (phrases, entities, expressions). NCP addresses this by introducing a higher-level prediction objective that captures these multi-token concepts.

Method

Concept Extraction via Vector Quantization

Given a sequence of token embeddings $e_1, e_2, \ldots, e_n$, NCP groups consecutive tokens into spans and maps each span to a discrete concept code using Vector Quantization:

1. Span representation: For a span $(e_i, \ldots, e_{i+k})$, compute a pooled representation $z = \text{Pool}(e_i, \ldots, e_{i+k})$
2. Quantization: Map $z$ to the nearest code in a learned codebook $C = \{c_1, c_2, \ldots, c_K\}$:

$q = \arg\min_{c_j \in C} \|z - c_j\|_2$

3. Concept embedding: The quantized code $c_q$ serves as the discrete concept representation

The codebook is trained end-to-end with commitment and codebook losses:

$\mathcal{L}_{\text{VQ}} = \|z - \text{sg}[c_q]\|_2^2 + \beta \|\text{sg}[z] - c_q\|_2^2$

where $\text{sg}[\cdot]$ denotes the stop-gradient operator.

Joint NCP + NTP Training

ConceptLM is trained with a combined objective:

$\mathcal{L} = \mathcal{L}_{\text{NTP}} + \lambda \mathcal{L}_{\text{NCP}}$

where:

• $\mathcal{L}_{\text{NTP}} = -\sum_t \log P(x_t | x_{<t})$ is the standard next token prediction loss
• $\mathcal{L}_{\text{NCP}} = -\sum_s \log P(c_s | x_{<s})$ is the next concept prediction loss over concept spans
• $\lambda$ balances the two objectives

Architecture

The model uses a shared transformer backbone with two prediction heads:

• Token head: Standard vocabulary projection for NTP
• Concept head: Projects to the VQ codebook space for NCP

The concept head predicts the discrete concept code for the upcoming span, encouraging the model to plan ahead beyond the immediate next token.

Experimental Results

From-Scratch Training (70M - 1.5B)

Across 13 benchmarks spanning language understanding, reasoning, and generation:

• ConceptLM consistently outperforms NTP-only baselines at all scales
• Gains are particularly pronounced on tasks requiring compositional understanding
• The improvement scales with model size, suggesting NCP captures complementary information

Continual Pretraining on Llama 8B

Applying NCP as a continual pretraining objective on an already-trained Llama 8B model:

• Further improvements over the NTP-trained baseline
• Demonstrates NCP is complementary to, not redundant with, NTP training
• Relatively small additional compute yields meaningful gains

Benchmark Analysis

Performance improvements are observed across diverse task types:

• Language understanding: Better grasp of multi-word expressions and phrases
• Reasoning: Improved compositional reasoning, likely from concept-level planning
• Generation: More coherent long-form text generation

Analysis

What Do Concepts Capture?

Analysis of the learned codebook reveals that concepts correspond to meaningful linguistic units:

• Named entities (multi-token names, locations)
• Common phrases and collocations
• Syntactic patterns (prepositional phrases, verb phrases)
• Domain-specific terminology

Codebook Utilization

The VQ codebook shows healthy utilization patterns:

• High codebook usage (most codes are actively used)
• Semantic clustering (similar concepts map to nearby codes)
• Stable training dynamics with the commitment loss

Ablation Studies

Key ablations reveal:

• Span length: Moderate span lengths (3-5 tokens) work best; too short reduces to NTP, too long loses granularity
• Codebook size: Larger codebooks improve performance up to a point, then plateau
• Loss weight $\lambda$: Balanced weighting between NTP and NCP is important; too much NCP weight can hurt token-level prediction

Significance

NCP represents a meaningful step beyond the dominant next-token prediction paradigm. By introducing concept-level prediction through vector quantization, ConceptLM captures hierarchical structure in language that flat token prediction misses. The consistent improvements across scales and the ability to enhance already-trained models through continual pretraining make this a practical and impactful contribution to language model pretraining methodology.