Research map / Large language models / Language, safety and interpretability

Task-Adaptive Pretrained Language Models via Clustered-Importance Sampling

ICLR 2025 · Large language models › Language, safety and interpretability · approach: LLM pruning and scaling data · 0 citations (OpenAlex)

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.

Builds task-specialist LMs by resampling a large generalist web corpus (RedPajama-v2) using cluster histograms estimated from a small specialist dataset (CRISP, clustered importance sampling), compared with classifier, DoGE, and cross-entropy difference selection.

Problem

Training specialist language models when only limited in-domain data is available, by adjusting the training distribution of a large generalist dataset.

Approach

Cluster the generalist dataset with hierarchical k-means on SBERT embeddings (up to 16.7M clusters), estimate the cluster histogram on the small specialist set, and sample pretraining batches by drawing clusters from that histogram and then examples uniformly within the cluster, scaling importance sampling to millions of clusters.

Key results

• CRISP with 262k clusters on 1.3B models reaches 71.2% on ARC-E, 38.6% on ARC-C and 33.4% on MMLU versus 58.4, 27.4 and 30.0 for the base RedPajama-v2 model.
• CRISP 1.3B on ARC improves over base with about a 30x training speedup, since the 1.3B CRISP model outperforms the 6.7B base model on ARC.
• CRISP PubMed perplexity is better than base on 97.3% of PubMed test documents while being worse on 95.9% of RedPajama-v2 documents.

Measured effects reported in the paper

Surprising findings

• Importance sampling with CRISP was ineffective after 256k steps when the generalist set was restricted to 120B tokens, due to high repetition (22.5 on average).
• Multi-task CRISP pretraining slightly outperformed targeting RWDB-R directly on RWDB-R accuracy.

Limitations

• Relies on the assumption that generalist and specialist per-cluster losses are equal, which can break when the embedding does not reflect loss-gradient similarity.
• Performance degrades for small specialist sets because cluster-frequency estimates are poor, and for small generalist sets because of repetition.
• MCQ fine-tuning on small train sets gives limited benefit, and the MCQ training loss is not closely tied to the accuracy metric.

Experimental setup

Datasets: RedPajama-v2 The Pile PubMed Central The Pile StackExchange The Pile Wikipedia ARC MMLU RewardBench Reasoning

Benchmarks: ARC-Easy ARC-Challenge MMLU RewardBench Reasoning PubMed Central perplexity StackExchange perplexity Wikipedia perplexity

Metrics: perplexity accuracy

Baselines: Base RedPajama-v2 pretraining logistic regression classifier DoGE cross-entropy difference

Builds on: DoGE importance sampling for continued pretraining (Xie et al.) SBERT k-means++

Closest papers by meaning

• Supervised Knowledge Makes Large Language Models Better In-context Learners (ICLR 2024 · similarity 0.88)
• Generalization v.s. Memorization: Tracing Language Models’ Capabilities Back to Pretraining Data (ICLR 2025 · similarity 0.87)
• Get more for less: Principled Data Selection for Warming Up Fine-Tuning in LLMs (ICLR 2024 · similarity 0.87)
• Dynamic Loss-Based Sample Reweighting for Improved Large Language Model Pretraining (ICLR 2025 · similarity 0.87)
• Two-stage LLM Fine-tuning with Less Specialization and More Generalization (ICLR 2024 · similarity 0.87)
• Enhancing Large Language Model Performance with Gradient-Based Parameter Selection (AAAI 2025 · similarity 0.87)