학습 중인 LLM의 다운스트림 성능을 빠르고 정확하게 평가하는 방법

Fast and Accurate Probing of In-Training LLMs' Downstream Performances

대규모 언어 모델(LLM)의 파라미터 크기와 테스트 시간을 확장하는 방식은 AI의 능력을 향상시켰지만, 기존의 생성 기반 평가 방식은 비용이 너무 많이 들고, LLM의 학습 중인 다운스트림 성능 평가의 지연 시간이 감당하기 어려울 정도로 커졌습니다. 하지만 학습 손실(퍼플렉시티)과 같은 간단한 지표가 항상 다운스트림 성능과 상관관계를 가지는 것은 아니며, 때로는 이러한 지표들의 추세가 실제 작업 결과와 달라지는 경우가 있습니다. 이러한 딜레마는 모델의 능력을 측정하는 데 있어 계산적으로 효율적이고 충분히 정확한 방법이 필요함을 시사합니다. 이러한 문제를 해결하기 위해, 우리는 다운스트림 성능을 모니터링하기 위한 경량화된 프로브를 사용하는 새로운 학습 중 평가 패러다임을 제안합니다. 이 프로브는 LLM 체크포인트(학습 중)의 내부 표현을 입력으로 받아, 성공 확률(예: pass@1)로 측정되는 다운스트림 작업에서의 체크포인트 성능을 직접 예측합니다. 우리는 다양한 프로브 아키텍처를 설계하고, OLMo3-7B의 체크포인트를 사용하여 다양한 다운스트림 작업에서 이들의 효과를 검증했습니다. 이 프로브는 체크포인트의 성능을 정확하게 예측할 수 있으며(평균 AUROC > 0.75), 체크포인트 간의 일반화 성능이 우수합니다(초기 체크포인트가 후기 체크포인트의 성능을 예측). 또한, 기존의 생성 기반 평가 방법으로 소요되는 약 1시간의 계산 시간을 약 3분으로 줄입니다. 요약하자면, 본 연구는 실용적이고 확장 가능한 학습 중 다운스트림 평가 패러다임을 제시하며, 이를 통해 LLM 개발 프로세스를 더욱 민첩하고 정보에 입각하며 효율적으로 만들 수 있습니다.

The paradigm of scaling Large Language Models (LLMs) in both parameter size and test time has pushed the boundaries of AI capabilities, but at the cost of making the traditional generative evaluation paradigm prohibitively expensive, therefore making the latency of LLM's in-training downstream performance evaluation unbearable. However, simple metrics like training loss (perplexity) are not always correlated with downstream performance, as sometimes their trends diverge from the actual task outcomes. This dilemma calls for a method that is computationally efficient and sufficiently accurate in measuring model capabilities. To address this challenge, we introduce a new in-training evaluation paradigm that uses a lightweight probe for monitoring downstream performance. The probes take the internal representations of LLM checkpoints (during training) as input and directly predict the checkpoint's performance on downstream tasks measured by success probability (i.e., pass@1). We design several probe architectures, validating their effectiveness using the OLMo3-7B's checkpoints across a diverse set of downstream tasks. The probes can accurately predict a checkpoint's performance (with avg. AUROC$>$0.75), have decent generalizability across checkpoints (earlier predicts later), and reduce the computation latency from $\sim$1 hr (using conventional generative evaluation method) to $\sim$3 min. In sum, this work presents a practical and scalable in-training downstream evaluation paradigm, enabling a more agile, informed, and efficient LLM development process.