One-sentence Summary

Researchers from Shanghai AI Lab, LUMIA Lab, and Shanghai Jiao Tong University introduce NCP-ArchPreview, an 8.9B-parameter latent-space language model that jointly trains next-token prediction with Next Concept Prediction (NCP) using a product-quantized concept vocabulary derived from hidden states, achieving superior pretraining efficiency and downstream performance while enabling lightweight domain adaptation and speculative decoding improvements.

Key Contributions

• NCP-ArchPreview is an 8.9B-parameter latent-space language model that jointly learns next-token and next-concept prediction, scaling latent-space modeling to trillion-token pretraining for the first time.
• NCP-ArchPreview reaches the final pretraining loss of OLMo-3-7B with only 51.3% of the training tokens, and after full pretraining outperforms that baseline by 2.45 downstream macro-average points, including a 5.99-point gain on GSM8K; it also approaches a parameter-aligned 8.9B Transformer using 85% of the computation and demonstrates a 1.74× compute efficiency improvement in scaling-law experiments.
• The learned concept space enables efficient domain adaptation by updating only the 17M-parameter VQ module, and injecting concept representations into a DFlash2 drafter improves speculative decoding mean accepted length by 4.17% with negligible overhead.

Introduction

In generative AI, the representation space a model learns can be as critical as its parameter count. While language models naturally develop high-level abstractions in their hidden states, standard next-token prediction treats these as a byproduct: supervision is confined to individual surface tokens, with no explicit objective guiding how semantic structure unfolds across longer spans. Prior hierarchical and latent-space architectures either fix the granularity of latent units or introduce new prediction targets that remain tied to token-level losses, leaving direct latent-space modeling at scale largely untested. The authors introduce NCP-ArchPreview, an 8.9B-parameter foundation model that jointly pretrains with next-token prediction and a novel next-concept prediction objective. By predicting discrete concepts over multi-token spans from the start of pretraining, and scaling to 5.73 trillion tokens, the work demonstrates that latent-space language modeling can substantially improve training efficiency, downstream performance, and compute Pareto-optimality compared to a strong token-only baseline.

Method

The authors introduce NCP-ArchPreview, a latent-space language model that augments standard next-token prediction with a concept-level predictive pathway operating in a learned discrete latent space. The model compresses token sequences into continuous concept representations, discretizes them through product vector quantization, and autoregressively predicts future concepts. These predictions are then fused back into the token stream to guide token generation. Hierarchical residual connections further mix information across layers and modules, while the entire system is trained end-to-end with a joint objective.

The backbone is organized into three modules: a Token Encoder, a Concept Module, and a Token Decoder. Given an input sequence $x_{1:T}$x1:T​, the Token Encoder produces token-level hidden states $\mathbf{h}_{1:T}$h1:T​. A mean-pooling operator groups every $k$k consecutive token states into a continuous concept representation $\mathbf{c}_m$cm​, yielding a concept sequence of length $M = \lfloor T/k \rfloor$M=⌊T/k⌋. These continuous concepts are then mapped to a finite discrete vocabulary via vector quantization, forming a structured target space for prediction. The Concept Module consumes the concept history and predicts the next concept as a differentiable combination of codebook entries. The predicted concept sequence is repeated to token resolution, causally shifted to prevent information leakage, and added to the token-level states before the Token Decoder performs next-token prediction.

To learn a discrete concept vocabulary, the authors employ product quantization. Each concept vector $\mathbf{c}_m$cm​ is split into $S$S segments of equal dimension, and each segment is assigned to the nearest entry in a dedicated codebook $\mathcal{E}^s = \{\mathbf{e}_1^s, \dots, \mathbf{e}_N^s\}$Es={e1s​,…,eNs​}. The quantized concept is the concatenation of the selected codewords. Product quantization yields $N^S$NS possible combinations while keeping individual codebooks small, greatly expanding the capacity of the discrete concept space. A vector quantization loss $\mathcal{L}_{\mathrm{VQ}}$LVQ​ moves the selected codebook entries toward the continuous representations using a stop-gradient on the encoder side, so that the codebooks track the latent distribution without directly altering the token-level hidden states.

The Concept Module predicts the next concept autoregressively from the concept history $\mathbf{c}_{<m}$c<m​. A stack of Transformer layers produces a latent state $\mathbf{u}_m$um​, and for each product quantization segment $s$s, a segment-specific prediction head outputs a probability distribution $\boldsymbol{\pi}_m^s$πms​ over the codebook entries. Instead of a hard discrete selection, the predicted segment is formed as the expectation $\hat{\mathbf{c}}_m^s = \sum_n \pi_{m,n}^s \mathbf{e}_n^s$c^ms​=∑n​πm,ns​ens​, and the full predicted concept $\hat{\mathbf{c}}_m$c^m​ is the concatenation of these segments. This keeps the prediction path fully differentiable while restricting the output to the linear span of the learned codebooks. The next-concept prediction is supervised by a mean-squared error loss $\mathcal{L}_{\mathrm{NCP}}$LNCP​ between $\hat{\mathbf{c}}_m$c^m​ and the detached continuous target $\mathbf{c}_m$cm​, which updates both the Concept Module and the Token Encoder through the preceding concept representations.

The predicted concepts are injected into the token stream by repeating each $\hat{\mathbf{c}}_m$c^m​ $k$k times and applying a causal shift $\Delta = k$Δ=k. For token positions $t \ge \Delta$t≥Δ, the concept signal $\mathbf{b}_t$bt​ is taken from the appropriately shifted concept index; earlier positions receive a zero vector. This ensures that no future-token information leaks into the prediction. The concept signal is added element-wise to the token hidden states: $\widetilde{\mathbf{h}}_t = \mathbf{h}_t + \mathbf{b}_t$ht​=ht​+bt​. The Token Decoder then models the next-token distribution $p(x_{t+1} \mid x_{\le t})$p(xt+1​∣x≤t​) from the fused states, trained with the standard causal next-token prediction loss $\mathcal{L}_{\mathrm{NTP}}$LNTP​.

To enhance information flow across different depths and granularities, the model incorporates hierarchical residual connections with two complementary mechanisms. Intra-Module Residual Connections (IRC) generalize the standard residual addition by allowing each layer to combine representations from multiple preceding depths. At layer $\ell$ℓ in a module, a candidate set $\mathcal{X}_\ell^s$Xℓs​ collects the initial state and all previous residual-updated states. A lightweight MLP applied to the current block output produces unnormalized coefficients, and the next hidden state is a weighted sum of the candidates. The MLP is initialized to select only the most recent residual-updated state, recovering the standard residual connection, while training can learn to incorporate earlier representations.

Cross-Module Residual Connections (CRC) transfer information between modules. When source and target modules operate at different sequence granularities, the source representations are first aligned by chunking or repetition. A softmax-normalized MLP over the target state produces depth-wise coefficients, and a weighted sum of layer-normalized source states is added to the target stream after scaling with a learned diagonal matrix. Three CRC pathways are used: TokenEncoder → ConceptModule, TokenEncoder → TokenDecoder, and ConceptModule → TokenDecoder. The last connection obeys the same causal shift as the predicted-concept pathway. Diagonal scalings are initialized to small values, so the model initially behaves like the original backbone and gradually learns to exploit cross-module signals.

Training is performed end-to-end, jointly optimizing the Token Encoder, Concept Module, Token Decoder, and codebooks. The total loss combines the next-token prediction loss, the next-concept prediction loss, and the vector quantization loss:

$$\mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{NTP}} + \alpha \mathcal{L}_{\mathrm{NCP}} + \beta \mathcal{L}_{\mathrm{VQ}},$$Ltotal​=LNTP​+αLNCP​+βLVQ​,

where $\alpha$α and $\beta$β control the auxiliary objectives. The NTP loss provides dense supervision for all parameters, the NCP loss encourages the encoder to retain information useful for concept prediction, and the VQ loss fits the codebooks to the concept distribution without backpropagating through the encoder. Optimization uses Moonlight Muon for matrix-valued parameters, with orthogonalized updates scaled by $\sqrt{\max(d_{\mathrm{in}}, d_{\mathrm{out}})}$max(din​,dout​)​, weight decay, and a cosine learning-rate schedule starting at $6 \times 10^{-5}$6×10−5. Embeddings, biases, and other non-matrix parameters are optimized with AdamW.

Experiment

The paper evaluates NCP-ArchPreview, a concept-predicting variant of OLMo-3-7B, using a two-stage curriculum and a comprehensive suite of 30 benchmarks spanning knowledge, reasoning, and code. NCP-ArchPreview consistently yields lower training loss and stronger downstream results than the baseline, with scaling experiments confirming a clear computational efficiency advantage. Ablations demonstrate that the Concept Module, hierarchical residual connections, and next-concept prediction loss each contribute meaningfully, while VQ-based adaptation updates only a small fraction of parameters to achieve competitive domain transfer with less forgetting and higher throughput. Concept representations also prove beneficial for multi-token prediction and block-parallel speculative drafting, and per-head normalization is identified as a remedy for training instabilities in the baseline configuration.

NCP-ArchPreview consistently outperforms the vanilla OLMo-3 model on all MMLU subsets across both training stages, with larger gains in Stage-1 than in Stage-2. The overall MMLU score improves by 2.58 points in Stage-1 and 1.82 points in Stage-2, driven by particularly strong advances in Humanities and Social Sciences. NCP-ArchPreview improves overall MMLU by 2.58 points in Stage-1 and 1.82 points in Stage-2 over the vanilla model. The largest absolute gains occur in MMLU-Humanities (+3.55 in Stage-1) and MMLU-Social Sciences (+2.79 in Stage-1), with all subsets showing positive improvements.

NCP-ArchPreview combines token encoder, token decoder, and concept module blocks to reach 40 block-equivalent parameters while using only 34 block-equivalent units of computation. It outperforms a vanilla Transformer with matched computation and approaches the performance of a vanilla model with matched parameters, despite requiring 15 percent less computation. The gains stem from the architecture rather than from simply adding parameters or compute. NCP-ArchPreview surpasses a vanilla Transformer that uses the same total computation (34F_blk). It nearly matches a vanilla Transformer with equal parameters (40P_blk) while using only 85 percent of the computation. Each concept module block adds parameters comparable to a standard block but incurs less than one quarter of the computation, enabling efficient scaling.

Adding intra-module residual connections (IRC) alone substantially reduces language modeling loss with a very small increase in analytical training FLOPs. Combining IRC with cross-module connections further improves loss without additional counted tensor contractions, while the full IRC + CRC design achieves the largest loss reduction at a higher FLOPs cost. A block-based attention residual variant underperforms all IRC-based configurations at a comparable compute budget. IRC alone delivers a loss reduction of 0.0273 with only 0.024% extra analytical FLOPs, offering a strong efficiency trade-off. IRC + CRC yields the best loss (0.0323 reduction) but requires 0.051% additional FLOPs, more than double the overhead of IRC alone. Adding input-level cross-module connections to IRC improves loss by 0.0023 over IRC only while keeping analytical FLOPs unchanged at 0.024%. Applying softmax normalization at all stages gives nearly identical loss to the cross-module variant (difference of 0.0003) at the same FLOPs. Block AttnRes with cross-module connections achieves only a 0.0180 loss reduction at 0.026% FLOPs, significantly smaller than any IRC-based variant.

VQ training adapts a model by updating only the codebooks and concept-prediction heads, using 17M trainable parameters without adding any new ones. In contrast, LoRA introduces 17M additional parameters to achieve the same trainable count, while full-parameter training updates all 8.9B existing weights. This makes VQ a parameter-efficient strategy that avoids increasing model size. VQ training updates 17M existing parameters and adds zero new parameters. LoRA matches the 17M trainable parameter count but adds 17M new parameters to the model. Full training updates all 8.9B parameters without adding any new ones.

Updating only the VQ codebooks and concept-prediction heads yields the highest code average and is the sole adaptation method that improves all four programming benchmarks. Full-parameter and LoRA fine-tuning both suffer large regressions on MBPP+, which lowers their overall code performance below the VQ result. Across all adaptation settings, the NCP-ArchPreview model retains higher general benchmark averages than the corresponding OLMo-3-7B variants. VQ adaptation improves every code task and achieves the highest code average among the adapted variants. Full-parameter training and LoRA both cause severe drops on MBPP+, reducing the full-training code average below its Stage-1 baseline. NCP-ArchPreview adapted with full training or LoRA outperforms the matched OLMo-3-7B controls on both code and general benchmarks.

The NCP-ArchPreview model consistently outperforms a vanilla Transformer on MMLU benchmarks, particularly in humanities and social sciences, while achieving parameter-matched performance with 15 percent less computation thanks to efficient concept module blocks. Intra-module residual connections substantially reduce language modeling loss with negligible overhead, and adding cross-module connections further improves the loss. Parameter-efficient VQ adaptation, which updates only existing codebooks and concept-prediction heads without adding new parameters, yields the best code benchmark results and avoids the regressions seen with full-parameter or LoRA fine-tuning, while NCP-ArchPreview retains higher general benchmark averages across all adaptation settings.