Step-TP: LLM 유도 텐서 프로그램 최적화를 위한 Chain-of-Thought 추론을 포함한 근거 기반 단계별 데이터셋

대규모 언어 모델 (LLMs)의 강력한 추론 능력에도 불구하고, 텐서 프로그램 (tensor programs)의 실행 효율성을 최적화하는 것은 정밀하고 결합 가능한 변환 결정 (transformation decisions)이 필요하기 때문에 여전히 어려운 과제로 남아 있습니다. 최근의 LLM 유도 방식들은 텐서 프로그램 최적화를 반복적인 결정 과정으로 구성하지만, 기존 데이터셋들은 토큰 효율성이 낮은 표현 방식을 사용하여 엔드 투 엔드 (end-to-end)로 최적화된 프로그램 쌍만을 제공하며, 검증 가능한 단계별 감독 (step-level supervision)과 해석 가능성 (interpretability)이 부족합니다. 그 결과, LLMs는 거대한 조합 최적화 공간 (combinatorial optimization spaces)에서 신뢰할 수 있는 단일 단계 결정을 내리는 데 어려움을 겪습니다. 우리는 구조화된 사고 사슬 (Chain-of-Thought, CoT) 추론과 함께 근거가 있고 원자적인 단계별 감독을 제공하는 텐서 프로그램 최적화용 사후 학습 (post-training) 데이터셋인 Step-TP를 소개합니다. Step-TP는 중간 프로그램 상태에 대해 폐쇄형 추론 루프 (closed reasoning loop)를 형성하여, 결과 모방이 아닌 신뢰할 수 있는 다단계 최적화를 가능하게 합니다. 이 설계는 네 가지 원칙에 따라 유도되었습니다: (i) TVM TIR로 결정론적으로 낮아지는(lowers to) 토큰 효율적이고 검증 가능한 중간 표현 (Intermediate Representation, IR); (ii) 복잡한 궤적을 해석 가능한 단일 단계 결정으로 분해하는 원자적이고 결합 가능한 최적화 전략; (iii) 명시적인 IR-to-IR 상태 전이와 결합된 구조화된 CoT 감독; (iv) 커버리지를 균형 있게 유지하면서 지름길 편향 (shortcut exploitation)을 방지하기 위한 전략 필터링. 데이터셋과 구현체는 GitHub 링크 https://github.com/LIUMENGFAN-gif/StepTP 에서 확인할 수 있습니다.