기술 요약: 코드용 회귀 언어 모델

문제 정의

본 논문은 소스 코드 또는 계산 그래프 표현으로부터 코드 실행의 수치적 결과 (메모리 풋프린트, 지연 시간, 정확도 등) 를 직접 예측하는 '코드 - 지표 회귀 (code-to-metric regression)' 과제의 해결책을 다룹니다. 역사적으로 이 작업은 프로그래밍 언어의 개방적 성격과 무겁고 도메인 특화적인 '특성 공학 (feature engineering)'에 대한 의존성으로 인해 방해받아 왔습니다. 기존 방법들은 종종 코드를 구문 트리, 인접 행렬과 같은 경직된 그래프 표현이나 명령어 수동 계산을 통한 표 형식 벡터로 변환해야 했으며, 이는 새로운 작업, 언어, 또는 커널 유형으로의 전이성을 제한했습니다. 핵심적인 어려움은 다양한 도메인 전반에 걸친 일반화 유연성을 잃지 않으면서도 고도로 비표형식적이고 그래프 기반인 데이터에 대한 정확한 회귀 모델을 학습하는 데 있습니다.

방법론

저자들은 현대의 대규모 언어 모델 (LLM) 패러다임과 정렬하여 코드 - 지표 회귀를 표준적인 다음 토큰 예측 문제로 취급하는 통합 프레임워크인 '회귀 언어 모델 (Regression Language Model, RLM)'을 제안합니다.

아키텍처 및 학습

• 통합 인코더 - 디코더: RLM 은 입력 텍스트를 처리하기 위해 고정된 인코더 (구체적으로 T5Gemma-S) 와 수치적 출력을 생성하기 위해 학습 가능한 디코더를 활용합니다. 인코더는 표준 언어 데이터로 사전 학습되며, 역전파는 오직 디코더에서만 수행되어 학습 비용을 크게 절감합니다.
• 텍스트 - 텍스트 회귀: 입력을 고정된 길이의 벡터로 매핑하는 전통적인 접근법과 달리, RLM 은 코드 (예: Python, C++, ONNX 그래프) 의 원시 텍스트 표현을 받아 토큰 시퀀스로 수치 값을 출력합니다.
• 전문 수치 토큰화: 정규화 없이 광범위한 동적 범위 ( $10^{-2}$ 에서 $10^6$ 까지) 의 지표를 처리하기 위해, 모델은 명시적인 자릿수별 토큰화를 사용합니다. 숫자는 부호, 지수, 가수 (mantissa) 를 위한 특수 토큰을 사용하여 표현됩니다 (예: $72.5$는 <+><-><1><7><2><5> 로 표현). 이는 수치적 불안정성과 데이터셋별 스케일링의 필요성을 제거합니다.
• 다중 목적 모델링: 디코더의 자기회귀적 특성을 활용하여, 모델은 정확도와 하드웨어 지연 시간과 같은 목적 간의 내재적 상관관계를 포착하면서 여러 지표를 순차적으로 조건부 예측할 수 있습니다 (예: $p(y_{latency} | y_{accuracy}, x)$).

데이터 소스

모델은 다음과 같이 구성된 다양하고 통합된 데이터셋으로 학습됩니다:

1. 고수준 프로그래밍: APPS(LeetCode 문제), CodeNet(37 개 언어의 1400 만 개 샘플), KernelBook(Triton GPU 커널) 의 코드.
2. 신경 아키텍처 탐색 (NAS): NASBench-101/201, FBNet, Once-for-All 등의 벤치마크에서 나온 52 만 개의 고유 아키텍처를 ONNX 중간 표현으로 표현한 데이터.

주요 기여

1. 통합 범용 모델: 단일 RLM 이 여러 언어 (Python, C++ 등) 의 메모리 사용량, GPU 커널 지연 시간, 신경망 정확도/지연 시간을 동시에 예측할 수 있음을 보여줌으로써 작업별 특성 공학의 필요성을 제거합니다.
2. 그래프 기반 방법 대비 우위: 밀집된 ONNX 텍스트 표현을 읽음에도 불구하고, RLM 은 NAS 벤치마크에서 기존에 GNN 이 우세했던 설계 공간 (예: Kendall-Tau 0.46 달성) 에서 최첨단 그래프 신경망 (GNN) 베이스라인을 능가합니다.
3. 효율성 및 확장성: 고정된 인코더 사용으로 비용 효율적인 학습이 가능합니다. 저자들은 상대적으로 작은 3 억 파라미터 모델이 높은 성능을 달성하며, 6 억 파라미터로 확장하면 추가적인 향상을 얻을 수 있음을 보여줍니다.
4. 정규화 없는 예측: 자릿수별 토큰화 전략은 수동 정규화 없이도 수백만 배의 범위에서 값을 예측할 수 있게 하여, MSE 기반 회귀 헤드 대비 중요한 장점을 제공합니다.

실험 결과

• 코드 성능: APPS 데이터셋에서 RLM 은 메모리 예측에 대해 0.93의 스피어만 순위 상관관계 ($\rho$) 를 달성합니다. CodeNet 에서는 17 개 언어에서 평균 $\rho > 0.5$를 달성하며, C++ 의 경우 0.748 에 도달합니다.
• NAS 성능: RLM 은 5 개의 고전적인 NAS 설계 공간에서 평균 Kendall-Tau 0.46을 달성하여 FLANT 및 GNN 기반 인코더와 같은 이전 최첨단 방법을 능가합니다. 정확도와 Pixel3, Eyeriss, Intel CPU, Nvidia GPU 와 같은 하드웨어별 지연 시간을 동시에 성공적으로 예측합니다.
• 절대 실험 (Ablation Studies):
  • 고정된 인코더: 무작위 초기화에 비해 사전 학습된 인코더는 수렴과 최종 성능을 크게 향상시킵니다.
  • 디코더 대 회귀 헤드: 디코더 기반 출력 (0.800 스피어만 $\rho$) 은 명시적 회귀 헤드 (정규화 없이 0.478, 정규화 시 0.717) 를 크게 능가합니다.
  • 토큰화: ONNX 그래프를 위한 맞춤형 소형 토크나이저는 성능을 향상시킵니다 (기본 토크나이저 대비 0.723 대 0.533).
  • 컨텍스트 학습 (ICL): 대규모 모델 (예: GPT-5) 을 사용한 ICL 은 샷 (shot) 이 증가함에 따라 향상되지만, 인-가중치 학습된 RLM 보다 일관되게 낮으며 컨텍스트 창 제한으로 인해 훨씬 더 높은 추론 비용을 초래합니다.

의의 및 주장

본 논문은 계산 그래프 회귀를 범용 다음 토큰 예측 문제로 단순화하는 길을 열었다고 주장합니다. 언어 모델이 코드와 그래프의 순수 텍스트 표현을 효과적으로 회귀할 수 있음을 보여줌으로써, 이 분야는 경직되고 수동으로 설계된 특성 공학에서 더 유연하고 통합된 접근 방식으로 전환할 수 있음을 시사합니다.

저자들은 이를 다음과 같은 분야에서 유용한 도구로 위치시킵니다:

• 프로그램 탐색: 효율적인 코드 솔루션 선택 가속화.
• 하드웨어 - 소프트웨어 공동 설계: 다양한 하드웨어 플랫폼에서의 지연 시간 예측.
• 컴파일러 최적화: 컴파일러 패스에 대한 성능 지표 추정.

논문은 raw 코드에서 수치적 결과를 예측하는 데 이 방법이 유망함을 보여주지만, 전체 과학 실험의 결과를 예측하는 것으로 확장하는 것은 향후 연구의 열린 질문으로 남긴다는 점에서 더 넓은 함의에 대해서는 겸손한 태도를 유지합니다. 주요 기여는 후처리나 전문적인 특징 추출 없이 코드와 계산 그래프를 위한 범용 회귀기로 언어 모델의 일반적 능력을 확립하는 데 있습니다.