Probabilistic Programs of Thought

논문 제목: Probabilistic Programs of Thought (PPoT)

저자: Poorva Garg, Renato Lui Geh, Daniel Israel, Todd Millstein, Kyle Richardson, Guy Van den Broeck (UCLA, Allen Institute for AI)

1. 문제 제기 (Problem)

대형 언어 모델 (LLM) 은 코드 생성, 수학 문제 해결, 구조화된 출력 생성 등 다양한 작업에 널리 사용되고 있습니다. 특히 '생각의 프로그램 (Programs of Thought)' 패러다임 하에서 LLM 은 복잡한 추론을 위해 코드를 생성하고 실행합니다.

기존의 일반적인 접근 방식은 다음과 같습니다:

1. LLM 에 프롬프트를 입력하여 코드를 생성합니다.
2. 생성된 코드를 실행하여 정답을 검증합니다.
3. 코드가 틀리면 다시 LLM 에게 요청하여 새로운 샘플을 생성합니다 (Best-of-n, Beam Search 등).

한계점:

• 높은 컴퓨팅 비용: $n$개의 프로그램을 샘플링하려면 $n$번의 LLM 생성 (GPU 추론) 이 필요합니다. 모델 크기가 커질수록 이는 prohibitively expensive(금지할 정도로 비싼) 비용이 됩니다.
• 비효율성: LLM 은 토큰 확률 분포를 내재하고 있음에도 불구하고, 한 번 샘플링된 토큰을 확정적인 값으로 취급하고 버린 후, 새로운 샘플을 위해 다시 전체 모델을 실행해야 합니다. 이는 토큰 수준의 불확실성 (entropy) 을 활용하지 못하게 합니다.

2. 방법론 (Methodology)

저자들은 LLM 이 생성한 프로그램 내부에 숨겨진 확률 분포를 노출시켜, **확률적 프로그램 (Probabilistic Program)**으로 변환하는 새로운 테스트 시간 프레임워크인 PPoT를 제안합니다.

핵심 아이디어

LLM 이 생성한 결정론적 (deterministic) 프로그램을 그대로 실행하는 대신, 프로그램 내의 특정 토큰 (숫자, 연산자 등) 을 확률 변수 (Random Variables) 로 변환합니다. 이때 각 확률 변수의 분포는 LLM 의 **다음 토큰 확률 (next-token probabilities)**로 파라미터화됩니다.

작동 원리

1. LLM 샘플링: LLM 에서 초기 프로그램 (예: 수학 문제 해결 코드) 을 한 번 생성합니다.
2. 토큰 분석 및 변환: 생성된 프로그램 내에서 불확실성이 높은 토큰 (주로 숫자, 비교 연산자, 산술 연산자 등) 을 식별합니다. 이 토큰들을 고정된 값이 아닌 확률 변수 $X_i$로 변환합니다.
   • 예: return 0 + 2에서 0이 틀린 경우, 이를 X ~ categorical(PM)으로 변환하여 $0$ 대신 다른 숫자가 나올 확률을 포함시킵니다.
3. 확률적 프로그램 생성: 변환된 코드는 확률적 프로그래밍 언어 (PPL) 구문을 가진 프로그램이 됩니다. 이 프로그램은 단일 LLM 호출로 생성된 것임에도 불구하고, 지수적으로 많은 수의 결정론적 프로그램을 암시적으로 표현합니다.
4. 저비용 샘플링: 생성된 확률적 프로그램은 CPU 만 사용하여 매우 저렴하게 여러 번 샘플링할 수 있습니다. GPU 추론 없이도 이 프로그램에서 새로운 프로그램 변형을 생성하고 실행하여 정답을 찾을 수 있습니다.
5. 조건부 샘플링: 원래 LLM 이 생성한 잘못된 프로그램을 재현하지 않도록 조건부 (conditioning) 를 걸어 더 다양한 샘플을 추출합니다.

이론적 근거

• 독립성 가정 (Assumption 1): 재샘플링된 토큰과 그 이후의 토큰들이 조건부 독립이라고 가정합니다. 이 가정 하에서 PPoT 를 통해 생성된 샘플들의 경험적 분포는 원래 LLM 의 분포에 수렴함이 증명됩니다 (Theorem 1).
• 계산 효율성: PPoT 는 LLM 의 분포를 '추론 가능한 (tractable)' 형태로 근사화하여, 추가적인 GPU 연산 없이도 분포를 탐색할 수 있게 합니다.

3. 주요 기여 (Key Contributions)

1. PPoT 프레임워크 제안: 구조화된 출력 생성 (코드, 수학 등) 을 위한 새로운 테스트 시간 디코딩 기법을 제안했습니다. 이는 단일 LLM 호출로 효율적으로 더 많은 샘플을 생성할 수 있게 합니다.
2. 분포적 타당성 (Distributional Soundness): 특정 독립성 가정 하에서 PPoT 가 LLM 의 원래 분포와 일치함을 이론적으로 증명했습니다.
3. 실험적 검증: GSM8k(수학 추론), Plot2Code(코드 생성), CRUXEval(코드 이해 및 실행) 의 3 가지 벤치마크에서 Qwen2.5 모델 (0.5B, 3B, 7B) 을 사용하여 PPoT 의 성능을 검증했습니다.
   • 성능 향상: 추가적인 LLM 생성 없이도 Pass@k 성능이 일관되게 향상되었습니다 (최대 7% 정확도 증가).
   • 계산 효율성: PPoT 샘플링은 GPU 연산이 필요 없으며, CPU 오버헤드만 발생하여 LLM 직접 샘플링 대비 계산 비용을 획기적으로 줄였습니다.
4. 확장성: PPoT 는 코드 생성뿐만 아니라 구조화된 텍스트 출력 (예: 리스트, 객체) 에도 적용 가능함을 보였습니다.

4. 실험 결과 (Results)

• GSM8k (수학 문제): Qwen2.5 Coder (7B) 모델에서 LLM 샘플 20 개를 사용할 때, PPoT 를 통해 각 LLM 샘플당 5 개의 추가 샘플을 생성하면 (총 100 개 후보), 정확도가 82.71% 에서 **93.02%**로 크게 향상되었습니다. 이는 LLM 만으로 20 개 샘플을 생성했을 때보다 훨씬 높은 성능을 내며, GPU 연산은 20 회만 사용했습니다.
• Plot2Code (그래프 생성): LLM 이 런타임 오류를 발생시키는 코드를 생성하더라도, PPoT 는 확률적 분포를 통해 올바른 차수나 매개변수를 가진 코드를 샘플링하여 오류를 수정하고 Ground Truth 와 유사한 그래프를 생성했습니다.
• CRUXEval (코드 이해): 구조화된 출력 (리스트 등) 에 대해서도 부분 샘플링 (subset sampling) 기법을 적용하여 성능을 개선했습니다.
• 계산 비용: PPoT 를 사용한 경우, LLM 호출 횟수 ($k$) 는 동일하게 유지하면서 샘플 수를 $k \times (1+m)$로 늘릴 수 있어, 동일한 정확도를 달성하는 데 필요한 LLM 호출 횟수를 크게 줄일 수 있음을 보였습니다. (예: 20 개의 PPoT 샘플은 약 39 개의 추가 LLM 샘플과 동등한 효과를 냄).

5. 의의 및 결론 (Significance)

• 비용 효율적인 추론: LLM 의 추론 능력을 극대화하면서도 GPU 비용을 절감할 수 있는 새로운 패러다임을 제시했습니다. 이는 대규모 모델의 실용적 배포에 중요한 의미를 가집니다.
• 불확실성 활용: LLM 이 생성한 토큰의 확률 분포를 단순히 버리는 것이 아니라, 이를 확률적 프로그램으로 변환하여 체계적으로 활용함으로써 모델의 잠재력을 끌어냈습니다.
• 미래 연구 방향: 본 연구는 LLM 생성물을 확률적 프로그램의 추상화로 보는 새로운 연구 방향을 열었습니다. 향후 더 복잡한 제어 흐름이나 서브-식 (sub-expressions) 을 확률 변수로 확장하거나, 기대값 계산 등 다른 확률적 쿼리에 적용할 수 있는 가능성이 열렸습니다.

요약하자면, PPoT는 LLM 이 생성한 코드를 "고정된 정답"이 아닌 "확률적 분포를 가진 프로그램"으로 재해석함으로써, 적은 GPU 비용으로 더 많은 시도 (샘플링) 를 가능하게 하여 코드 생성 및 추론 작업의 정확도를 획기적으로 높인 혁신적인 방법론입니다.