LLM Pretraining with Continuous Concepts
Authors: Jihoon Tack
Jack Lanchantin
Jane Yu
Andrew Cohen
Ilia Kulikov
Janice Lan
Shibo Hao
Yuandong Tian
Jason E Weston
Xian Li
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models
Pretraining
Concepts
Sparse Autoencoders
Primary area: foundation or frontier models, including LLMs

Abstract: Next token prediction has been the standard training objective used in large language model pretraining. Representations are learned as a result of optimizing for token-level perplexity. We propose Continuous Concept Mixing (CoCoMix), a novel pretraining framework that combines discrete next token prediction with continuous concepts. Specifically, CoCoMix predicts ``continuous concepts'' learned from a pretrained sparse autoencoder and mixes them into the model's hidden state by interleaving with token hidden representations. Through experiments on multiple benchmarks, including language modeling and downstream reasoning tasks, we show that CoCoMix is more sample efficient and consistently outperforms standard next token prediction and knowledge distillation. We find that combining both concept learning and interleaving in an end-to-end framework is critical to performance gains. Furthermore, CoCoMix enhances interpretability and steerability by allowing direct inspection and modification of the predicted concept, offering a transparent way to guide the model’s internal reasoning process.

Published as a conference paper at ICLR 2026
LLM PRETRAINING WITH CONTINUOUS CONCEPTS
Jihoon Tack Jack Lanchantin' Jane Yu', Andrew Cohen', Ili Kulikov', Janice Lan', Shibo Hao Yuandong Tian', Jason Weston', Xian Li' FAIR at Meta, "KAIST, SUC San Diego
ABSTRACT
Next token prediction has been the standard training objective used in large lan- guage model pretraining. Representations are learned as result of optimizing for token-level perplexity. We propose Continuous Concept Mixing (CoCoMix), a novel pretraining framework that combines discrete next token prediction with con- tinuous concepts. Specifically, CoCoMix predicts "continuous concepts" learned from pretrained sparse autoencoder and mixes them into the model's 's hidden state by interleaving with token hidden representations. Through experiments on multi- ple benchmarks, including language modeling and downstream reasoning tasks, we show that CoCoMix is more sample efficient and consistently outperforms standard next token prediction, knowledge distillation and inserting pause tokens. We find that combining both concept learning and interleaving in an end-to-end framework is critical to performance gains. Furthermore, CoCoMix enhances interpretability and steerability by allowing direct inspection and modification of the predicted concept, offering a transparent way guide the model's internal reasoning process.
INTRODUCTION
Recent progress large language models (LLMs) has revolutionized natural language processing (Brown et al., 2020; Dubey et al., 2024) and thus became core technology in various real-world applications, such as coding assistants (Roziere al., 2023), search engines (Xuan- Kuan--uuy& et al., 2023), and personal AI assistants (Gao et al., 2023). Central to these breakthroughs is the simple paradigm of next token prediction, which leverages massive amounts of unlabeled text to uncover rich linguistic patterns (Radford al., 2018: 2019). However, natural language tokens are often superficial (c.g.. function words like "the" or "a"), necessitating substantial training for models to acquire high-level reasoning and conceptual understanding while also hindering their ability to tackle long-horizon tasks such as planning (LeCun, 2022; Bachmann and Nagarajan, 2024).
To this end, recent studies have investigated methods that go beyond token-level signals by leveraging richer information to train models. For instance, some approaches target more expressive predic- tion objectives, such as predicting multiple tokens at once better capture semantic relationships (Gloeckle et al., 2024; DeepScek-At 2024), while others augment the input with rich signals, e.g.. self-generated thought tokens (Zelikman al., 2024), fixed pause tokens (Goyal et al., 2024) prior to next token prediction. Moreover, emerging evidence suggests that LLMs inherently encode high-level concepts and reasoning processes their latent representations (Deng et al., 2023; Yang et al., 2024), indicating that replacing discrete language tokens with continuous latent representations has promise in improving reasoning efficiency (Hao et al., 2024). While token-level modeling remains important for coherent text generation, the key challenge is to enrich supplement these language tokens so that LLMs can learn more abstract reasoning abilities and long-range dependencies.
This raises key question: can augment next token prediction objective to explicitly model concepts in a latent representation space, thereby bridging semantic abstraction and fine-grained token-level guidance?
To this end, we draw inspiration from recent findings that Sparse Autoencoders (SAEs) can effectively isolate meaningful latent features in LLMs by capturing the high-level semantic concepts (Cunning- ham et al., 2023; Bricken et al., 2023). As SAEs are trained to reconstruct the model's ''s hidden state with sparsity constraints, it encourages focusing on a compact set of concept dimensions (Templeton

---

Published as a conference paper at ICLR 2026
Using attribution to select concepts
"Continuous concept"
fcon
Gradient flow
f
conceppt
min CE:
Selected Concepts I
Predicted concepts
h hcoo
h
Extracting concepts from pretrained SAE model's hidden state
Leamning to predict concepts
Mixing/Iiterreaaving continuous concepts into the hidden state
Figure 1: Overview of CoCoMix. We SAE to extract concepts from pretrained model's hidden state hion and then select important concepts based on the attribution score (i.e., measuring the influence the output). These selected concepts used labels I for concept prediction by minimizing the cross-entropy loss .). The predicted concepts z are then compressed into a compact vector, forming continuous concept which mixed into the model's hidden state by interleaving with token hidden representations. Here, the token hidden representation, and continuous concepts, respectively.
and boxes indicate text tokens,
et al., 2024). This makes it possible to highlight the pretrained model's concepts- the core semantic directions that underpin model prediction-wile avoiding unnecessary features.
We propose Continuous Concept Mixing (CoCoMix), novel and effective language modeling framework using continuous concepts. Specifically, extract semantic concepts using a pretrained SAE and select the most influential ones based on attribution scores, which quantify each concept's influence the model's 's output. The model is then trained predict these selected concepts from its hidden state using a cross-entropy loss. Once multiple concepts are predicted, we compress them into a single continuous concept and mix (or insert) into the hidden states by interleaving with token embeddings, thereby directly contributing to the next token prediction. An additional benefit is that one can probe or analyze the predicted concepts, enabling controllable generation and improving model interpretability. An overview of CoCoMix is shown in Figure 1.
We demonstrate the efficacy of CoCoMix through extensive evaluations across multiple language modeling benchmarks and pretraining model sizes, ranging from million-scale to billion-scale parameter sizes. For instance, when applied to a 1.38B-sized model, CoCoMix achieves comparable performance with the next token prediction with 21.5% fewer training tokens. Moreover, CoCoMix demonstrates substantial improvements in weak-to-strone supervision scenarios, where concepts extracted from a small model can even be used ground truth labels to supervise the training of larger model. Finally, we show that the insertion compressed concept vectors enables us to steer and control the model by probing the predicted concept during generation.
COCOMIX: CONTINUOUS CONCEPT MIXING
In this section, we present Continuous Concept Mixing (CoCoMix), framework that extends next- token prediction with continuous concepts. We first review our problem setup of interest and the concept extraction process using a sparse autoencoder (SAE) in Section 2. Then, describe the concept selection framework and explain how the model predicts and interleaves these concepts into its hidden state to improve language modeling in Section 2.2.
PRELIM: PROBLEM SETUP AND SPARSE AUTOENCODER
Problem setup. We consider language greguage-modeing task vocabulary V, with training corpus D C V*, where each sequence x (æ), .2....) E D drawn i.i.d. from the data distribution. We require pretrained language model for concept extraction Moon seon hcon which predicts next
2

---

Published as a conference paper at ICLR 2026
token probability Meoa( x2t)). Here, extract the continuous concepts from the hidden state /eoo. where heen outputs the hidden state at a chosen layer of depth Lean, and Scea maps that representation to the next token distribution. We aim to train M to model the same autoregressive Zee, Cer), but by augmenting with continuous concepts c.
factorization M(x) =
Sparse Autoencoder. We first propose to extract high-level concepts from a pretrained model's hidden state. We use an SAE, which maps the input to a high-dimensional sparse activation, with the objective of reconstructing the original input (Lee et al., 2006). When applied to an LLM, the SAE decomposes the hidden state into multiple dimensions, each of which can be viewed as a distinct concept capturing semantically meaningful features (Yun et al., 2021; Bricken et al., 2023). We consider the Topk SAE (Makhzani and Frey, 2014), which uses a Topk activation to enforce sparsity.
Formally, for given imput sequence and corresponding pretrained model's hidden state at position t, denoted as h = Reon (x), E Rd..s, SAE consists of a linear encoder E : R - R and a linear decoder D RC - Rdim, where is the dimension of the concept space. The reconstruction
process of SAE is as:
c (h, C₁ TopK(c."")) hoo D(c.),
where c²re E RC is the pre-activation concept vector, TopK(-) zeros out all but the largest KSAE entries, and hoom is the reconstruction. The SAE is trained by minimizing the reconstruction loss: ||hçoa - h [2 By enforcing TopK sparsity, the SAE isolates the most critical dimensions in c that explain the pretrained model's features. Here, each activated element of c, interpreted as a concept.
2.2 CONTINUOUS CONCEPT MIXING
Now, describe the core training pipeline of CoCoMix, which consists of concept selection framework (refer to the left figure in Figure 1) and two training steps to learn and utilize continuous concepts (refer to the right figure Figure 1). First, we select important concepts using the attribution score, which measures the influence of each concept on the output. Then, we propose predicting the selected concepts from the model's hidden state using cross-entropy loss, allowing the model to implicitly learn which concepts should be encoded as hidden representation. Finally, we utilize the predicted concepts to create "continuous concept", which is interleaved in the hidden states, enabling the model to explicitly learn to use both the continuous concepts well as the token hidden states, Intuitively, the model selectively learns which concepts are useful for the next token prediction and how it should mix them with the token representations.
Target concept selection using attribution. While the extracted concept C₂ represents the core features of the current input context, it may not directly indicate which concepts matter most for predicting the ground truth next token 2++1 (Templeton et al., 2024). To this end, we propose utilizing attribution, method measure the causal influence of each concept on the output (Baehrens et al., 2010: Sundararajan et al., 2017). Specifically, the attribution score measures the influence based on the local linear approximation of the effect of changing the concept value. In this paper, we use simple attribution score that measures the influence by multiplying the loss gradient with given input (Simonyan and Zisserman, 2014; Shrikumar et al., 2016).
Concretely, for given input x and corresponding concept C₁ we define the attribution score a E RC;
a C₁ o Vo, - log feon (I++1] D(c:), her),
(I)
where denotes element-wise multiplication. Based the computed attribution score, we select salient concepts and incorporate them into language modeling.
Predicting the selected concepts. For a given attribution score a,, we first select the indices of the concept that have a high score and use these as discrete labels for concept prediction. Let I = (i.,..., ) be the set of indices corresponding to the top Katr values of a, and h, = h(()), Rd be the model's hidden state of the given input x at the same token position. Then, we learn to predict these labels using a linear prediction head M that outputs logit z, M (h.) E BC by minimizing the following cross-entropy loss CE(.,.) -) is as follows:
Lconceps
CE (zi, 1). Σ
(2)
Kam
3

---

Published as a conference paper at ICLR 2026
Mixing continuous concepts with token embeddings. To encourage the model to internalize the concept more effeetively, we propose "mixing" (i.e., interleaving) the predicted concept with the existing token hidden representations. As model predicts multiple concepts once, propose to compress them into a compact vector through learnable mapping, which we refer to as a continuous concept. This compressed vector is interleaved with token vectors.
Formally, for a given concept prediction logit 2,, we sparsify the logit using TopK activation to predict the concepts, then compress them into a continuous concept vector C₂ Rd.
C2 W W TopK(z)) b,
where W E RdxC and b Rd project the TopK-sparse vector d-dimensional embedding. We then append é, to the model's hidden sequence é, as (h1, c h,, G2), which is fed into the remaining transformer blocks f. Note that this design not only improves the model's performance but also offers interpretability and controllability through analyzing and steering the predicted concept 2t, which can be probed or tuned during the model's s generation process. Furthermore, by analyzing the weights of the compression layer W, one can identify which concept is useful for the next token prediction.
This approach shares similarities with intervention techniques in the mechanistic interpretability literature, which modify the hidden state by adding concept vector to the original hidden state (Zou et al., 2023; Wu et al.. 2024). However, unlike intervention methods that directly manipulate the hidden state, our approach treats the predicted concept as a separate unit of information by interlcaving it in the hidden state. This design allows the model to process the concept independently, leveraging the pretrained model's internal reasoning reflected in the prediction.
Training objective. The final training objective for CoCoMix combines the standard next token prediction loss and the concept prediction term with a tunable coefficient À as follows:
log f [I++1 lgg, Cgt) ACconcept
3 EXPERIMENTS
We provide an empirical evaluation of CoCoMix by investigating the following questions:
Can CoCoMix improve the performance of next token prediction? (Figure and Figure 3)
Does CoCoMix show improvement in weak-to to-strong supervision setup compared to other knowl- edge distillation methods? (Table 1 and Figure 4)
Does CoCoMix introduce model interpretability and steerability? (Figure 5)
How does each proposed component of CoCoMix contribute the performance? (Figure
Before answering each question, we outline the experimental protocol (more details in Appendix A).
pretrained open-source SAE that is trained the 1 124M-sized GPT-2 (Gao et al., 2024), which has demonstrated strong effectiveness despite a somewhat small model size (Chaudhary and Geiger, 2024). Here, we consider training CoCoMix with three different numbers of active parameters, including 69M, 386M, and 1.38B (which is an academic-scale pretraining (Allen-Zhu, 2025)), with context length of 1024. For the analysis and ablation study, we mainly conducted experiments on the 69M model. Note that this activated parameter count includes those in the new layer (i.e., the concept predictor), whereas for other baselines, we match the same number of active parameters by increasing the hidden state dimension size d. CoCoMix utilizes fewer FL.OPs than Pause token (as discussed in Section 3.3) but more FL.OPs than NTP, due to the interleaving of continuous concepts. We use the Open WebText dataset (Radford al., 2019) as the pretraining corpus to use the same distribution used to train Maan- All experiments are conducted with 20B training tokens, except for the main experiment in Section 3.1, which uses 200B tokens.
Training setup. We
Baselines. We consider the standard pretraining next token prediction (NTP) procedure, and the commonly used distillation method, knowledge distillation (KD) (Hinton al., 2015), which is widely used in pretraining (Gu et al., 2024; Sanh et al., 2019; Team, 2024). We excluded KD baselines that require multiple models (i.e., more than single teacher model) to be trained (Gu et al., 2024).
4

---

Published as a conference paper at ICLR 2026
15
NTP
20
NTP
+-
NTP
Volidriinn Perplexiny
(ours)
Avarage Parplostyy
CoCoMix (ours)
56
(ours)
14-
Accurssy (2))
13-
12
14-
52
21.555 loss tolsons
11
Traning colenns
Training
(B)
Training tokens
(c) Average downstream task Acc Figure 2: CoCoMix vs. NTP performance at different training checkpoints. Each model contains total of 1.38B parameters. Each model trained on the OpenWebText dataset. For CoCoMix, the concepts are extracted from a 124M-sized model (10x smaller than the base model). The plots show improvements in: (a) validation perplexity, (b) average perplexity on LAMBADA, Wiki'ext- 103, and (c) average accuracy on HellaSwag, PIQA, SIQA, Arc-Hasy, and WinoGrande. Here, CoCoMix achieves the same performance of NTP with 21.5% fewer tokens.
(a) Validation PPL.
(b) Average downstream task PPI.
Perplxay
Perpleccy
Perpoocy
Ascracy (30
400
15-
-11
-1.0
30
69M
38AM
188
69M
384M
1.388
69M
386M
1388
6991
386M
1.888
(a) Open Text, PPL (L) (b) LAMBADA, PPL (4) (c) WikiText- 103, PPL. (L) (d) HellaSwag, Acc-n (t)
05
19.2
Accursy (8)
Acursty (8)
Accurscy (8)
Accusscy (1)
+2
15
60
-0.1
50m
59M
386M
1.368
69M
386M
1.385
69M
386M
1.388
69M
386M
1.388
PIQA, Ace (t)
SIQA, Ace (t)
(g) Are-Easy, (t)
(h) WinoGrande, Acc (t)
Figure 3: CoCoMix vs. NTP performance at different model sizes. We consider various model sizes, including 69M, 386M, and 1.38B parameters and train on 200B Open tokens. We evaluate the models on Webeet validation perplexity and downstream datasets LAMBADA, WikiText-100 103, HellaSwag, PIQA. SIQA, Are-Easy, and WinoGrande. Overall, CoCoMix achieves stronger performance than NTP in virtually all datel-d combinations considered.
For KD, we minimize the KL divergence between teacher and student output while balancing the KL term with the NTP loss. Note that using synthetic datasets for distillation (Kim and Rush, 2016; Yu et al., 2024) is excluded due to inefficiency (needs to generate more than 20B tokens).
Evaluation setup. For evaluation, we consider the validation perplexity of the pretraining dataset and 7 downstream tasks to benchmark commonsense reasoning and reading comprehension. This includes LAMBADA (Paperno et al., 2016), WikiText- 103 (Merity et al., 2017). HellaSwag (Zellers al., 2019), PIQA (Bisk et al., 2020), Social IQA (SIQA) (Sap et al., 2019), ARC-easy (Clark et al., 2018), and WinoGrande (Sakaguchi et al., 2020) datasets. We also consider Open WebMath (Paster al., 2023) pretraining dataset to demonstrate that concepts learned from pretrained LLM can still be applied to CoCoMix, even when the model was trained on a different corpus.
3.1 MAIN RESULTS
In this section, we illustrate two core results: i) the comparison with NTP on a relatively large-scale pretraining setup and ii) the comparison with KD baseline, especially weak-to-stron supervision scenarios where concepts extracted from a small model are used to guide a larger model.
Improving NTP with CoCoMix at scale. We present the main result by applying CoCoMix to the NTP by training NTP and CoCoMix on 200B tokens. As shown in Figure 3, CoCoMix consistently and significantly improves the performance of overall downstream tasks various model sizes.
5

---

Published as a conference paper at ICLR 2026
40
NTP
-
NTP
- NTP
Average Perplesity
Average Accuracy (50
Volidation Perplexity
CoCoMix (ours)
CoCoMix (ours)
CoCoMix (ours)
30-
46
25
14
Traning coleons
Training
(B)
Training sokens
(a) Weak-to-strong: Avg. PPI. (b) Weak-to-strong: Avg. Acc
(c) Distribution shift: Math domain
Figure CoCoMix vs. Knowledge Distillation (KD). For our weak-to-strong supervision setup, we train 386M model where the teacher of KD (or concept extraction for CoCoMix) is 124M-sized model: we report (a) average perplexity over Web LAMABADA, and WikiText and (b) average accuracy over HellaSwag, PIQA, SIQA, Are-Easy, and WiniGrande dataset. For (c)
distribution shift setup, we train all methods on Open AppeWWeMath,, a math specific pretraining corpus.
Table 1: CoCoMix vs. Next Token Prediction (NTP) vs. Knowledge Distillation (KD). We report performance on the OpenWebTee (OWT) training set, as well as downstream tasks, including LAMBADA (LAMB), WikiText- 103 (Wiki), HellaSwag (HellaS), PIQA, Social Interaction QA (SIQA), Arc-Easy (Arc-E), and WinoGrande (WinoG). We train three different sizes of models where 124M model is used as a teacher. All models are trained on 20B tokens sampled from the OpenWebTer dataset. The bold indicates the best result.
Method
Total Params
OWT PPL
LAMB (P)
Wiki (P)
HellaS Acc-n
PIQA Acc (1)
SIQA Ace (t)
Are-E
WinoG
Avg PPL (U)
Avg
(t)
(?)
Ace (t)
NTP
69M
25.3
107.6
52.3
27.4
59.4
36.6
39.7
50.7
61.8
42.7
KD
69M
25.2
99.3
51.0
27.4
59.8
36.2
39.8
50.7
58.5
42.8
CoCoMix
37.2
39.3
51.0
42.9
NIP
386M
16.3
26.3
29.9
64.1
38.4
47.3
50.9
46.8
KD
386M
16.4
24.6
29.1
33.6
63.6
38.0
48.5
51.4
23.4
47.0
CoCoMix
19.3
47.9
52.3
1
16.6
50.5
KD
1.38B
14.2
16.6
24.9
37.4
66.7
39.0
50.1
52.3
18.5
49.1
CoCoMix
1.38B
13.9
15.4
24.9
38.4
66.9
39.5
50.8
53.0
18.1
49.7
Our results also indicate that larger models (e.g., 386M and 1,38B) can benefit from using concepts extracted from smaller 124M model, showing effective weak-to-stron supervision. Moreover, as shown in Figure 2, CoCoMix consistently improves the performance over the NTP on a billion-scale model. For instance, CoCoMix achieves similar performance to NTP while using 21.5% fewer tokens, demonstrating high sample efficiency. Finally, it is worth noting that the performance gain of using CoCoMix is increasing over the training steps, demonstrating strong generalization performance.
Comparison with KD baseline. We also compare CoCoMix with KD baseline across multiple scenarios, including (1) stronger teacher model teaching smaller student model: (2) weak-to-strong supervision, where a weaker teacher teaches larger student model; and (3) distribution shift, where the student trained on a corpus different from the teacher's pretraining distribution. As shown in Table 1, CoCoMix demonstrates improvements over KD in all considered model configurations. In particular, CoCoMix shows significant performance gain e.g., improving average perplexity of 2.8 in 386M, while KD does not show great improvement. This arises because weaker teacher can introduce noisy or suboptimal knowledge, especially when the student surpasses the teacher in capability (Rawat et al., 2024). This trend is also observable in Figure 4, where models trained with KD fall behind standard training midway through training as the student outpaces the teacher (especially in the distribution shift scenario). In contrast, CoCoMix selectively utilizes useful concepts, resulting in a consistent performance gain.
weak-to-strongg supervision setup,
3.2 INTERPRETABILITY AND STEERABILITY COCOMIX
Another advantage of CoCoMix is its interpretability and model steering. Specifically, the model is trained to predict concepts in its hidden state, we can analyze which concepts it focuses on based the concept predictions. Furthermore, by amplifying the magnitude of the predicted

---

Published as a conference paper at ICLR 2026
Input Prompt best platform for buying ticketa the
concapt 23843: site address related
Steer concept 15687: yearlmonth related
GPT2 SAE CoCoMix (ours)
Ganeration: vernacilar ticketing
Goneration: vernacular websita, Ticketrastar.co con. The site is a great resource for finding the best deals tickets.
Caneration: by and the new platforn by fans, The naw ticket- by a group of fans has been known the "ticket by 'ticket the
the year
syster, said the official, "The ticket prices are fixed and the tackets are available at the sane Time."
Generation: Sphona app-
Generation: urehin-cheers, online store new, fast-growing selection premium tickets
Generation: Facebook page. It's place where you can buy tickets next one. The other option is to buy ticket for a month or two
Figure 5: Qualitative demonstration of the concept steering effeet. CoCoMix and GPT2 models are 350M and 124M parameter transformers, respectively, trained on the OpenWebTes. dataset. For CoCoMix, we manipulate the predicted concept logit z, while for GPT2, we adjust the SAE concept space e by increasing the activation of specific concept index. This illustrates the impact of targeted concept steering on the respective model outputs.
concept Z₁, one can control the output generation of the model. Following Templeton et al. (2024), we multiply 2, of desired concept element by constants ranging from -10 10. To verify whether this steerability works as intended, we steer the activation of the same concept in the pretrained model's SAE latent space and confirm whether the output exhibits the corresponding concept. Here, we use a 386M parameter model trained with CoCoMix, where the pretrained model is GPT-2. As shown in Figure 5, when the concept related to "website address" amplified, both models start generating actual website addresses. This demonstrates that model has successfully learned the GPT-2 aligned concepts. More examples of steering can be found in Appendix B.3.
3.3 ANALYSIS OF COCOMIX'S EFFECTIVENESS
In this section, we provide a detailed analysis of CoCoMix to validate the effect of each proposed component. Unless otherwise specified, we use a 69M model and train on 20B tokens.
Effectiveness of the attribution score. We first analyze whether the attribution score effectively extracts important concepts. To demonstrate this, we train CoCoMix using the activation value Ca for concept extraction, i.e., Lconcape (C:) in Equation 2, instead of the attribution score a, Remark that the activation value also well reflects the importance of the concept (Bricken et al., 2023). As shown in Figure 6a, using attribution scores significant significantly improves performance, improving sample efficiency by 17.5% compared to activation value based selection. We believe it will be interesting future direction to explore other selection criteria for improving CoCoMix's 's performance or removing undesirable concepts reduce bias, e.g., selectively removing unsafe concepts for safe training.
Comparison with direct hidden state predictions. To evaluate the importance of predicting the concept extracted from SAE, we compare CoCoMix with direct hidden state prediction strategies (i.e., predict the full hidden state without projecting into the concept space). To have a comparison under the same architecture as CoCoMix, we replace the concept prediction head M with a two-layer multilayer perceptron (MLP), denoted g(-). which predicts the pretrained LLM's hidden state h con directly from the hidden state of the model The predicted representation, g(h). then compressed into continuous embedding for insertion to have the same architecture CoCoMix. Here, we use continuous loss including, f f2. and the cosine distance (c.g.. |h g for (2) to predict the hidden state. As shown in Figure 6b, direct hidden state prediction leads to a performance drop. We conjecture this to be due to SAE's ability to decompose the latent into semantically meaningful concepts while predicting all hidden states may include noisy components, emphasizing the effectiveness of our method.
Compression layer weight analysis. Now, we analyze the weight of the compression layer W in Equation 3 to show how CoCoMix utilize the predicted concept. To this end, we visualize the f₂ norm of each concept's weights of 386M CoCoMix: for the weight matrix W E RdxC where d is the
hidden dimension and C' is the number of concepts, the norm is defined as: W..ell2 = VV W2,6 As shown in Figure 6c, we found that portion of concept weights are near zero, e.g., 5.8% has a norm less than 1e-2. This indicates that CoCoMix ignores these concepts when compressing them into
7

---

Published as a conference paper at ICLR 2026
Activation
40
Direct hidden pred:
Validatoo Parplexity
Attributionn (Act" Grad)
Volidbiion Perplexiny
pred: fg
36
Direct hidden pred: cos
1500
Concept pred (ours)
Frequency
27
32
1000
26
28
500
25
17.5% less tokers
0.5
Training tokens (B)
Traiing tolons
Weight norm
(a) Effectiveness of the attribution
(b) Concept vs. direct hidden state
(c) Compression layer's weight
(S3M)
Adding
32
NTP + Pause
Validatio Perplexty
Concept pred (53M)
Concept mix (69M)
Validation Parplecky
Mixing (Intereexding
Validation Perplextyy
CoCoMix (ours)
Pause
32
il) CoCoMtx (69Pg)
28
28
28
20
10
Training solens
Training
(B)
Training tokens
(d) Component analysis
(e) Concept conditioning method
(I) Comparison with Pause Token
Figure 6: Analysis of CoCoMix. (a) Effectiveness of the attribution score for selecting concepts. (b) Comparison between concept prediction and direct hidden state prediction (i.e., predicting the hidden state with continuous loss rather than discretizing the hidden state with SAE). (c) The sparsity in the compression weight. (d) Component analysis by analyzing the contribution of concept prediction and mixing. (e) Design choices for concept conditioning by comparing adding the concept vector to the original hidden state and mixing (interlcaving the concept vector with token hidden representation). (f) Comparison between CoCoMix and the Pause token (i.e., adding learnable tokens).
continuous concept if it is not useful. We conjecture this enabled CoCoMix for strong weak-to-strong supervision as learned to ignore ineffective concepts extracted from weak model.
Component Analysis. We analyze the contributions of each component of our method: (a) concept prediction Equation 2, and (b) concept insertion Equation 4. The results in Figure 6d demonstrate that both components are critical for performance improvement. Specifically, applying the concept prediction loss alone yields modest reduction in perplexity. However, incorporating concept insertion alongside prediction enhances the effectiveness of the loss, resulting in further performance gains. This highlights the role of insertion in enabling the model to leverage the pretrained LLM's latent reasoning effectively. Notably, while concept insertion increases parameter count, it has limited impact on performance when used alone, emphasizing the critical role of concept prediction.
Computational cost analysis. One possible concern of CoCoMix is the increased computation due to the expanded token sequence in upper layers. To this end, we report the total FLOPS required for train- ing CoCoMix and baselines, including the NTP and the Pause token. As shown in Table 2, CoCoMix introduces a modest FLOPs increase compared NTP, but remains more efficient than the Pause token while achieving better perplexity in both low-scale and large-scale settings. Furthermore, it is worth not- ing that CoCoMix with 140.7B tokens matches the perplexity of NTP trained with 200B tokens, using 17.8% fewer FLOPs, highlighting its improved efficiency- perfoorancee trade-off.
Table 2:
Computational cost. Training FLOPS (x1018) and validation perplexity (Val. PPL) comparison of NTP, Pause token, and CoCoMix, using a 69M model.
Method
Train Tokens
FLOPPs
Val. PPL
NTP
2.88
25.3
+ Pause token CoCoMix
3.48
25.1
3.37
24.7
NTP
200 B
28.83
22.7
CoCoMix CoCoMix
141 B
23.69
22.7
200 B
33.68
22.0
Continuous concept mixing method. We explored two methods for introducing the continuous concept é into the model's hidden state h2, referred to as concept conditioning. The first method, insertion, insert the continuous concept in front of the token embedding, thus enabling the model to have a mixed input of the token and concept. The other option is intervention, which directly modifies the hidden state by adding the concept vector, i.e h, h, + ét As illustrated in Figure 6e,
8

---

Published as a conference paper at ICLR 2026
both methods enhance pretraining performance (compared NTP in Figure 6d), highlighting the importance of concept conditioning, where the insertion method performed better. By introducing a distinct concept vector, the insertion method enables the model to explicitly recognize and effectively utilize concepts during generation, enhancing its overall performance.
CoCoMix pause tokens. Furthermore, consider additional baseline that jointly uses pause token (Goyal et al., 2024). Specifically, the pause token uses an additional learnable token that is inserted into the token embedding, enabling the LLM think more before predicting the next token, which is similar to our continuous concept insertion. To this end, we insert the pause token for every input token on the same hidden state layer as CoCoMix to ensure comparable computation. Moreover, we also consider training the pause token with KD. As shown in Figure 6f, CoCoMix consistently outperforms the pause token, indicating our inserted (or interleaved) continuous concept indeed consists of useful information to improve the performance.
4 RELATED WORK
Beyond token-level guidance for language modeling, While next token prediction (NTP) remains the standard paradigm, recent works explore guidance beyond language tokens. For instance, some methods explore better target, such leveraging multi-token predictions to capture long context dependencies (Gloeckle et al., 2024; DeepSeek-Al. AI, 2024) or predicting sequence embedding (Lee al., 2024). Additionally, methods explore new types of inputs, e.g., using latents (Hao et al., 2024) or self- gelf-generated thought as inputs (Zelikman et al., 2024), which have shown improving reasoning capabilities. Only recently, concept-level modeling using local encoder-decoder architectures has also been explored to represent a language at a higher abstraction level (L.CM et al., 2024). Other methods add extra tokens in the input space to increase computation at inference time Nye al. (2021); Wei et al. (2022); Goyal et al. (2024); Lanchantin et al. (2024). In contrast, we propose a pretraining approach that integrates NTP with continuous concepts, connecting high-level semantics with fine- ime-grained token guidance.
Sparse Autoencoders (SAEs). SAEs extend the autoencoder by enforeing sparsity constraints in the latent space (Lee et al., 2006). The features learned by SAEs are often interpretable and disentangled, making them useful across various domains, including language modeling (Bricken et al., 2023). Additionally, SAES have gained attention in mechanistic interpretability due to their ability to capture coherent semantic concepts (Marks et al., 2024). This property has enabled practical advancements in identifying and manipulating semantic concepts and facilitating steering for controlled model outputs (Lieberum et al., 2024). Among SAE variants, TopK SAES (Makhzani and Frey, 2014) enforce explicit sparsity using TopK activation function, demonstrating effectiveness even for large models (Gao et al., 2024). In this work, we leverage SAE and, to the best of our knowledge, are the first to apply it to LLM pretraining, achieving strong performance while enhancing the interpretability and controllability of the trained model.
Knowledge distillation (KD). Our method also be related to KD, i.e., transfers the expertise of a teacher model to a student model to enhance performance (Hinton et al., 2015; Zagoruyko and Komodakis, 2017), as CoCoMix extracts high-level semantic features from pretrained model. Recently, KD for LLMs has garnered increasing attention, leveraging knowledge from a teacher to improve the generative and encoding capabilities of student (Sanh et al., 2019; Ko al., 2024). Especially, applying KD to LLM pretraining remains challenging due to the massive token scales (billions to trillions), forcing most current methods to resort to naive token-level probability matching (Team, 2024; Gu et al., 2024). Additionally, while pretrained models contain vast amount of learned information and are thus beneficial to use, reusing knowledge from smaller teacher models remains challenging (Burns et al., 2023). In this work, we show CoCoMix can even leverage the concept extracted from small models to train a large model showing weak- weak-to-strong supervision.
5 CONCLUSION
We propose CoCoMix, a new LLM pretraining framework that augments next token prediction with continuous concepts. By leveraging concepts extracted from pretrained SAE as targets, model predicts both the next token and the associated concept. Predicted concepts are then compressed into continuous concept vector, which is then mixed into the hidden state. This approach enhances
9

---

Published as a conference paper at ICLR 2026
interpretability and controllability by enabling direct probing of the distilled concepts. Experimental results show that CoCoMix consistently improves performance across benchmarks, particularly in challenging generalization scenarios such as weak-to-stee strong supervision.
ETHICS STATEMENT
We introduce novel LLM pretraining framework that augments continuous concepts for the next token prediction. This framework not only improves the performance but also enhances the inter- pretability and controllability of the LLM. By enabling direct probing and analyzing the CoCoMix's s concept space, one can achieve transparency in understanding and guiding model behavior. Addi- tionally, the framework offers the potential to selectively exclude harmful or undesirable concepts during LLM pretraining. contributing to the development of more responsible and ethical Al sys- tems. We believe this research provides new perspectives on language modeling through enhanced interpretability.
REPRODUCIBILITY STATEMENT
All implementation details, including the full of hyperparameters and experimental settings, are described in Section 3 and Appendix A.
REFERENCES
Z. Allen-Zhu. Physics of language models: Part 4.1, architecture design and the magic of canon layers, 2025.
G. Bachmann and V. Nagarajan. The pitfalls of next-token prediction. In International Conference on Machine Learning. 2024.
D. Bachrens, T. Schroeter, S. Harmeling, M. Kawanabe, K. Hansen, and K.-R. Müller. How to explain individual classification decisions. The Journal of Machine Learning Research, 2010.
Y. Bisk, R. Zellers, J. Gao, Y. Choi, al. Piqa: Reasoning about physical commonsense in natural language. In AAAI Conference on Artificial Intelligence, 2020.
T. Bricken, A. Templeton, J. Batson, B. Chen, A. Jermyn, T. Conerly, N. Turner, C. Anil, C. Denison,
A. Askell, R. Lasenby, Y. Wu, S. Kravec, N. Schiefer, T. Maxwell, N. Joseph, Z. HaffieldDDodds A. Tamkin, K. Nguyen, B. McLean, J. E. Burke, T. Hume, S. Carter, T. Henighan, and C. Olah.
Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. http://taansfomme-errritt http://00//mmoooomamiicffeeeeeee
T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, Askell, et al. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
C. Burns, P. Izmailov, J. H. Kirchner, B. Baker, L. Gao, L. Aschenbrenner, Y. Chen, A. Ecoffet, M. Joglekar, J. Leike, et al. Weak-to-strong generalization: Eliciting strong capabilities with weak supervision. arXiv preprint 211:::2312.099900 2023.
M. Chaudhary and A. Geiger. Evaluating open-source sparse autoencoders on disentangling factual knowledge in gpt-2 small. arXiv preprint arXiv:2409.04477, 2024.
Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint AXXi:1803.00499 2018.
H. Cunningham, A. Ewart, I. Riggs, R. Huben, and L. Sharkey. Sparse autoencoders find highly interpretable features in language models. arXiv preprint wrXiv:2309.08600, 2023.
DeepSeek-AL. Deepseek-v3 technical report. arXiv preprint arXiv:2412 /9437, 2024.
10

---

Published as a conference paper at ICLR 2026
Y. Deng, K. Prasad, R. Fernandez, P. Smolensky, V. Chaudhary, and S. Shieber. Implicit chain of thought reasoning via knowledge distillation. arXiv preprint arXiv:311.001660, 2023.
A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Yang. A. Fan, et The llama 3 herd of models. arXiv preprint arXivv2407.21783 2024.
D. Gao, L. Ji, L. Zhou, K. Q. Lin, Chen, Z. Fan, and M. Z. Shou. Assistgpt: A general multi-modal assistant that can plan, execute, inspect, and learn. arXiv preprint arXivv2336.088640, 2023.
L. Gao, D. la Tour, H. Tillman, G. Goh, Troll, A. Radford, Sutskever, J. Leike, and Wu. Scaling and evaluating sparse autoencoders. arXiv preprint arXiv:2406.04493 04090, 2024.
E. Gloeckle, B. Y. Idrissi, B. Rozière, D. Lopez-Paz, and G. Synnaeve. Better & faster large language models via multi-token prediction. In International Conference on Machine Learning. 2024.
S. Goyal, Z. Ji, A. S. Rawat, A. K. Menon, S. Kumar, and V. Nagarajan. Think before you speak: Training language models with pause tokens. In International Conference on Learning Representations, 2024.
Y. Gu, H. Zhou, F. Meng. J. Zhou, and M. Huang. Miniplm: Knowledge distillation for pre-training language models. arXiv preprint arXiv:2410.17215. 2024.
S. Hao, S. Sukhbaatar, D. Su, X. Li, Z. Hu, J. Weston, and Y. Tian. Training large language models to reason in a continuous latent space. arXiv preprint 4TXiv:2412.0676,, 2024.
G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.022 02531, 2015.
Y. Kim and A. M. Rush. Sequence-level knowledge distillation. In Annual Conference of the Association for Computational Linguistics, 2016.
Ko, S. Kim, Chen, and S.-Y. Yun. Distillm: Towards streamlined distillation for large language models. In International Conference on Machine Learning, 2024.
J. Lanchantin, S. Toshhiial, J. Weston, S. Sukhbaatar, al. Learning self-notes. Advances in Neural Information Processing Systems, 36, 2024.
and memorize with
I.M, L. Barrault, P.-A. Duquenne, M. Elbayad, A. Kozhevnikov, B. Alastruey, P. Andrews, M. Coria, G. Couairon, M. Costa-jussà, et al. Large concept models: Language modeling in sentence representation space. arXiv preprint axXiv:2412.08821, 2024.
Y. LeCun. A path towards autonomous machine intelligence. Open Review, 2022.
H. Lee, Battle, R. Raina, and A. Ng. Efficient sparse coding algorithms. Advances in Neural Information Processing Systems, 2006.
H. Lee, D. Kim, J. Jun, S. Joo, Jang, K.-W. On, and M. Seo. Semiparametric token-sequence co-supervision. In Anmual Conference of the Association for Computational Linguistics, 2024.
T. Lieberum, S. Rajamanoharan, A. Conmy, L. Smith, N. Sonnerat, V. Varma, J. Kramár, A. Dragan, R. Shah, and N. Nanda. Gemma scope: Open sparse autoencoders everywhere all at once on gemma 2. arXiv preprint arXiv:2408.00 051407 2024.
I. Loshchilov. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101. 2017.
A. Makhzani and B. Frey. K-sparse autoencoders. In International Conference on Learning Repre- sentations, 2014.
S. Marks, C. Rager, J. Michaud, Belinkov, D. Bau, and Mueller. Sparse feature eir- cuits: Discovering and editing interpretable causal graphs in language models. arXiv preprint arXiv:2403. 19647, 2024.
S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. In International Conference Learning Representations, 2017.
11

---

Published as a conference paper at ICLR 2026
M. Nye, A. J. Andreassen, G. Gur-Ari, H. Michalewski, J. Austin, D. Bieber, D. Dohan, A. Lewkowycz, M. Bosma, D. Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:21 2021.
D. Paperno, G. Kruszewski, A. Lazaridou, Q. N. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The lambada dataset: Word prediction requiring broad discourse context. In Annual Conference of the Association for Computational Linguistics, 2016.
K. Paster, M. D. Santos, Z. Azerbayev, and J. Ba. Openwebmath: An open dataset of high-quality mathematical web text. arXiv preprint orXiv:2310.06766, 2023.
A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever. Improving language understanding by generative pre-training. 2018.
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. 2019.
A. S. Rawat, V. Sadhanala, A. Rostamizadeh, A. Chakrabarti, W. Jitkrittum, V. Feinberg, S. Kim, H. Harutyunyan, N. Saunshi, Z. Nado, et al. A little help goes a long way: Efficient Ilm training by leveraging small Ims. arXiv preprint arXiv:2410. 18779 2024.
B. Roziere, J. Gehring, F. Gloeckle, S. Sootla, Gat, X. E. Tan, Y. Adi, Liu, R. Sauvestre, T. Remez, et al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial winograd schema challenge at scale. In AAAI Conference on Artificial Intelligence, 2020.
V. Sanh, Debut, J. Chaumond, and T. Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. In Advances in Neural Information Processing Systems, 2019.
M. Sap, H. Rashkin, D. Chen, R. Le Bras, and Y. Choi. Social IQa: Commonsense reasoning about social interactions. In Conference on Empirical Methods in Natural Language Processing, 2019.
A. Shrikumar, P. Greenside, A. Shcherbina, and A. Kundaje. Not just black box: Learning important features through propagating activation differences. In International Conference on Machine Learning, 2016.
K. Simonyan and A. Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint arXiv: 1409//5500 1556, 2014.
M. Sundararajan, A. Taly, and Q. Yan. Axiomatic attribution for deep networks. In International Conference on Machine Learning, 2017.
G. Team. Gemma 2: Improving open language models practical size. arXiv preprint 0rXiv:2408.00118, 2024.
A. Templeton, T. Conerly, J. Marcus, J. Lindsey, T. Bricken, B. Chen, A. Pearce, C. Citro, E. Ameisen, A. Jones, H. Cunningham, N. L. Turner, C. McDougall, M. MacDiarmid, C. D. Freeman, T. R. Sumers, E. Rees, J. Batson, A. Jermyn, S. Carter, C. Olah, and T. Henighan. Scaling monosemanticity: Extracting interpretable features from claude 3 sonnet. Trans- former Circuits Thread, 2024. URL https: woweerssoooeeeireeutt wtw.rrcuiss..ub22244 pub/2024/ scaling wwanng-mmoommmmtiiiini.o..
J. Wei, X. Wang, D. Schuurmans, M. Bosma, F. Xia, E. Chi, L., D. Zhou, et al. Chain-of -f-thuughht prompting elicits reasoning large language models. Advances neural information processing systems, 35:24824-2488 2022.
Z. Wu, A. Arora, Z. Wang, A. Geiger, D. Jurafsky, C. D. Manning, and C. Potts. Reft: Representation finetuning for language models. In Advances in Neural Information Processing Systems, 2024.
D. Xuan-Quy L. Ngoc-Bich, P. Xuan-Dung, N. Bac-Bien, and V. The-Duy. Evaluation of chatgpt and microsoft bing ai chat performances on physics exams of vietnamese national high school graduation examination. arXiv preprint arXiv:2306.045388 2023.
12

---

Published as a conference paper at ICLR 2026
S. Yang, E. Gribovskaya, N. Kassner, M. Geva, and S. Riedel. Do large language models latently perform multi-hop reasoning? In Annual Conference of the Association for Computational Linguistics, 2024.
Yu, J. Xu, J. Weston, and Kulikov. Distilling system into system 1. arXiv preprint arXiv:2407.06023, 2024.
Z. Yun, Chen, B. A. Olshausen, and Y. LeCun. Transformer visualization via dictionary learn- ing: contextualized embedding as linear superposition of transformer factors, arXiv preprint 1:X1::2103.1599, 2021.
S. Zagoruyko and N. Komodakis. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. In International Conference on Learning Representations, 2017.
E. Zelikman, G. Harik, Y. Shao, Jayasiri, N. Haber, and N. D. Goodman. Quiet-star: Language models can teach themselves to think before speaking. In Conference on Language Modeling, 2024.
R. Zellers, A. Holtzman, Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can machine really finish your sentence? In Annual Conference of the Association for Computational Linguistics, 2019.
A. Zou, L. Phan, S. Chen, J. Campbell, P. Guo, R. Ren, A. Pan, X. Yin, M. Mazeika, A.-K. Dombrowski, et al. Representation engineering: A top-down approach to ai transparency. arXiv preprint ar arXiv:2310.0 440,, 2023.
13

---

Published as a conference paper at ICLR 2026
A EXPERIMENTAL DETAILS
Architecture details. Our method and baseline both utilize the GPT-2-based Transformer architecture and tokenizer (Radford al., 2019), with a context length of 1024, We consider three model sizes, defined by the number of activated parameters; 69M, 386M, and 1.38B. For CoCoMix, the hidden state dimensions d are 512, 1024, and 2028, with 8, 24, and 24 layers, respectively. The number of attention heads is set to 8, 16, and 16 for the three configurations. Baseline models are configured to match the number of activated parameters by employing the same number of layers and attention heads, with hidden state dimensions of 624, 1072, and 2096, respectively. We leverage an open- source pretrained Topk Sparse Autoencoder (SAE) (Gao et al., 2024) for concept extraction, where Koonsept is set to 32 and the concept activation size is 32,76s 7,76.. Consequently, our models introduce activated parameters on top of the base Transformer parameters. The GPT-2 model (a teacher model for KD and concept extraction model for CoCoMix) has 12 layers, hidden dimension size of 768, and 124M parameters. The concept extraction layer Leon is configured as the 6th middle layer of GPT-2 for all model sizes. For CoCoMix, the concept prediction is done at the 4th layer for the 69M model and the 6th layer for the larger 386M and 1. 1.38B models. The reported training FLOPs in Table 2 are based the base model, and show low FLOPs thanks to the high sparsity of the compression layer (only activated with a concept number of 32), which introduces small cost.
an additional (C - K concept))
Training details. We mainly followed the training details outlined in (Brown et al., 2020). For the main results presented in Section 2, models were trained on 200B tokens over approximately 382K training steps, while other experiments were conducted on 20B tokens. We used OpenWeby Text as the training dataset to match the same training corpus as GPT-2. The learning rate schedule included a warm-up phase over the first 1/300 of the total training steps, followed by cosine decay to 10% of the maximum learning rate at the end of training. The maximum learning rates were set to 6c-4, 3e-4, and 2c-4 for the 69M, 386M, and 1.38B models, respectively. The batch sizes were configured to 0.5M, 0.5M, and 1M tokens for the three model sizes. A weight decay of 0.1 was applied across all configurations, and we utilized the AdamW optimizer (Loshchilov, 2017) with 3 = 0.9 and 32- = 0.95. For training CoCoMix, the concept prediction loss Equation 4 was scaled by hyperparameter x = 0.1, and K was (NTP) loss is our primary objective, we chose À that is smaller than to ensure that the concept prediction loss has a similar scale-would not dominate but instead serve as an auxiliary objective. We found that this value empirically worked well with our small-scale model (69M), and thus we consistently used 0.1 throughout all experiments. While further tuning could potentially yield better results, this initial choice was sufficient to clearly demonstrate our method's effectiveness. For training CoCoMix under a 69M transformer for 20B tokens, we have used 16 A100 GPUs for less than 8 hours.
4. Here, as the next token prediction
For the KD baseline, we employed the vanilla KD loss, where the output probabilities of the teacher model Mee and the student model M were matched using the Kullback-Leible (KL) divergence. Specifically, given an input x, the KD loss is defined as M |x) + AKD KL (e)Meeaal)))))MMMM-a where ABD = 0.1 consistently demonstrated strong performance.
14

---

Published as a conference paper at ICLR 2026
Valdddo hespacty
70-
Cacolla
Pepbesry
55
26-
28.07
20
100
150
200
150
200
150
203
220
(a) OpenWebText, PPL (1) (b) LAMBADA, PPL (L) (c) WikiText-103. PPL (L) (d) HellaS Acc-n (t)
Acoroyy
Acoray
Acarst
100
1.50
200
150
200
130
202
220
(e) PIQA, Acc (1)
(1) SIQA, Acc (t)
(g) Arc-Easy, Ace (1)
(h) WinoGrande, Ace (t)
Figure 7: CoCoMix vs. NTP performance at different training checkpoints on 69M parameter
model. Each model is trained on the 200B tokens sampled from the Open WebText dataset. The plot shows the result of (a) WebTet, (b) LAMBADA, (c) WikiText- 103, (d) HellaSwag, (e)
PIQA, (f) SIQA, (g) Are-Easy, and (h) WinoGrande datasets. We the concepts extracted from 124M-sized model for training CoCoMix.
Wilicce feclesty
NTP
CoCoffee (oun)
35-
16-
Pepporty
Prpposii
78
16
25
100
150
200
100
150
200
100
130
300
100
130
100
Traning tokema (8)
Training soliura (8)
Traning sohara (8)
Training takam (B)
(a) Weexee PPL (1) (b) LAMBADA. PPL (L) (c) WikiText 103, PPL (U) (d) HellaSwag, Acc-n (t)
Kauny NN
knng
Nouncr
100
150
2d0
150
200
100
150
300
150
200
Traning tokera (8)
Traning (8)
Traning sokera (8)
TTainin takems (B)
(c) PIQA, Acc (t)
(f) SIQA, Acc (t)
(g) Are-Easy, Acc (1)
(h) WinoGrande, Ace (t)
Figure 8: CoCoMix vs. NTP performance at different training checkpoints 368M parameter model. Each model is trained on the 200B tokens sampled from the OpenWeb WebText dataset. The plot shows the result of (a) OpenWebTexx (b) LAMBADA, (c) WikiText- 103, (d) HellaSwag, (e)
PIQA, (f) SIQA, (g) Arc-Easy, and (h) WinoGrande datasets. We use the concepts extracted from 124M-sized model for training CoCoMix.
B ADDITIONAL RESULTS
B.1 DETAILED PERFORMANCE DURING TRAINING
In this section, we present the performance tracking during training 200B tokens, including validation perplexity and the perplexity and accuracy of various downstream tasks, including LAM- BADA, WikiText-1033 103, HellaSwag, PIQA. SIQA, Arc-Easy, and WinoGrande datasets. As shown in Figure 7, Figure 8, and Figure 9, we compare CoCoMix with the next token prediction (NTP) across different active parameter sizes: 69M, 386M and 1.38B. In most of the graphs, CoCoMix consistently demonstrates performance gains. Notably, our results show that CoCoMix achieves stable improvements in perplexity across all tasks. Furthermore, CoCoMix exhibits sample efficiency;
15

---

Published as a conference paper at ICLR 2026
NTP Cocoffle (oart)
52-
ACCUniC Fom 750
Pepplaii
22
100
150
200
100
150
200
100
150
200
150
100
Traning tokema 18)
Traning soliam (8)
Traning sohera (8)
Trannng tokom (B)
(a) Texx PPL (1) (b) LAMBADA, PPL (L) (c) WikiText- 103, PPL (U) (d) HellaSwag. Acc-n (1)
Houres CR
Acurssy
100
150
200
100
200
150
150
300
150
200
Taaning tokens (8)
Traning toliens (8)
Traning token (8)
Thaknin tokens (B)
(c) PIQA, Ace (t)
SIQA, Ace (f)
g Are-Easy, Ace (T)
(h) WinoGrande, Ace (t)
Figure 9: CoCoMix vs. NTP performance at different training checkpoints 1.38B parameter model. Each model is trained on the 200B tokens sampled from the OpenWeb WebText dataset. The plot shows the result of (a) Open WexTex, (b) LAMBADA, (c) WikiText-109 103, (d) HellaSwag, (e) PIQA, (f) SIQA, (g) Arc-Easy, and (h) WinoGrande datasets. We use the concepts extracted from 124M-sized model for training CoCoMix.
Table 3: Comparison of CoCoMix with ablation baselines. We report validation perplexity on WebnWe using a 69M-parameter Transformer model.
Method
Validation PPL
NTP (token-only)
25.3
Interleave token twice
25.1
(ii) Insert concept every steps
25.0
(iii) Concept-only (no token hidden states)
25.1
CoCoMix (Ours)
24.7
for instance, in the 1.38B model, CoCoMix reaches the same WpenWebText validation perplexity as NTP while requiring approximately 43B fewer tokens (a 21.5% improvement in token efficiency).
B.2 MORE COMPARISON WITH ADDITIONAL BASELINES
To further demonstrate the effectiveness of CoCoMix, we consider three additional baselines. Specifi- cally, we compare with: (i) duplicating each token hidden state, resulting in interleaved sequence of (h,, h1, h2, h2 ... h, h:): ; (ii) interleaving continuous concepts every T steps instead at every time step, yielding (h,... hy, êr. hy+1. h2r, C.r,,... ); where use T = 4 and (iii) using only continuous concepts without token hidden states, i.e., (θ1, 622 ",,
This comparison isolates key components of CoCoMix and highlights their contributions overall performance. We present the result in Table 3. First, duplicating token hidden states to match the sequence length of CoCoMix yields worse performance despite higher computational cost (it is worth noting that the copied baseline increases the forward GFLOPs even higher than CoCoMix, 164.46 + 169.52 for computing 1024 tokens), indicating that gains are not solely due to sequence length. Second, reducing the interlcaving frequency of concept vectors leads to performance degradation, suggesting that frequent integration of semantic signals is crucial. Lastly, a model using only continuous concepts outperforms token-only models but still underperforms CoCoMix, confirming that the synergy between token-level and concept-level representations essential for optimal effectiveness.
16

---

Published as a conference paper at ICLR 2026
B.3 ADDITIONAL STEERABILITY RESULTS
Irput Prompt: best platform for buying
No Steer
Steer 12783 $ dollar related
Steer 19193: Phone related
GPT2 SAE CoCoMix (ours)
Generation: vernauular ticketing syster, said the offictal. "Tha tacket prices are fixed and the tickets are available at the sane time.
Generations vernacular website, a great resourbe
Generation: ichone.
fanding out about the event. The event is free to attend, but you will need to pay for your ticket. The cost 528 for
great way buy ticketa the gane. it is easy to use. and it is very fast. You can buy ticket to a gane, for
Generation: iphone
Generation: $18.00 price. The price $18 ticket. bought $5.50
Generation: inhons. The iPhone way get
the iPhone. good way people buy seartphone. iPhone lot features, but it's not
The price good ticket sore then
(a) Main figure prompt, '$ dollar' and 'Phone' concept
Input Prompt: Latent
No Stear
Steer concupt 15687: yearlmonth related
Steer concupt 19193: Phone related
GPT2 SAE CoCohlix (ours)
Generation:
technique that allows conplex systen way
Generations
the comrunity the and of year.
powerful data is stored in relational database, and the model is then used to greate a nodel of a par particular type of data
Generation:
the
most the
that is easy to understand easy
For the user to menipulate.
Generation: ixportant part of the The first is design nodel that to
Generation: good idea. It's is way of deing series of month and then you to the level and thang. You two-vear-ele
Ceneration: great to get started with Android device. The first thing your corputer.
to do install the located
easy
worry
(b) New prompt, 'ycar/month' and "Phone" concept
Input Prompt: Latent concept modeling a
No Steer
Steer concapt 6781: politiclla related
Steer concapt 8984: aggressive tone
CoCoMti (ours)
Generation: technique that allows
Generation: hige-level policy: The polscy
Generation: thing? I'm what is
model complex ia easy underatand easy
White House does
<RFNOVED=
user manipulate.
that
not discussed the natter with
GPT2 SAE
Generation: very Lant part of the design process. The first step is to create model that is easy to use and for the understand.
Generation: comnonn practice in the world of computer science. The new laM, which was introduced in 2810, requires that all data public
Generation: thing of the past.. And that s what we re talking about here. concept?
treated
(c) New prompt, 'Politic/law' and 'aggressive tone' concept
More qualitative demonstration of the concept steering effect. CoCoMix and GPT2 models are 350M and 124M parameter transformers, respectively. For CoCoMix, we manipulate the predicted logit z, while for GPT2, we adjust the SAE concept space e by increasing the activation of specific concept index
Figure 10:
To further analyze the steerability enabled by CoCoMix, conducted experiments using both the same prompt as in the main figure and a new prompt (in Figure 10). For consistency, we first applied steering on additional concepts identified during our analysis-"S dollar" and "Phone"-usin the same prompt as in Figure 5. These experiments confirmed that the model could effectively modulate output based on these newly identified concepts, producing coherent and aoleeet-aligned generations. Next, to verify whether the identified concepts generalize to different contexts, we experimented with a new prompt: "Latent concept modeling is a" and steered the model using the previously identified concepts "month/year" and "Phone." The results showed that the model successfully reproduced outputs aligned with these concepts, further supporting the robustness of our method. Additionally, we explored whether new concepts could be identified and steered using the same prompt. In this case, we identified two new concepts: "politics/law" and "aggressive tone." Steering the model with these new concepts demonstrated that the outputs could be effectively controlled to exhibit characteristics aligned with the corresponding concepts. These findings further highlight the flexibility and interpretability of our approach.
17

---

Published as a conference paper at ICLR 2026
C FUTURE WORK AND LIMITATIONS
Several recent approaches, including ours, leverage additional compute to support more advanced forms of reasoning, such test-time optimization latent thinking via pause tokens. While this often leads to improved performance, it naturally increases computational cost. As shown in Section 3.3, CoCoMix achieves favorable compute efficiency, demonstrating stronger performance FLOP compared to prior latent thinking methods. This opens up an interesting direction for future work: designing more informative and structured inputs that better utilize the compute budget of the Transformer to improve the performance. Also, future work could explore learning continuous concepts during pretraining without the need for distillation.
18


Reviews:
[{"review_id": "aPtp6SCjQ4", "reviewer": "Reviewer_SsAp", "summary": "The authors propose a new pre-training paradigm based on \"continuous concepts\".\n\nSpecifically, the authors introduce an additional prediction head in a \"standard\" (transformer) LLM to produce continuous tokens based on \"concepts\". In order to train this additional prediction head the authors propose an additional training objective based on Sparse AutoEncoders (SAEs) and importance attribution using gradients.\n\nThe authors then validate their approach across various model sizes and baselines.", "questions": "Q1: How can you make sure your results are not attributed to randomness in hyper-param tuning? Can you demonstrate anything to this effect?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The main strenghts of the paper:\n1. Great pre-training analysis with a novel architecture.\n2. A working pre-training recipe that seems to improve performance\n3. A way of introducing steerable concepts into the models generation. This can open the door to a lot of interesting research.", "weaknesses": "Main weaknesses:\n1. Model sizes are limited (but understandable).\n2. Hyper-parameter tuning was not discussed in detail and perhaps some of the scores can be attributed to poor hyper-params.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eq2Go4qepp", "reviewer": "Reviewer_XEZ8", "summary": "This work includes a range of experiments that demonstrate the model’s performance across multiple dimensions, such as downstream accuracy, efficiency, and steerability. Overall, the experiments are clearly presented and thoughtfully discussed. The chosen baselines are relevant, though somewhat limited. The analysis of the CocoMix model’s effectiveness is generally solid and supported by clear results.", "questions": "1. Figure 2 PPL scores -- are they computed using NTP objective for both models, or using NTP for the NTP baseline and the CocoMix PPL for the proposed method? If it's the latter, these scores are not really comparable...\n2. For reproducibility reasons, I suggest the authors provide a link to their repo containing the code for the implemented methods and the experiments.\n3. Does CoCoMix scale to larger LLMs?\n4. Are the differences reported in the paper statistically significant?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The CoCoMix model is well described, and the method is easy to follow.\n\n2. The proposal of combining next token prediction with continuous concepts in the pretraining paradigm is novel. This idea of integrating an interpretability mechanism (concept) into pretraining frameworks through SAE comes with significant originality. Such pretraining innovations remain rare in the field, and the model's effectiveness suggests potential incremental impact.\n\n3. This work includes a range of experiments that demonstrate the model’s performance across multiple dimensions, such as downstream accuracy, efficiency, and steerability. Overall, the experiments are clearly presented, and the chosen baselines are relevant, though somewhat limited. The analysis of the CocoMix model’s effectiveness is generally solid.", "weaknesses": "1. Concept Interpretability: This work has been based on the assumption that the latent representation layer in SAE corresponds to human-interpretable concepts. Since this is central to the interpretability claims, more content addressing this assumption would be helpful, beyond what is described in the steerability section. How exactly does CoCoMix capture real, continuous mixture of concepts?\n\n2. Model Design Justifications: Some architectural choices could be better justified if the authors add more explanations or ablation studies. For example, what's the significance of compressing top concepts into a \"continuous\" one, instead of focusing on the top concept?\n\n3. The authors claim that their improvement on downstream tasks shows that CocoMix is better than NTP (Figure 3), but these differences in performance are minimal at best, and I'm not sure if they're significant. Is it worth the overhead?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wNyftpWJhI", "reviewer": "Reviewer_LpV1", "summary": "In this paper, the authors design a new method to improve the next-token prediction and final performance. Specifically, they design CoCoMix, a new pretraining framework for large language models that augments standard next-token prediction with an additional objective: predicting “continuous concepts” derived from a sparse autoencoder and injecting those concepts directly into the model’s hidden states, interleaved with normal token representations. It improves both sample efficiency and downstream performance on language modeling and reasoning tasks. Beyond accuracy, CoCoMix also makes models more interpretable and steerable, because the predicted concepts can be inspected and edited to transparently influence the model’s internal reasoning. The provided figures and examples make this paper easier to understand. Overall, the quality of this paper will be further improved after addressing concerns listed below.", "questions": "[1 ] To be honest, after reading Figure 1 alone or combined with the text in the introduction, the reviewer is still confused about how the extracted concepts benefit the next token prediction. More explanations might be helpful. \n\n[2] During the target concept selection process, it the attribution conducted in each training batch? \n\n[3] Another question is about clarification, after the concept selection, why it is necessary to conduct concept prediction? \n\n[4] In experiments, the model sizes are relatively small. The reviewer understand this might be because the limited budget in model training. Could you discuss whether the findings in this paper could generalize to larger models? \n\n[5] It will be great to disclose the training costs, e.g., GPU numbers, types, GPU hours. \n\n[6] While there are performance gain, will it brought other costs like training time when applying COOMIX?\n\n[7] Sumamrizing the performance gain will improve this paper further. For example, in the abstract, intro, captions of Figure 2 and 3, it will be great to introduce the performance gain. \n\n[8] The steering example is interesting. COuld you introduce how do we know which concept is related to “website” or “money”? The transparency will improve the readability of this paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The topic is highly interesting and might generate broad impact to the LLM community. \n\nI like the visualizations of figures, which are clear and improve the readability of this paper. The reviewer appreciates the authors for doing this.  \n\nIn experiments, the analysis is solid and comprehensive.", "weaknesses": "[1 ] To be honest, after reading Figure 1 alone or combined with the text in the introduction, the reviewer is still confused about how the extracted concepts benefit the next token prediction. More explanations might be helpful. \n\n[2] During the target concept selection process, it the attribution conducted in each training batch? \n\n[3] Another question is about clarification, after the concept selection, why it is necessary to conduct concept prediction? \n\n[4] In experiments, the model sizes are relatively small. The reviewer understand this might be because the limited budget in model training. Could you discuss whether the findings in this paper could generalize to larger models? \n\n[5] It will be great to disclose the training costs, e.g., GPU numbers, types, GPU hours. \n\n[6] While there are performance gain, will it brought other costs like training time when applying COOMIX?\n\n[7] Sumamrizing the performance gain will improve this paper further. For example, in the abstract, intro, captions of Figure 2 and 3, it will be great to introduce the performance gain. \n\n[8] The steering example is interesting. COuld you introduce how do we know which concept is related to “website” or “money”? The transparency will improve the readability of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)