{"id":"70340b1e-f2a8-439b-8810-740da83b9e93","arxiv_id":"2502.08524","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":5,"one_line_summary":"A language model trained to predict and interleave teacher-derived SAE concepts into its hidden states beats plain next-token prediction and knowledge distillation on several benchmarks.","lead":"CoCoMix is a pretraining method that asks a language model to predict sparse 'concepts' extracted from a pretrained teacher's internal representations, and then interleaves those predicted concept vectors into its hidden states alongside token representations.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 21.5% token-savings headline ignores that CoCoMix appends one concept vector after every token, roughly doubling per-token FLOPs versus NTP; until a FLOPs- or wall-clock-normalized comparison is reported, the sample-efficiency claim is unsubstantiated.","rationale":"The reader's verdict is CONDITIONAL, and I agree with that verdict, but I weight the supporting concerns differently. The reader's designated weakest assumption is the fixed 124M-teacher/SAE concept space (no sensitivity analysis over SAE sparsity, K_attr, or teacher size). I regard that as a robustness gap rather than the most likely point of failure: the paper already provides several independent supports for the concept-supervision mechanism — Fig. 6a (attribution vs. activation selection), Fig. 6b (concepts vs. direct hidden-state prediction), Fig. 6d–e (component and conditioning ablations), and the distribution-shift result in Fig. 4c (OpenWebMath) — so the claim that the teacher's concepts carry useful signal does not rest on a single configuration. The most load-bearing vulnerability is instead the efficiency metric used for the central claim. CoCoMix interleaves one concept vector per token, doubling the sequence length through 18 of 24 layers at the 1.38B scale; per-token FLOPs are therefore near 2× NTP's, so a re-derivation of the efficiency claim on a FLOPs- or wall-clock-normalized axis is needed before 'more sample efficient' (abstract) and '21.5% fewer training tokens' (§3.1, Fig. 2a) can be accepted as stated. My crude reading of Fig. 2 suggests CoCoMix may still be roughly at parity or modestly ahead at equal cumulative compute on some metrics, so the likely outcome is a weakened but not reversed claim; that is exactly why the check is decisive rather than rhetorical. The single-seed issue is real but secondary: effect sizes at 386M/1.38B are large and directionally consistent across sizes and training checkpoints (Figs. 7–9), which makes pure-noise explanations unlikely, though the precise 21.5% number is uncertain. I also note the 'consistently outperforms' wording is too strong given per-task regressions in Table 1 (e.g., PIQA 386M: NTP 64.1 vs. CoCoMix 63.6; Arc-Easy 69M: 39.7 vs. 39.3). The paper does disclose that CoCoMix uses more FLOPs than NTP and includes the pause-token control, which is credit-worthy; the remaining problem is that the headline and the central claim are still framed entirely in tokens. Verdict: UNCHANGED (CONDITIONAL), with the primary condition sharpened to a mandatory compute-normalized efficiency comparison.","tokens_in":18226,"tokens_out":22415,"duration_ms":215224,"concrete_test":"Compute the per-token FLOPs of the 1.38B CoCoMix and NTP configurations from the architecture in Appendix A (interleaving doubles the sequence length through layers 7–24; attention ≈ 4Ld² + 2L²d and MLP ≈ 8Ld² per layer; add the concept head and compression layer, each d·C with C=32,768). Then, using the training curves in Fig. 2, evaluate CoCoMix at the token count at which its cumulative FLOPs equal NTP's cumulative FLOPs at 200B tokens. If CoCoMix's validation perplexity and average downstream accuracy at that point do not match or beat NTP's final values, the 21.5% token-saving claim does not constitute a compute saving and the headline must be revised to a FLOPs-normalized statement.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's flagship efficiency claim (abstract; §3.1; Fig. 2a) is that CoCoMix matches NTP's final OpenWebText validation perplexity at the 1.38B scale using 21.5% fewer training tokens. The comparison is token-based, but CoCoMix does substantially more compute per token, and the paper discloses this only in passing ('CoCoMix utilizes fewer FLOPs than Pause token … but more FLOPs than NTP', p. 5). The reason is architectural: in §2.2 the compressed concept ĉ_t is appended after every token, so the remaining transformer blocks process (h_1, ĉ_1, …, h_t, ĉ_t) — a sequence of length 2T instead of T. At 1.38B the insertion point is layer 6 of 24 (Appendix A), so layers 7–24 process twice as many vectors; since attention is superlinear in sequence length while MLP cost is linear, each of those layers costs roughly 2–2.7× an NTP layer, for a per-token FLOPs multiplier near 2× even after the baselines' larger hidden dimension (d=2096 vs d=2028) is credited. At the claimed crossing point (~157B tokens), CoCoMix has therefore consumed roughly 1.5–1.7× NTP's cumulative compute, so '21.5% fewer tokens' is not evidence of sample efficiency in the standard compute-normalized sense. This is load-bearing because 'more sample efficient' is the abstract's central claim and the quantified headline result. The concern is not that the method is worthless: the pause-token control (Fig. 6f), which pays the same doubled-sequence cost, and the ablations (Fig. 6a–e) support that the concept-supervision content drives the gains. But the NTP comparison must be restated on a compute-normalized axis, and because both curves are single runs, the 21.5% figure itself carries unknown variance.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper proposes CoCoMix, a pretraining framework that augments next-token prediction with a concept-prediction loss derived from a frozen pretrained GPT-2 and a TopK sparse autoencoder (SAE). For each token position, the top K_attr concepts by a gradient-times-activation attribution score are used as discrete labels for a linear concept head on the student; the predicted concepts are then sparsified, compressed into a continuous vector, and interleaved into the student's hidden-state sequence. Experiments on 69M, 386M, and 1.38B models pretrained on OpenWebText report improved validation perplexity and downstream accuracy relative to NTP, knowledge distillation, and pause-token baselines, including a headline claim of matching NTP's final perplexity with 21.5% fewer training tokens at 1.38B. The paper also demonstrates qualitative concept steering and claims weak-to-strong supervision benefits.","tokens_in":18596,"tokens_out":9710,"duration_ms":109825,"significance":"If the results hold under a compute-normalized comparison, CoCoMix is a valuable contribution: it shows that SAE-derived concept targets can provide a training signal beyond token-level loss, and the component ablations in Fig. 6 usefully isolate concept prediction from concept interleaving. The public code link and the pause-token control (which pays a similar doubled-sequence cost) are strengths that partially address the extra-compute confound. However, the headline sample-efficiency claim currently rests on token counts only and on apparently single runs, so the significance is conditional pending FLOP- or wall-clock-normalized evidence and variance reporting.","major_comments":[{"comment":"The central claim that CoCoMix is 'more sample efficient' and reaches NTP's final validation perplexity with 21.5% fewer tokens is based on token counts, not compute. Because the compressed concept vector ĉ_t is interleaved after every token (Eq. 3) and the insertion point for the 1.38B model is layer 6 of 24, layers 7–24 process sequences of length 2T instead of T; each such layer therefore incurs roughly 2× the MLP cost and up to 4× the attention cost, offset only slightly by the smaller hidden dimension (d=2028 vs d=2096 for the baseline). At the claimed crossing point near 157B tokens, CoCoMix has plausibly consumed substantially more cumulative FLOPs than NTP at 200B tokens. The sentence on p.5—'CoCoMix utilizes fewer FLOPs than Pause token … but more FLOPs than NTP'—is the only disclosure and is not quantified. Please add FLOP- or wall-clock-normalized learning curves, or an iso-compute comparison, and state whether the teacher forward/backward pass for computing concept labels is included in the cost; otherwise the efficiency claim should be restated as 'token efficiency at fixed perplexity' rather than sample efficiency.","section":"§3.1, Fig. 2a; §2.2; Appendix A"},{"comment":"No error bars, random seeds, or significance tests are reported, and the headline 1.38B result appears to be a single run. Several reported differences are small and inconsistent across tasks: for example, at 69M in Table 1, CoCoMix is slightly worse than NTP on ARC-Easy (39.3 vs 39.7) and effectively ties on PIQA and HellaSwag; at 386M, PIQA is a tie. Given the language in §3.1 that improvements are 'consistent and significant,' the paper should either report variance over multiple seeds or temper this claim. This is load-bearing because the 'consistently outperforms' statement is a central contribution of the paper.","section":"Table 1, Figs. 2–4"},{"comment":"All experiments use a single fixed teacher/SAE configuration: one open-source TopK SAE with K_concept=32 and C=32768, a 124M GPT-2 teacher, extraction layer L_con=6, K_attr=4, and λ=0.1. No sensitivity analysis is provided over SAE sparsity, K_attr, teacher size, or insertion layer. Since the concept-prediction targets are entirely defined by this teacher and SAE, the reported gains could be idiosyncratic to that particular choice. A sweep over K_attr and SAE sparsity on the 69M/20B setup, plus at least one different teacher size or SAE variant, would substantially strengthen the claim that concept-level supervision—rather than a specific SAE configuration—is responsible for the improvements.","section":"§2.2, Appendix A"}],"minor_comments":[{"comment":"Please clarify the gradient variable in Eq. (1): the notation suggests ∇_{c_t} while the input multiplier is cpre_t, and it is unclear whether gradients through the TopK nonlinearity are included. Also specify the K used in the TopK of Eq. (3) (K_attr? K_concept? something else), and state how positional information is handled for the interleaved concept vectors.","section":"Eq. (1)–(3)"},{"comment":"The 'activated parameter count' matching is not fully defined: please state explicitly which parameters are counted for CoCoMix (concept predictor, compression layer, inserted vectors) and how the baseline hidden dimensions are chosen. The FLOP contribution of the concept-target computation should also be reported, since it is part of CoCoMix's training cost.","section":"Section 3, Appendix A"},{"comment":"The caption says 'We report performance on the OpenWebText (OWT) training set,' but the evaluation section describes validation perplexity and Fig. 2a labels the metric as validation perplexity. Please correct the caption.","section":"Table 1 caption"},{"comment":"For the direct hidden-state prediction baselines, please specify the compression mechanism used to turn the predicted full hidden state into the inserted continuous embedding; the caption currently says only that the predicted representation is 'compressed into a continuous embedding,' which is not enough to ensure the comparison is architectural.","section":"Fig. 6b"},{"comment":"The text in several figures appears as garbled glyphs in the submitted version. Please ensure that fonts and special characters render correctly in the final PDF.","section":"Figures 1, 5, 10"}],"recommendation":"major_revision","confidential_remarks":"The paper is timely and the code release is a strength. The main risk is not circularity but an overclaimed efficiency measure: the 21.5% token-savings headline needs to be reframed or supplemented with a compute-normalized comparison. The pause-token control and component ablations suggest the method has real substance, so I would not reject on the current evidence; a major revision with FLOP-normalized curves, variance reporting, and sensitivity analyses would address the load-bearing concerns."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"First, the thing you should know: the 21.5% token-savings claim is not a sample-efficiency claim as stated. CoCoMix inserts a concept vector after every token, so the upper transformer layers process a sequence of length 2T instead of T. The paper acknowledges this only in passing (\"more FLOPs than NTP\"), and your stress-test estimate of roughly 1.5–1.7x total compute at the crossing point looks right. So the abstract overstates the practical benefit.\n\nWhat's genuinely new: the combination of SAE concept extraction, attribution-based selection, and interleaving predicted concept vectors as a pretraining objective. The ablations in Fig. 6 are well-designed and support the claim that the concept content, not just the extra sequence slots, drives the gains: attribution beats activation, concept prediction beats direct hidden-state prediction, and CoCoMix beats a pause-token baseline that pays the same doubled-sequence cost. Code is available, which is a real plus. The weak-to-strong supervision result is also interesting.\n\nSoft spots: the missing compute-normalized comparison is the big one. The authors need to report either FLOPs or wall-clock at equal perplexity, or at least show the token-efficiency curve against a compute x-axis. Second, there are no error bars or multiple seeds anywhere; the 21.5% figure is a single run. Third, the concept space is entirely a construct of the teacher's frozen SAE, and the paper provides no sensitivity analysis over SAE sparsity, K_attr, teacher size, or insertion layer. Minor: \"consistently outperforms\" is an overclaim—some individual tasks show small degradations (e.g., SIQA at 69M).\n\nWho it's for: people working on pretraining objectives, distillation, or SAE-based interpretability. It's a serious paper with a clear method and believable ablations, but the main efficiency claim needs recomputation on a fair axis. I'd send it to peer review with a request for major revision: compute-normalized comparison, uncertainty quantification, and hyperparameter sensitivity.","headline":"Clever pretraining idea, but the token-efficiency headline ignores a ~2x compute overhead, so the paper needs a compute-normalized comparison before the main claim holds.","tokens_in":19231,"tokens_out":2448,"would_cite":true,"duration_ms":24160,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"The paper proposes Continuous Concept Mixing (CoCoMix), which augments next-token prediction with continuous concept prediction and shows this improves sample efficiency and downstream performance.","keywords":["language model pretraining","next-token prediction","sparse autoencoders","concept prediction","knowledge distillation","weak-to-strong supervision","model interpretability","sample efficiency"],"falsifier":"A direct falsifier: train CoCoMix with the teacher's concept indices replaced by random indices at each position, keeping everything else identical; if the validation perplexity and downstream accuracy remain at the reported levels, then the claimed benefit does not come from semantic concepts. A cheaper version is the paper's own attribution ablation: when top-4 concepts are chosen by raw activation rather than attribution, sample efficiency drops by 17.5%, but the random-selection condition is not reported, so it would settle whether any concept signal — or just any extra interleaved vector — is doing the work.","tokens_in":17997,"feed_emoji":"🧠","tokens_out":6453,"duration_ms":63164,"temperature":0.7,"pith_summary":"The paper asks whether next-token prediction alone is the right training signal for language models, and answers with a pretraining method that adds a second, semantic target: predicting \"continuous concepts\" distilled from a pretrained teacher. The concepts are extracted by a sparse autoencoder from the middle-layer hidden state of a 124M GPT-2 model; at each position the student must predict the teacher's four most output-influential concept indices. The predicted concepts are compressed into one vector and interleaved into the student's hidden state, so the model reads both a token and a concept at every step. The central empirical claim is that this makes pretraining more sample-efficient and more accurate: a 1.38B model reaches the same OpenWebText validation perplexity as standard next-token training with 21.5% fewer tokens, and downstream accuracy improves across seven reasoning benchmarks. A sympathetic reading is that concept-level supervision works especially well in weak-to-strong settings, where a smaller teacher guides a larger student better than probability-matching distillation.","feed_headline":"Predicting concepts while predicting words cuts training tokens 21.5%","feed_subtitle":"A 1.38B model matches standard perplexity with 21.5% fewer tokens and beats distillation on reasoning benchmarks.","key_machinery":"The load-bearing object is the continuous concept vector $\\hat{c}_t = W\\,\\mathrm{TopK}(z_t) + b$, where $z_t$ is the student's concept logit at position $t$, $\\mathrm{TopK}$ keeps the largest entries, and $W,b$ are learned. It is produced from a linear prediction head trained with $\\mathcal{L}_{\\mathrm{concept}}$, a cross-entropy loss whose labels are the teacher's top-4 concept indices selected by the attribution score $a_t = c^{\\mathrm{pre}}_t \\odot \\nabla_{c_t}(-\\log f_{\\mathrm{con}}(x_{t+1}\\mid D(c_t), h_{<t}))$. That formula is what lets the method prefer concepts that causally influence the teacher's next-token prediction, rather than merely the largest activations. The concept vector is then interleaved as a separate unit between token hidden states, so downstream transformer layers can use it as an additional input; the paper's ablations show this interleaving, not just the prediction loss, is required for the reported gains.","core_discovery":"On the paper's own terms, CoCoMix's discovery is that the discrete token stream is not the only useful supervision channel in pretraining: a model trained to also predict sparse-autoencoder concepts — selected by a gradient-times-activation attribution score rather than by raw activation strength — learns faster and better. The mechanism is end-to-end: the student's hidden state at a chosen layer is fed to a linear head that outputs concept logits, a cross-entropy loss pushes those logits toward the teacher's top-4 concepts, and a learnable compression layer turns the top-k predicted logits into a continuous vector interleaved with the token hidden states. Ablations identify both pieces as necessary: concept prediction alone gives a modest gain, interleaving alone is weak, and the combination is what produces the reported improvements. The same recipe also outperforms knowledge distillation and pause-token training, and the predicted concept logits can be amplified during generation to steer the model's output, which the authors present as evidence that the model has internalized the teacher's concept space.","pith_inferences":["An extension the paper leaves to future work is learning the concept space online during pretraining instead of fixing it to a pretrained teacher's sparse autoencoder; if that works it would remove the distillation dependency, but it would need a stabilization mechanism for the concept targets.","The 17.5% sample-efficiency gain of attribution over raw activation suggests that selection quality matters more than concept coverage; a natural test is whether cheaper selection rules, such as random top-4 indices or frequency-based concepts, retain part of the gain.","Because the method interleaves extra vectors into the hidden state, it is close in spirit to pause-token training; a natural next test is whether CoCoMix's advantage persists when the pause token is given the same compute or when the concept vector is generated by the student itself."],"forward_implications":["Standard next-token prediction is not the only effective training objective: adding a continuous concept-prediction loss lowers validation perplexity at matched token budgets.","Weak-to-strong supervision works better with concept targets than with output-probability distillation, so a smaller teacher can guide a larger student.","The trained model can be steered at inference time by amplifying specific predicted concept logits, which produces outputs aligned with those concepts.","The gains grow with training length and survive distribution shift to a math corpus, where knowledge distillation falls behind standard training."],"supporting_citations":[{"why":"Supplies the pretrained open-source TopK sparse autoencoder used for concept extraction and the scaling analysis that motivates it.","marker":"Gao et al., 2024"},{"why":"Defines the TopK activation and k-sparse autoencoder construction that CoCoMix adopts.","marker":"Makhzani and Frey, 2014"},{"why":"Provides the GPT-2 teacher model, the OpenWebText pretraining corpus, and the tokenizer/architecture used throughout.","marker":"Radford et al., 2019"},{"why":"Defines the knowledge-distillation loss that is the paper's main comparison baseline.","marker":"Hinton et al., 2015"},{"why":"Supplies the gradient-times-input attribution idea adapted in the concept selection score.","marker":"Simonyan and Zisserman, 2014"},{"why":"Provides the element-wise gradient-input attribution formulation used to compute attribution scores.","marker":"Shrikumar et al., 2016"},{"why":"Defines the pause-token training baseline that CoCoMix is compared against.","marker":"Goyal et al., 2024"},{"why":"Motivates interpreting sparse-autoencoder dimensions as semantic concepts and supplies the concept-steering procedure used in the interpretability experiments.","marker":"Templeton et al., 2024"}],"fun_headline_variants":["CoCoMix: predicting concepts cuts training tokens 21.5%","Mix in continuous concepts for faster, better LLM pretraining","Concept prediction boosts sample efficiency in LLM pretraining","Steerable LLM training via continuous concept mixing","Continuous concepts beat next-token-only pretraining"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing assumption is that the fixed sparse-autoencoder concept space of a 124M GPT-2 teacher, together with the top-4 gradient-times-activation attribution picks, provides useful supervision for students of a different size and even for out-of-corpus data like OpenWebMath.","fun_headline_variants_meta":{"raw":{"variants":["CoCoMix: predicting concepts cuts training tokens 21.5%","Mix in continuous concepts for faster, better LLM pretraining","Concept prediction boosts sample efficiency in LLM pretraining","Steerable LLM training via continuous concept mixing","Continuous concepts beat next-token-only pretraining"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000197,"raw_usage":{"total_tokens":1351,"prompt_tokens":916,"completion_tokens":435,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":532,"completion_tokens_details":{"reasoning_tokens":354}},"tokens_in":532,"tokens_out":435,"duration_ms":4542,"temperature":1.0,"reasoning_tokens":354,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-08T04:46:37.074566+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"A direct falsifier: train CoCoMix with the teacher's concept indices replaced by random indices at each position, keeping everything else identical; if the validation perplexity and downstream accuracy remain at the reported levels, then the claimed benefit does not come from semantic concepts. A cheaper version is the paper's own attribution ablation: when top-4 concepts are chosen by raw activation rather than attribution, sample efficiency drops by 17.5%, but the random-selection condition is not reported, so it would settle whether any concept signal — or just any extra interleaved vector — is doing the work.","supporting_citations":[],"review_version":1}