AttentionInfluence: 약한 데이터에서 강한 데이터로의 사전 학습 데이터 선택을 위한 어텐션 헤드 영향력 활용

AttentionInfluence: Adopting Attention Head Influence for Weak-to-Strong Pretraining Data Selection

May 12, 2025

저자: Kai Hua, Steven Wu, Ge Zhang, Ke Shen

cs.AI

초록

최근, LLM의 복잡한 추론 능력을 향상시키기 위해 추론 집약적인 사전 학습 데이터를 수집하는 데 대한 관심이 증가하고 있다. 기존 접근 방식은 일반적으로 이러한 데이터를 식별하기 위해 지도 학습 분류기에 의존하며, 이는 인간이나 LLM에 의한 라벨링을 필요로 하여 종종 도메인 특정 편향을 유발한다. 어텐션 헤드가 문맥 내 추론에 중요한 역할을 한다는 점에 착안하여, 우리는 지도 신호 없이도 간단하면서도 효과적인 AttentionInfluence라는 방법을 제안한다. 우리의 접근 방식은 간단한 어텐션 헤드 마스킹 작업을 통해 소규모 사전 학습 언어 모델이 강력한 데이터 선택자로 작동할 수 있게 한다. 구체적으로, 우리는 검색 헤드를 식별하고 이러한 헤드를 마스킹할 때의 손실 차이를 계산한다. 우리는 AttentionInfluence를 1.3B 파라미터의 밀집 모델에 적용하여 241B 토큰으로 구성된 SmolLM 코퍼스에 대한 데이터 선택을 수행하고, SmolLM 코퍼스와 73B 토큰으로 구성된 선택된 하위 집합을 혼합하여 1T 학습 토큰과 WSD 학습률 스케줄링을 사용하여 7B 파라미터의 밀집 모델을 사전 학습한다. 우리의 실험 결과는 여러 지식 집약적이고 추론이 많은 벤치마크(즉, MMLU, MMLU-Pro, AGIEval-en, GSM8K, HumanEval)에서 1.4pp에서 3.5pp에 이르는 상당한 개선을 보여준다. 이는 작은 모델이 더 큰 모델의 최종 성능을 향상시키는 효과적인 약한-강한 스케일링 특성을 입증하며, 추론 중심 데이터 선택을 위한 유망하고 확장 가능한 경로를 제시한다.

English

Recently, there has been growing interest in collecting reasoning-intensive pretraining data to improve LLMs' complex reasoning ability. Prior approaches typically rely on supervised classifiers to identify such data, which requires labeling by humans or LLMs, often introducing domain-specific biases. Due to the attention heads being crucial to in-context reasoning, we propose AttentionInfluence, a simple yet effective, training-free method without supervision signal. Our approach enables a small pretrained language model to act as a strong data selector through a simple attention head masking operation. Specifically, we identify retrieval heads and compute the loss difference when masking these heads. We apply AttentionInfluence to a 1.3B-parameter dense model to conduct data selection on the SmolLM corpus of 241B tokens, and mix the SmolLM corpus with the selected subset comprising 73B tokens to pretrain a 7B-parameter dense model using 1T training tokens and WSD learning rate scheduling. Our experimental results demonstrate substantial improvements, ranging from 1.4pp to 3.5pp, across several knowledge-intensive and reasoning-heavy benchmarks (i.e., MMLU, MMLU-Pro, AGIEval-en, GSM8K, and HumanEval). This demonstrates an effective weak-to-strong scaling property, with small models improving the final performance of larger models-offering a promising and scalable path for reasoning-centric data selection.