# LLM Pretraining with Continuous Concepts

## Meta Info

URL: <https://arxiv.org/abs/2502.08524>

First author: Jihoon Tack

Affiliation: FAIR at Meta, KAIST, UC San Diego

## Code

<https://github.com/facebookresearch/RAM/tree/main/projects/cocomix>&#x20;

<figure><img src="https://2903491203-files.gitbook.io/~/files/v0/b/gitbook-x-prod.appspot.com/o/spaces%2Flm9XqccbkIHTz4uHcA6m%2Fuploads%2FaM4LH8iWAzizybk7A6XH%2Fimage.png?alt=media&#x26;token=aa2c2f1a-05f7-4001-ab4c-4034e08a66be" alt=""><figcaption><p>Overview Framwork</p></figcaption></figure>

## Understand the paper

### Background

* natural language tokens are often superficial (e.g., function words like “the” or “a”), and need a more efficient way for high-level reasoning and conceptual understanding
* LLMs inherently encode high-level concepts and reasoning processes in their latent representations
* Sparse Autoencoders (SAEs) can effectively isolate meaningful latent features in LLMs by capturing the high-level semantic concepts

### Key Question

Can we augment the next token prediction objective to explicitly model concepts in a latent representation space, thereby bridging semantic abstraction and fine-grained token-level guidance?

### Objection

* Enrich or supplement natural language tokens so that LLMs can learn more abstract reasoning abilities and long-range dependencies.

### Method

* We extract semantic concepts using a pretrained SAE and select the most influential ones based on attribution scores, which quantify each concept’s influence on the model’s output.
* The model is then trained to predict these selected concepts from its hidden state using a cross-entropy loss.&#x20;
* Once multiple concepts are predicted, we compress them into a single continuous concept and mix (or insert) them into the hidden states by interleaving with token embeddings, thereby directly contributing to the next token prediction

### Comment

* While Coconut (Training Large Language Models to Reason in a Continuous Latent Space) directly feeds the hidden layers as ”reasoning thought”, CocoMix (this work) chooses the most significant ”concepts” from the hidden layers.&#x20;
* Still, CocoMix does not emphasize the reasoning chain, but on the extra condition for generating the next token.
* Next step: refine the reasoning chain?
