Probabilistic Programs of Thought

1. Problem Statement

Large Language Models (LLMs) are increasingly used for code generation, mathematical reasoning, and structured output tasks. The standard approach to improving performance in these domains involves test-time decoding strategies like Best-of-n or beam search, where the model generates $n$ distinct programs, executes them, and selects the best one based on a verifier.

The Core Limitation:
Generating $n$ independent samples requires $n$ separate, GPU-intensive forward passes through the LLM. As model sizes and the required number of samples ($n$) increase, this process becomes prohibitively expensive in terms of compute time and cost. Existing methods discard the rich probability distribution (next-token logits) generated during the initial sampling, treating the output as a single deterministic string rather than a distribution over possible strings.

2. Methodology: Probabilistic Programs of Thought (PPoT)

The authors propose PPoT, a novel test-time framework that treats an LLM-generated program not as a static string, but as a probabilistic program. This allows for the generation of exponentially many new samples from a single LLM call using cheap CPU-based probabilistic inference.

Core Concept

1. LLM as a Probabilistic Program: The authors formalize the LLM's generation process as a probabilistic program where each token $X_i$ is a categorical random variable parameterized by the LLM's next-token probabilities $P_M(\cdot | X_{<i}, t)$.
2. Compilation: Given a deterministic program generated by the LLM, the system identifies specific "high-entropy" or critical tokens (e.g., digits, arithmetic operators, comparison symbols) and converts them into random variables ($X_i$).
   • The distribution of these variables is set to the original next-token probabilities provided by the LLM.
   • Tokens not selected for conversion remain fixed constants.
3. Probabilistic Inference: The resulting structure is a probabilistic program that compactly represents a distribution over exponentially many deterministic programs.
4. Resampling: Instead of calling the LLM again, the system performs probabilistic inference (sampling) on this compiled program. This generates new, valid program variations by resampling the random variables while keeping the rest of the structure intact.

Algorithmic Workflow (Algorithm 2)

1. Generate: The LLM generates one initial program ($k=1$) and returns the token sequence along with the next-token probability distributions.
2. Analyze & Compile: A token_analysis step selects a subset of token indices $L$ (e.g., numbers and operators) to treat as random variables. A compile step constructs a probabilistic program where these tokens are random variables with distributions derived from the LLM's logits.
3. Sample: The system samples $m$ new programs from this probabilistic program.
4. Execute & Verify: All $m$ new programs are executed and verified.
5. Result: For the cost of 1 LLM generation, the system obtains $1 + m$ candidate solutions.

Theoretical Soundness

The paper addresses the fact that resampling a token in an autoregressive sequence technically requires resampling all subsequent tokens to maintain the true joint distribution $P_{code}$.

• Assumption 1 (Independence): The authors assume that tokens succeeding a resampled token $X_j$ are independent of $X_j$ given the preceding tokens.
• Theorem 1: Under this independence assumption, the empirical distribution of samples generated by PPoT converges almost surely to the true LLM distribution $P_{code}$. The authors argue this assumption is robust enough for practical code generation tasks where local token changes (e.g., changing a + to a -) drastically alter execution results.

3. Key Contributions

1. Novel Decoding Technique: Introduction of PPoT, which leverages the perspective of probabilistic programming to perform efficient test-time inference on LLM-generated code.
2. Distributional Soundness: Formal proof that under specific independence assumptions, PPoT samples approximate the LLM's underlying distribution without requiring additional GPU compute.
3. Efficiency: Demonstrates that PPoT can generate hundreds of samples with negligible CPU overhead, avoiding the GPU bottleneck of standard Best-of-n approaches.
4. Empirical Validation: Extensive evaluation across three diverse benchmarks showing consistent performance gains.

4. Experimental Results

The authors evaluated PPoT on Qwen2.5 models (0.5B, 3B, 7B) across three benchmarks:

• GSM8k (Mathematical Reasoning):
  • Result: With $m=5$ PPoT samples per LLM generation, accuracy improved by 2–7% across all model sizes.
  • Efficiency: PPoT achieved the accuracy of 20 LLM samples using only 8 LLM samples (a 60% reduction in GPU calls).
• Plot2Code (Code Generation from Images):
  • Result: Improved mean text match scores by 0.5–6%.
  • Qualitative: PPoT successfully corrected runtime errors in LLM-generated code (e.g., fixing vector dimension mismatches) by resampling specific tokens.
• CRUXEval (Structured Output/Code Understanding):
  • Result: Improved accuracy by 3–7%.
  • Method: Used subset sampling to resample tokens in structured outputs (e.g., lists, strings) to find valid inputs that satisfy a function's output.

Computational Cost:

• The compilation and sampling steps in PPoT are performed on the CPU.
• Runtime analysis shows that the time taken for $m=20$ PPoT samples is nearly indistinguishable from $m=0$ (standard LLM sampling), confirming that the accuracy gains come at near-zero computational cost.

5. Significance and Future Directions

• Paradigm Shift: PPoT shifts the paradigm from "generate more samples via brute force" to "extract more value from a single sample via probabilistic reasoning."
• Scalability: As LLMs grow larger and inference costs rise, PPoT offers a scalable way to improve reliability without linearly increasing GPU costs.
• Generalizability: While demonstrated on code, the framework applies to any structured output where a verifier exists (e.g., JSON generation, logical reasoning).
• Future Work: The authors suggest expanding the scope of random variables beyond simple tokens to include sub-expressions and control flow structures, and using these probabilistic programs for computing expectations or conditioning on complex constraints.

In summary, Probabilistic Programs of Thought provides a mathematically grounded, computationally efficient method to "distill" the uncertainty inherent in LLM outputs, allowing for the exploration of a vast solution space with minimal additional compute resources.