KL-Budgeted Poisson-Race Sampling
Implementation & benchmark of arXiv:2608.26996 — A note on shifting distributions via Poisson races
Source paper: A note on shifting distributions via Poisson races arXiv:2608.26996 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace repeated rejection sampling from a neural proposal P with a Poisson race that selects one proposal candidate with exactly the target distribution Q. Candidate i receives a Poisson arrival time A_i and race score A_i divided by Q(X_i)/P(X_i); the smallest score wins. When Q is close to P in KL divergence, the selected index has low expected logarithmic rank, providing an adaptive compute budget instead of a fixed number of candidates.
Formulas
Mathematical statement
The paper defines a P-to-Q simulator as an index I such that X_I has distribution Q when X_1,X_2,... are iid from P, with cost(I)=E[log_2 I]. Its theorem states that if D(Q||P)<infinity, there exists a simulator satisfying E[log_2 I] <= D(Q||P)+1.45||Q-P||_1, while every simulator satisfies E[log_2 I] >= 0.5 max{D(Q||P),||Q-P||_1}. Here D(Q||P)=sum_x Q(x)log_2(Q(x)/P(x)) is KL divergence and ||Q-P||_1=sum_x|Q(x)-P(x)| is total variation in unhalved L1 form. A directly implementable Poisson race draws iid proposal marks X_i~P and iid exponential increments E_i~Exp(1), forms A_i=sum_{j=1}^i E_j, and returns I=argmin_i A_i/w(X_i), where w(x)=Q(x)/P(x). The marked process for symbol x has proposal rate P(x); rescaling its clock by w(x) produces rate P(x)w(x)=Q(x), so the winning mark is Q-distributed. If w_max=max_x w(x) is known, exact scanning can stop after i when A_{i+1}/w_max >= S_min, where S_min=min_{j<=i} A_j/w(X_j), because all future arrivals are no earlier than A_{i+1}.
Implementation notes
Integrate the method at an inference-time categorical sampling point, such as next-token generation or a discrete latent-variable posterior. Let a cheap proposal network produce logits l_P and an accurate target network produce logits l_Q over vocabulary size V. Compute log probabilities with log-softmax: logp and logq. Draw proposal tokens x_i~Categorical(exp(logp)), draw exponential increments e_i=-log(U_i) for U_i~Uniform(0,1), and maintain A_i=A_{i-1}+e_i. For each token compute logw=logq[x_i]-logp[x_i] and log-score=log(A_i)-logw; keep the candidate with minimum log-score. Precompute logw_max=max_x(logq[x]-logp[x]). After observing candidate i, draw or expose the next arrival A_next=A_i+Exp(1), and stop exactly when log(A_next)-logw_max >= best_log_score. Pseudocode: A=0; best_log_score=inf; best_x=None; while True: A += Exp(1); x=sample(p); s=log(A)+logp[x]-logq[x]; if s<best_log_score: best_log_score=s; best_x=x; A_next=A+Exp(1); if log(A_next)-logw_max>=best_log_score: break; return best_x. Compute the race in log space to avoid overflow. The mathematical quantities are the score, stopping rule, and exactness guarantee; estimate KL, mean log2(index), and wall-clock cost empirically. First test on a small WikiText-2 Transformer, comparing direct target sampling, ordinary rejection sampling, and Poisson-race sampling with a smaller proposal model. Measure marginal distribution error, target evaluations, candidate index, throughput, and perplexity. Success requires distribution error below 1e-3 and lower target-computation cost at matched perplexity, with mean log-index empirically below the KL-plus-total-variation upper bound.
Verification
Stage 1 · Toy mechanism gate: Failed ✗
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built and numerically verified the finite-categorical Poisson race with the exact look-ahead stopping rule. The output marginal matches Q within Monte Carlo noise and mean log2 rank increases with KL, with a small candidate-count reduction versus rejection sampling (1.599 vs 1.604 close; 4.768 vs 4.813 moderate; 9.629 vs 9.663 farther). However, the observed close-case mean log2 rank (0.489) is above the stated KL-plus-L1 upper bound (0.124), and the pure-Python race is slower than rejection sampling; no neural inference or target-computation speedup was demonstrated.
- Agent confidence
- 8/10
- Baseline
- Rejection sampling: mean candidate index 1.604, 4.813, 9.663 for close/moderate/farther cases; TV errors 0.0097, 0.0065, 0.0062; runtime 3.50, 9.97, 18.40 s.
- Idea
- Poisson race: mean candidate index 1.599, 4.768, 9.629; mean log2 index 0.489, 1.753, 2.646; TV errors 0.0079, 0.0066, 0.0102; runtime 4.53, 11.63, 22.23 s. Identical P=Q always stops at index 1.
Limitations: Only a small finite categorical toy test was run; no Transformer, language-model logits, GPU implementation, batched/vectorized implementation, FLOP-matched target-network benchmark, or comparison at matched perplexity was tested. The theorem is existential, so failure of this direct race to meet its upper bound does not disprove the paper's theorem, but it limits the claimed practical conclusion here.
How to run: python3 poisson_race_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- poisson_race_experiment.py 4.1 KB View Raw
- report.md 1.6 KB View
- results.json 3.0 KB View Raw