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
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.
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.
Still, CocoMix does not emphasize the reasoning chain, but on the extra condition for generating the next token.
Next step: refine the reasoning chain?
Last updated