이 논문은 AI를 활용하여 과학 문제에 대한 새로운 "state of the art"(SOTA)를 발견하는 방법을 제시합니다. 기존의 연구(예: AlphaEvolve)가 정지된(frozen) LLM을 프롬프트(prompt)하여 검색을 수행하는 것과 달리, 이 논문은 LLM이 특정 테스트 문제에 대한 경험을 바탕으로 계속 훈련할 수 있도록 "test-time"에 강화 학습(Reinforcement Learning, RL)을 수행하는 방법을 제안하며, 이를 "Test-Time Training to Discover" (TTT-Discover)라고 명명합니다.

핵심 방법론: TTT-Discover

TTT-Discover는 표준 RL과는 다른 두 가지 중요한 목표를 가지고 있습니다. 첫째, 정책(policy)은 단 하나의 문제를 해결하면 되고 다른 문제로 일반화할 필요가 없습니다. 둘째, 평균적인 보상 최대화가 아닌 단 하나의 '최고의' 솔루션을 찾는 것이 목표입니다. 이러한 특성을 활용하여 TTT-Discover는 가장 유망한 솔루션을 우선시하도록 학습 목표와 검색 서브루틴을 설계합니다.

1. 환경(Environment) 정의: 각 과학 문제는 Markov Decision Process (MDP)로 정의되는 환경을 구성합니다.
   • 문제 설명($d$): LLM 정책에 컨텍스트로 제공됩니다.
   • 상태($s$): 후보 솔루션을 나타내며, 예를 들어 커널 구현 코드, 수학적 구성물 등이 될 수 있습니다.
   • 보상 함수($R(s)$): 솔루션의 품질을 측정하는 연속적인 값입니다(예: 역 런타임, 점수 등).
   • 행동($a$): LLM 정책 $\pi_\theta(\cdot | d, s)$에 의해 생성되는 코드와 '사고 토큰'(thinking tokens)을 포함합니다.
   • 전이 함수($T$): 행동 $a$를 파싱(parse)하고 실행하여 새로운 상태 $s'$를 생성합니다($s' = T(a)$).
   • 발견(Discovery): $R(s) > r_{\text{sota}}$ (기존 SOTA 보상)인 상태 $s$를 찾는 것을 의미합니다.

2. 학습 목표(Learning Objective): Entropic Objective:

표준 RL 알고리즘이 기대 보상을 최대화하는 반면, TTT-Discover는 최대 보상을 선호하는 엔트로피 목표 $J_\beta(\theta)$를 사용합니다.
$$J_\beta(\theta) = E_{s \sim \text{reuse}(\mathcal{H})} \left[ \log E_{a \sim \pi_\theta(\cdot|s)} \left[ e^{\beta(s)R(s,a)} \right] \right]$$
여기서 $\beta(s)$는 초기 상태(initial state)에 따라 적응적으로 설정되어 유도된 정책의 KL divergence를 제약합니다. 이는 학습 초기의 불안정성을 피하고 학습 후반의 수렴 문제를 해결하는 데 도움이 됩니다. 그라디언트 $\nabla_\theta J_\beta(\theta)$는 다음과 같이 계산됩니다.
$$\nabla_\theta J_\beta(\theta) = E_{s \sim \text{reuse}(\mathcal{H}), a \sim \pi_\theta(\cdot|s)} \left[ \frac{e^{\beta(s)R(s,a)}}{\pi_\theta(\cdot|s)[e^{\beta(s)R(s,a)}]} \nabla_\theta \log \pi_\theta(a|s) \right]$$
TTT-Discover는 이 그라디언트에 KL 페널티를 적용하여 장점(advantage)을 형성합니다: $A(a;s) = w_\beta(s)(a) - 1 - \lambda \log \frac{\pi_\theta(a|s)}{\pi_{\theta_0}(a|s)}$, 여기서 $w_\beta(s)(a) = \frac{e^{\beta(s)R(s,a)}}{\pi_\theta(\cdot|s)[e^{\beta(s)R(s,a)}]}$.

3. 상태 재사용(State Reuse): PUCT-inspired Rule:

이전 솔루션을 재사용하는 것은 효과적으로 시도(attempt)에 추가적인 타임스텝(timestep)을 더하여 더 복잡한 솔루션이 나타날 수 있도록 합니다. TTT-Discover는 PUCT(Polynomial Upper Confidence Trees)에서 영감을 받은 규칙을 사용하여 초기 상태를 선택합니다. 각 상태 $s$는 다음과 같이 점수가 매겨집니다.
$$Q(s) + c \cdot P(s) \cdot \sqrt{1 + T / (1 + n(s))}$$
여기서 $Q(s)$는 해당 상태 $s$를 초기 상태로 하여 생성된 자식(children) 중 '최대' 보상(평균이 아님)을 나타냅니다. 이는 '최고의' 결과를 찾는 발견 문제의 목표에 부합합니다. $P(s)$는 버퍼 내에서 $s$의 보상 순위에 비례하며, $n(s)$는 $s$ 또는 그 자손이 확장된 횟수, $T$는 총 확장 횟수, $c$는 탐색 계수(exploration coefficient)입니다.

4. 구현 세부 사항:
   • LLM: 오픈 모델인 gpt-oss-120b를 사용합니다.
   • 훈련 환경: Tinker API를 사용합니다.
   • 파인 튜닝: LoRA (rank 32)를 사용합니다.
   • 훈련 과정: 50번의 훈련 스텝을 수행하며, 각 스텝에서 512개의 롤아웃(rollout)을 생성합니다.
   • 비용: 문제당 수백 달러 수준으로 저렴합니다.

실험 및 결과

TTT-Discover는 수학, GPU 커널 엔지니어링, 알고리즘 설계, 생물학 등 다양한 분야의 연속 보상(continuous rewards) 문제에 대해 평가되었습니다. 모든 시도된 문제에 대한 결과를 보고하며, 거의 모든 문제에서 새로운 SOTA를 달성했습니다.

5. 수학 (Mathematics):
   • Erdős의 Minimum Overlap Problem: AlphaEvolve의 이전 SOTA인 0.380924를 능가하는 0.380876의 새로운 상한(upper bound)을 발견했습니다. TTT-Discover는 600개의 비대칭(asymmetric) 스텝 함수(step function)를 찾아냈습니다.
   • Autocorrelation Inequalities (첫 번째): ThetaEvolve의 SOTA인 1.50314를 넘어선 1.50286의 새로운 상한을 입증했습니다. 30000개의 스텝 함수를 사용했습니다.
   • Autocorrelation Inequalities (두 번째): SOTA를 달성하지 못했습니다 (0.959 vs. 0.961).

6. GPU 커널 엔지니어링 (GPUMode TriMul):
   • H100 GPU에서 1161µs 런타임의 커널을 찾아내어 이전 SOTA인 1371µs (최고 인간) 대비 2배 빠른 성능을 보였습니다.

7. 알고리즘 설계 (AtCoder):
   • 과거 AtCoder 알고리즘 대회 문제에서 SOTA 점수를 갱신했습니다 (567,062 vs. 566,997).

8. 생물학 (단일 세포 분석 denoising):
   • denoising 문제에서 SOTA 성능을 달성했습니다 (0.71 vs. 0.64).

TTT-Discover는 이전 최고 결과들이 클로즈드(closed) 프런티어(frontier) 모델을 요구했던 것과 달리, 오픈 모델(gpt-oss-120b)로 모든 결과를 달성했으며, 코드가 공개되어 재현 가능합니다. 이는 비슷한 아이디어를 공유하는 동시 연구(concurrent work)인 MiGrATe 및 ThetaEvolve와 비교했을 때도 TTT-Discover가 더 나은 성능을 보여줍니다.