컴파일러 월드 모델을 향하여: 효율적인 텐서 프로그램 탐색을 위한 잠재 역학 (Latent Dynamics) 학습

텐서 프로그램 최적화 (Tensor program optimization)는 현대 머신러닝 시스템에 필수적이지만, 그 탐색 공간 (search space)은 매우 방대합니다. 기존의 자동 스케줄러 (auto-schedulers)는 학습된 비용 모델 (cost models)을 통해 측정 비용을 줄이지만, 대개 각 후보를 정적인 코드 스냅샷 (static code snapshot)으로 평가하며, 해당 코드를 생성한 스케줄링 궤적 (schedule trajectory)을 무시합니다. 이로 인해 이들은 액션 의존성 (action dependencies)에 민감하지 못하며 표면적인 코드 변화에 취약합니다. 우리는 스케줄링 평가를 프로그램 상태 (program states)에 대한 액션 조건부 잠재 역학 (action-conditioned latent dynamics)으로 모델링하는 extit{월드 모델 기반 (world-model-inspired)} 평가기를 제안합니다. 초기 프로그램에서 시작하여, 경량화된 전이 모델 (transition model)을 통해 연속적인 잠재 공간 (latent space)에서 스케줄링 액션을 전개 (roll out)함으로써, 비용이 많이 드는 AST 변이 (AST mutation)와 반복적인 코드 인코딩 (code encoding)을 피합니다. 최종적인 동적 표현 (dynamic representation)은 액션 및 하드웨어 특징 (features)과 결합되어 후보들의 순위를 매깁니다. TVM AutoScheduler에 구현된 우리의 방법은 동일한 64회 측정 (64-trial) 예산 하에서 GPU에서는 Ansor 대비 대표 서브그래프 (representative-subgraph) 지연 시간 (latency)을 1.37배, CPU에서는 1.54배 개선합니다. 또한 10배 적은 측정 횟수를 사용하면서도 기하 평균 (geometric mean) 기준 2.2% 이내의 오차로 Ansor-10K와 일치하며, PyTorch/PyTorch-opt(cuDNN) 대비 전체 모델 추론 (full-model inference) 속도를 기하 평균 4.61배/3.67배 가속화합니다.