SFTProbe가 추론 데이터 품질을 손실 기하로 예측해 큐레이션을 단축한다

무엇인가

이 논문은 추론 LLM을 위한 SFT 데이터 큐레이션 문제를 다룬다. 추론 트레이스는 수천에서 수십만 토큰에 이르고 실패·재검토·우회 시도가 섞여 있어, 기존 instruction tuning용 난이도·다양성 필터(퍼플렉시티, 학습 가능성, 임베딩 유사도 등)를 그대로 쓰기 어렵다. 그래서 현장에서는 여러 교사 모델로 대량 생성한 뒤 LLM 심사로 난이도와 다양성을 거르는 generate-then-filter 파이프라인을 쓰는데, 이 과정이 수 주 걸린다. 저자들은 전체 생성과 필터링을 거치지 않고 추론 데이터 품질을 직접 예측할 수 있는지 묻는다.

어떻게 동작하나

SFTProbe의 핵심은 사전학습 LLM이 낮은 랭크와 비등방성 곡률을 가진 손실 분지 안에 있다는 기하학적 관점이다. 먼저 아주 작은 데이터 부분집합에 LoRA를 짧게 적응시켜 학습이 주로 일어나는 저차원 부분공간을 찾고, LoRA 어댑터의 QR 분해로 직교 기저 방향들을 얻는다. 각 방향으로 손실이 매끄럽게 증가하다가 깨지는 지점까지 조금씩 밀어 보고, 손실 프로파일에 큐빅 스플라인을 적합해 곡률 부호가 바뀌는 변곡점을 찾는다. 이렇게 하면 방향마다 두 개의 체크포인트가 생기는데, 논문은 사전학습 지점과 변곡점에서의 손실 차이가 SFT 중 전체 파라미터 그래디언트 차이를 타이트하게 상한 지을 수 있음을 정리로 보인다. 따라서 손실 서명을 클러스터링하면 전체 학습을 하지 않고도 그래디언트가 비슷한, 즉 중복되는 예제를 묶을 수 있다는 것이 다양성 판단의 근거다.

무엇과 다른가

구체적으로 각 예제의 손실 서명은 추론 트레이스 전체가 아니라 약 5%에 해당하는 초기 1k 토큰만을, 방향별 두 체크포인트를 합친 2r개 체크포인트에서 평가해 만든다. 이 서명을 표준화한 뒤 FAISS로 k개 클러스터로 묶고, 각 클러스터에서 같은 수를 뽑아 다양성을 맞춘다. 난이도는 임의의 한 방향으로 변곡점까지 섭동한 체크포인트에서 초기 100토큰 손실로 점수화한다. 어려운 문제일수록 초기 손실 곡률이 크다는 관찰에 기반한다. 소스별로 난이도 점수를 두 클러스터로 나눠 쉬운 쪽을 버리고, 평균 난이도가 높은 소스에서 더 많이 샘플링하도록 softmax 비율로 예산을 배분한다. 교사 선택도 별도 학습 없이, 후보 교사들이 문제를 100토큰 이내로 바꿔 쓴 표현을 사전학습 학생 모델에서 평가해 평균 재표현 손실이 가장 낮은 교사를 고르는 방식이다.

어떻게 쓰나

의료 추론 실험에서는 M23k 데이터셋에서 1k개를 선택해 Llama3.1-8B-Instruct로 학습했을 때 최고 베이스라인 대비 1.1% 성능 향상을 보고했다. 또 37,698개의 오염 제거된 질문에서 교사를 QwQ-32B로 고르고 1,000개 추론 예제를 생성해, SFTProbe로 큐레이션한 1k 데이터가 전체 큐레이션된 M23k와 0.6% 이내로 맞먹었다고 한다. 이는 데이터 선택 상황과 새 데이터 생성 상황 모두를 보여주는 실험이다.

전제와 한계

수학 추론에서는 NuminaMath의 192k 질문 부분집합에서 교사를 고르고 추론 트레이스를 생성해 OpenThinker-7B와 OpenThoughts-Math 기반 선택 방법들과 비교했다. 원문은 최고 베이스라인 대비 3.4% 성능 향상을 보고한다. 더 눈에 띄는 수치는 5k 생성 예제가 89k 예제로 구성된 OpenThoughts-Math와 맞먹고, 7.5k 생성 예제는 그보다 정확도 2.5% 높다는 점이다. 같은 수학 성능을 OpenThinker-7B 대비 53배 적은 연산으로 달성했고, 긴 추론 트레이스가 이미 생성된 선택 방법들과 비교해도 실행 시간이 더 적게 든다고 주장한다.

물리 추론에서는 1967년부터 2018년까지의 물리 올림피아드 문제와 정답을 정제해 3,084개 문제-해답 쌍을 만들고, PhysicsEval 문제와 합쳐 40k 문제 규모의 혼합 데이터셋을 구성했다. 이로부터 고품질 물리 추론 SFT 데이터셋을 만들어 학습한 모델이 IPhO 2025 은메달 수준인 PRIME-RL/P1-30B-A3B에 근접했다고 실험 섹션은 서술한다. 초록은 은메달 모델을 능가한다고 표현하지만, 본문 Table 1 설명은 ‘근접’에 가깝다. 절제 실험에서는 LoRA 랭크를 4에서 2로 줄여도 성능 차이가 유의하지 않았고, LoRA 학습은 100~400회 반복으로 충분했다. 다양성 필터에서 변곡점 대신 학습 궤적 체크포인트를 쓰면 0.7% 하락, 짧은 LoRA 궤적 체크포인트를 쓰면 추가로 0.8% 하락해 변곡점 기반 손실 서명의 기여를 보였다.

개발자 관점에서 SFTProbe는 추론 모델용 SFT 데이터를 만들 때 대규모 생성-필터 파이프라인 비용을 줄이는 후보 도구다. 특히 교사 모델 후보가 여러 개 있고, 도메인 데이터가 의료·수학·물리처럼 정답 검증이 가능하며, 데이터 예산이 제한된 상황에서 손실 기반으로 난이도와 다양성을 빠르게 점수화하는 용도로 쓸 수 있다. 다만 원문이 깔고 있는 전제는 손실이 매끄럽고 곡률이 유계이며, SFT 그래디언트가 저차원 부분공간에 주로 갇혀 있고, 근포물선 손실 분지가 성립한다는 것이다. 품질 추정에 전체 트레이스가 아니라 초기 1k 토큰과 100토큰 수준의 짧은 구간만 쓰므로, 긴 추론 후반의 오류나 장기 호라이즌 특성은 반영되지 않을 수 있다. 또한 완전한 생성 없는 방법은 아니어서 교사 모델의 짧은 재표현과 손실 평가용 프리픽스 생성이 필요하고, 최종 선택된 예제는 결국 생성해야 하며, 교사 선택은 GPU 메모리에 올릴 수 있는 후보로 제한된다. 원문은 별도의 한계 절을 길게 두지 않지만, 이런 전제와 실험 범위가 실무 적용 전 확인해야 할 지점이다.