Research map / Large language models / Language, safety and interpretability
SimReg: Achieving Higher Performance in the Pretraining via Embedding Similarity Regularization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes SimReg, a token-embedding similarity regularizer (cosine contrastive loss over same-label vs different-label tokens within a sequence) added to next-token pretraining, reporting faster convergence and small zero-shot gains for LLaMA-style dense and Mixtral MoE models.
Problem
Next-token pretraining yields context-dependent token embeddings with high intra-class variance and inter-class similarity, so cross-entropy alone does not sufficiently separate token representations, limiting representation learning efficiency.
Approach
Adds a contrastive embedding-similarity term on the final-layer token embeddings, treating same-token-ID positions in a sequence as positives (plus self-similarity) and other tokens as negatives, with cosine similarity, a temperature, softplus scaling, and chunk-wise computation for sequence parallelism; a margin-based analysis links it to lower cross-entropy.
Key results
- SimReg improves zero-shot average accuracy by +1.04 on LLaMA-350M and +1.14 on LLaMA-1.3B over the cross-entropy baseline (per Table 1).
- Reports over 30% training convergence speedup (about 33% on 1.3B; over 37% on 3B and 7B at 52B tokens) with final loss reduced by about 0.03-0.05.
- Chunk-wise SimReg adds under 2% runtime and under 1% memory overhead on a 7B model on H800 GPUs.
Surprising findings
- Intermediate-layer supervision gives negligible gains; applying SimReg only at the last layer is sufficient, per the authors.
- Chunk-wise SimReg matches or even slightly outperforms full SimReg on the 7B model.
Limitations
- Hyperparameters (tau, lambda) require grid search and a heuristic scaling rule with model dimension; only validated on a few sizes.
- Gains on downstream zero-shot tasks are modest (about 1%) and inconsistent across individual tasks (e.g., some tasks drop).
Experimental setup
Datasets: C4
Benchmarks: ARC-Easy ARC-Challenge BoolQ HellaSwag OpenBookQA PIQA MMLU WinoGrande SciQ
Metrics: zero-shot accuracy average accuracy training loss validation perplexity time per iteration memory
Baselines: SimCSE CPretrain CReg SCE
Builds on: SimCSE supervised contrastive learning AdaGC
Closest papers by meaning
- Analyzing Similarity Metrics for Data Selection for Language Model Pretraining (NeurIPS 2025 · similarity 0.90)
- DEPT: Decoupled Embeddings for Pre-training Language Models (ICLR 2025 · similarity 0.87)
- Repurposing Language Models into Embedding Models: Finding the Compute-Optimal Recipe (NeurIPS 2024 · similarity 0.87)
- Data-Constrained Language Model Pretraining: Improved Regularization and Scaling Laws (NeurIPS 2026 · similarity 0.87)
- EmbedLLM: Learning Compact Representations of Large Language Models (ICLR 2025 · similarity 0.87)
- Model-Aware Tokenizer Transfer (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
