Suppose you want to draw a random subset of exactly $n$ items from a universe of $N$ items, where each item $i$ has a positive weight $\w_i$.We assume $\w_i \in (0, \infty)$. This is without loss of generality: items with $\w_i = 0$ are excluded from the universe (they are never selected), and items with $\w_i = \infty$ are deterministically included (reduce $n$ and $N$ by 1 each). The general case reduces to the finite positive case by preprocessing. The conditional Poisson distribution assigns each size-$n$ subset a probability proportional to the product of its weights:✓

$$ \Ps(S) \propto \mathbf{1}\big[ |S| = n\big] \prod_{i \in S} \w_i $$

The normalizing constant $\Zw{\bw}{n} \defeq \sum_{|S|=n} \prod_{i \in S} \w_i$ is a weighted generalization of the binomial coefficient, which recovers $\binom{N}{n} = \Zw{\bw}{n}$ when $\bw = \mathbf{1}^N$.✓

Inclusion probabilities. The inclusion probability $\pip_i \defeq \sum_{S} \Ps(S)\, \mathbf{1}[i \in S]$. Higher weight means higher inclusion probability, but the relationship is nonlinear because the other weights also matter—doubling $\w_i$ does not double $\pip_i$, as the other items push back through the size constraint $|S| = n$. Later in this article, we provide an interactive widget for exploring the nonlinear relationship between $\bw$ and $\bpip$.

Why is this distribution special? The conditional Poisson distribution is an exponential family with natural parameters $\theta_i \defeq \log \w_i$ and sufficient statistics $\mathbf{1}[i \in S]$. Among all distributions over size-$n$ subsets with prescribed inclusion probabilities $\pip_i = P(i \in S)$, it is the unique maximum-entropy one✓—making the fewest assumptions beyond the marginals (Jaynes, 1957; Chen, Dempster & Liu, 1994), in the same sense that the Gaussian is max-entropy for given mean and variance. The log-normalizer $\log \Zw{\bw}{n}$ is convex in $\btheta$, so many properties follow mechanically: inclusion probabilities are the gradient ($\pip_i = \partial \log \Z / \partial \theta_i$) and fitting $\btheta$ to target inclusion probabilities is a convex optimization problem. The distribution is also called the exponential fixed-size design for this reason (Tillé, 2006).

Relationship to Poisson sampling. In Poisson sampling,Named after mathematician Siméon Denis Poisson. Although poisson is the French word for fish, no fishing metaphor is intended. each item $i$ is included independently with probability $p_i$, so the sample size $|S| = \sum_{i=1}^{N} \mathbf{1}[i \in S]$ is random. The conditional Poisson distribution conditions on $|S| = n$ exactly—fixing the sample size while preserving the relative inclusion odds. Under Poisson sampling, each item's inclusion probability is simply $\pip_i = p_i$; conditioning on $|S| = n$ makes $\pip_i$ depend on all the other weights too, which is what makes computing $\bpip$ nontrivial. The weight $\w_i$ is the odds of the $i$^th coin: $\w_i = p_i / (1 - p_i)$, equivalently $p_i = \w_i/(1+\w_i)$.✓

Sampling without replacement. The following construction shows how conditional Poisson can be used for sampling without replacement. Draw $n$ items independently from the categorical distribution over weights, and keep only the samples where all draws are distinct:

The resulting distribution over size-$n$ subsets is exactly $\Ps(S)$.✓ This rejection sampler is not a practical sampling algorithm;The acceptance probability of this construction is $n! \cdot \Zw{\bw}{n} / \W^n$.✓ it simply establishes what the distribution is. We will derive sampling algorithms that sample from $\Ps(S)$ efficiently. In machine learning, Meister, Amini, Vieira & Cotterell (2021) use it to make beam search stochastic: at each step, the beam is a conditional Poisson sample of the candidates, drawn without replacement, instead of the highest-scoring ones.

What this post covers. The computational challenges are: computing $\Zw{\bw}{n}$ and $\Ps(S)$, computing $\bpip$ from $\bw$, drawing exact samples $S \sim \Ps$, and the inverse problem of finding $\bw$ from target $\bpip$. This post gives efficient algorithms for all four—in $\mathcal{O}(N \log^2 n)$ time using a polynomial product tree—and times them against the $\mathcal{O}(Nn)$ dynamic program used by existing software. The appendices cover numerical stability, timings over a wider range of sizes, and further identities.

Software. The code is available as a Python library. As far as I can tell, this is the only publicly available library for conditional Poisson sampling in Python (or any language outside of R's survey-sampling packages). The existing R implementation—UPmaxentropy in the sampling package—uses $\mathcal{O}(Nn)$ dynamic programming, and its documentation cites Chen & Liu (1997). The product-tree algorithm used here does not appear in any prior software that I could find.

Computing $\Zw{\bw}{n}$, $\bpip$, and Samples

The key idea is that $\Zw{\bw}{n}$ is hiding inside a product of polynomials:

$$ (1 + \w_1 \z)(1 + \w_2 \z) \cdots (1 + \w_N \z) = \sum_{k=0}^{N} \Zw{\bw}{k}\, \z^k $$

Why does this work? When you expand the product, each factor $(1 + \w_i \z)$ contributes either $1$ (item $i$ excluded) or $\w_i \z$ (item $i$ included). A single term in the expansion picks one choice per factor, giving $\prod_{i \in S} \w_i \cdot \z^{|S|}$ for some subset $S$. The exponent of $\z$ counts how many items were included. Summing over all $2^N$ terms and collecting by powers of $\z$, the coefficient of $\z^k$ is exactly $\sum_{|S|=k} \prod_{i \in S} \w_i = \Zw{\bw}{k}$.

So the $n$^th coefficient is exactly $\Zw{\bw}{n}$, the normalizing constant.✓ Below, we compute it in two ways: by multiplying in the factors one at a time, in $\mathcal{O}(Nn)$ time, and by multiplying them in a balanced binary tree, in $\mathcal{O}(N \log^2 n)$ time. The second is a standard technique from computer algebra known as the subproduct tree (see von zur Gathen & Gerhard (2013)).

Polynomials as arrays. In code, a polynomial like $1 + 3\z + 2\z^2$ is just the array [1, 3, 2]. The $\z$ is a formal variable—we never substitute a number into it; it just tracks which coefficient is which. Multiplying two polynomials corresponds to convolving their coefficient arrays—for example, $(1 + 2\z)(1 + 3\z)$ convolves [1, 2] with [1, 3] to get [1, 5, 6]. It turns out that convolution can be done efficiently using the fast Fourier transform (FFT), costing $\mathcal{O}(d \log d)$ where $d$ is the degree, rather than $\mathcal{O}(d^2)$ for the schoolbook method. This is the key fact that makes the product tree $\mathcal{O}(N \log^2 n)$ instead of $\mathcal{O}(Nn)$.

Notation. We write $\llbracket f \rrbracket(\z^k)$ for the coefficient of $\z^k$ in a formal power series $f(\z) = \sum_k a_k \z^k$, i.e., $\llbracket f \rrbracket(\z^k) = a_k$. This is sometimes written $[\z^k]\, f(\z)$; we use the Scott bracket notation to avoid ambiguity with other uses of square brackets.

Computing the Normalizing Constant $\Z$

One factor at a time. The simplest way to expand the product is to multiply in one factor at a time, discarding coefficients above degree $n$, which never affect the coefficient of $\z^n$. After $m$ factors, the coefficient of $\z^k$ is $\Zw{\w_1, \ldots, \w_m}{k}$, and multiplying by the next factor gives the weighted Pascal recurrence:

$$\Zw{\w_1, \ldots, \w_m}{k} = \Zw{\w_1, \ldots, \w_{m-1}}{k} + \w_m \cdot \Zw{\w_1, \ldots, \w_{m-1}}{k-1}$$

with base cases $\Zw{\cdot}{0} = 1$ and $\Zw{\cdot}{k} = 0$ for $k < 0$ or $k > m$. This is a weighted generalization of Pascal's identity $\binom{m}{k} = \binom{m-1}{k} + \binom{m-1}{k-1}$: item $m$ is either excluded (first term) or included (second term).✓ Each step updates $n + 1$ coefficients, so this dynamic program (DP) computes $\Z$ in $\mathcal{O}(Nn)$ time.

The product tree. Multiplying by one degree-1 factor at a time gives FFT nothing to speed up. The product tree multiplies the same factors in a balanced order, so that most of the work is in multiplying large polynomials, where FFT pays off. Each leaf of a complete binary tree holds one degree-1 polynomial $(1 + \w_i \z)$, and each internal node multiplies its children's polynomials: the polynomial at node $a$, which we write $P_a(\z)$, is the product of $(1 + \w_i \z)$ over the leaves $i$ below $a$. The root holds the full product, whose $n$^th coefficient is $\Zw{\bw}{n}$.

Complexity. At each level of the tree, the total size of the polynomials being multiplied is $\mathcal{O}(N)$, and each multiplication is done via FFT in $\mathcal{O}(d \log d)$ time where $d$ is the degree. The recurrence is

$$T(N) = 2\,T(N/2) + \mathcal{O}(N \log N)$$

which solves to $T(N) = \mathcal{O}(N \log^2 N)$. We can do better, because we only need the coefficient at $\z^n$: intermediate polynomials can be truncated to degree $n$ without affecting the result. (Convolution of two polynomials truncated to degree $n$ still gives the correct coefficients up to degree $n$.) With truncation, the bottom $\log_2 n$ levels of the tree each cost $\mathcal{O}(N \log n)$, for $\mathcal{O}(N \log^2 n)$ in total. Above those levels, every polynomial has degree $n$ and the number of nodes halves at each level, so the remaining levels add only $\mathcal{O}(N \log n)$. The total is $\mathcal{O}(N \log^2 n)$.

Computing Inclusion Probabilities $\bpip$

The inclusion probability is the gradient of the log-normalizer (an exponential family identity):

$$\pip_i = \frac{\partial \log \Zw{\bw}{n}}{\partial \theta_i}$$

So any program that computes $\log \Zw{\bw}{n}$ gives $\bpip$ by running backpropagation on it—no new algorithm needed. By the Baur-Strassen theorem (1983), the gradient costs at most a small constant factor more than the forward computation.Baur & Strassen (1983) prove a factor of 3 for nonscalar operations on polynomials; for general reverse-mode automatic differentiation, Griewank & Walther (2008) put the factor between 3 and 4 in their cost model, or at most 5 if memory access can be arbitrarily slow (Section 4.6). Backpropagation through the Pascal DP gives $\bpip$ in $\mathcal{O}(Nn)$ time, and backpropagation through the product tree gives it in $\mathcal{O}(N \log^2 n)$.

Why backpropagation matters. Written out directly, $\pip_i$ requires the "leave-one-out" product $\prod_{j \neq i}(1 + \w_j \z)$—the full product with item $i$ removed. Recomputing it from scratch for each $i$ takes $N$ forward passes, $\mathcal{O}(N^2 n)$ with the DP. Backpropagation shares the work among the items, so all $N$ inclusion probabilities cost about as much as one forward pass. The product tree is just the fastest forward pass we know of; the gradient follows mechanically. In the PyTorch implementation, it comes from torch.autograd—no hand-coded tree traversal needed.

Drawing Samples

Sampling reuses the product tree (no gradient computation needed). Walk the tree top-down, splitting a quota of $n$ between children at each internal node:

This is exact, not approximate—it follows from the weighted Vandermonde identity: $\Zw{(\ba;\, \bb)}{k} = \sum_{j=0}^{k} \Zw{\ba}{j} \cdot \Zw{\bb}{k-j}$, which says the number of ways to choose $k$ items from $A \cup B$ is the convolution of the two groups' counts. Each term $q_j$ is proportional to the conditional probability of assigning $j$ items to the left child and $k - j$ to the right. At the leaves, quota 1 means "include this item"; quota 0 means "exclude."

The tree is built once ($\mathcal{O}(N \log^2 n)$) and reused across samples. Each sample costs $\mathcal{O}(n \log N)$: at each of the $\mathcal{O}(\log N)$ levels, only nodes with nonzero quota are visited, and the total quota per level is $n$.

The DP's sampler. The Pascal DP multiplies the factors in a tree too, a lopsided one: node $m$ joins the product of the first $m - 1$ factors (its left child) with the leaf $(1 + \w_m \z)$. At such a node, a quota of $k$ can be split only two ways: all $k$ to the left (item $m$ excluded), or $k - 1$ to the left and $1$ to item $m$ (included). The two Vandermonde terms are the two terms of the Pascal recurrence, so item $m$ is included with probability $\w_m \Zw{\w_1, \ldots, \w_{m-1}}{k-1} / \Zw{\w_1, \ldots, \w_m}{k}$.✓ The walk down this tree decides one item per level, reading these numbers from the DP's table, so each sample costs $\mathcal{O}(N)$.

All computations are verified against brute-force enumeration in the test suite.✓

Basic Usage

Now that we've seen how the product tree works under the hood, here's the library interface.

>>> import torch
>>> from conditional_poisson import ConditionalPoissonTorch
>>> cp = ConditionalPoissonTorch.from_weights(n=3, w=[0.68, 1.02, 0.55, 1.63, 0.67, 2.82])

>>> cp.log_normalizer
3.3621224365508637

>>> cp.incl_prob
tensor([0.3813, 0.5038, 0.3236, 0.6424, 0.3771, 0.7718], dtype=torch.float64)

>>> torch.stack([cp.sample() for _ in range(5)])
tensor([[1, 3, 5],
        [0, 3, 5],
        [1, 3, 4],
        [0, 1, 3],
        [1, 2, 5]])

>>> # Inverse problem: find weights from target inclusion probabilities
>>> cp_fit = ConditionalPoissonTorch.fit(cp.incl_prob, n=3)
>>> (cp_fit.incl_prob - cp.incl_prob).abs().max()
tensor(8.3039e-11, dtype=torch.float64)

The Poisson Approximation

The inclusion probabilities $\pip_i = P(i \in S)$ always sum to $n$,✓ and each $\pip_i \in [0, 1]$.✓ Items with larger weights get higher inclusion probabilities, but the relationship is nonlinear—the other weights push back through the size constraint $|S| = n$.

Poisson approximation. Scaling every weight by the same factor $r > 0$ does not change the distribution: it multiplies the weight of every size-$n$ subset by $r^n$, which cancels in the normalization.✓ It does change the Poisson sampling design behind the weights. With weights $r\bw$, the coins have probabilities

$$p_i(r) \defeq \frac{\w_i r}{1 + \w_i r},$$

so the coins of the introduction are $p_i = p_i(1)$. We choose $r$ so that the Poisson sample has the right expected size, $\sum_{i=1}^{N} p_i(r) = n$, and call this $r$ the contour radius, for a reason explained in Appendix A. If each item were included independently with probability $p_i(r)$, we would get

$$\pip_i \;\approx\; p_i(r) = \frac{\w_i \, r}{1 + \w_i \, r}.$$

The actual inclusion probabilities are close but not identical—Hájek (1964, Theorem 5.2) showed that $\pip_i = p_i(r) + \mathcal{O}(1/N)$. This means initializing $\theta_i^{(0)} = \text{logit}(\piptgt_i)$ gives a warm start that is off by at most $\mathcal{O}(1/N)$ for the fitting optimizer.

The contour radius in computation. The same $r$ keeps both algorithms numerically stable. The coefficients of $\prod_{i=1}^{N}(1 + \w_i \z)$ can be too large or too small for floating-point numbers, and their sizes can differ by hundreds of orders of magnitude: the DP's table can overflow, and the FFT, which rounds every coefficient relative to the largest one, can lose $\Z$ in the rounding error. Rescaling the weights by $r$ turns the coefficients, up to a common factor, into the distribution of the size of a Poisson sample with coins $p_i(r)$, and puts the coefficient we need, at $\z^n$, at the peak of that distribution.✓ Both of our implementations do this; Appendix A explains how, and why it works.

Fitting Weights to Target Probabilities

A common use case: you know the inclusion probabilities you want and need to find weights that produce them.✓

Objective. Given target inclusion probabilities $\bpiptgt$, we find weights $\bw$ (equivalently, log-weights $\btheta$ where $\theta_i = \log \w_i$) by maximizing

$$ \ell(\btheta) \defeq \bpiptgt^{\top} \btheta - \log \Zw{\bw}{n} $$

This is concave, so any local maximum is global.

Gradient. $\nabla_{\btheta} \ell(\btheta) = \bpiptgt - \bpip(\btheta)$.✓ At the optimum, $\bpip(\btheta) = \bpiptgt$ exactly, so the gradient is zero. Each evaluation of $\ell(\btheta)$ and $\nabla \ell(\btheta)$ costs $\mathcal{O}(N \log^2 n)$: one pass through the product tree + backpropagation.

Optimizer. We minimize $-\ell(\btheta)$ with L-BFGS, which needs only $\ell(\btheta)$ and its gradient—no Hessian. The Poisson approximation provides a warm start, $\theta_i^{(0)} = \text{logit}(\piptgt_i)$, which has initialization error $\mathcal{O}(1/N)$ per item. It also provides the starting estimate of the curvature: if each item were included independently with probability $\piptgt_i$, the normalizer would be $\prod_{i=1}^{N} (1 + \w_i)$, and the Hessian of its log with respect to $\btheta$, at $\btheta^{(0)}$, would be the diagonal matrix with entries $\piptgt_i (1 - \piptgt_i)$.✓ L-BFGS starts from the inverse of this matrix instead of the identity.

Line search. Near the optimum, a step improves $\ell(\btheta)$ by roughly the square of the gradient, which soon falls below the rounding error in $\ell(\btheta)$ itself (and that error grows with $N$). A line search that compares values of $\ell(\btheta)$ then rejects good steps and stalls with an error around $10^{-8}$. The gradient stays accurate to nearly machine precision, so our line search also accepts a step that meets the approximate Wolfe conditions of Hager & Zhang (2005), which replace the comparison of values of $\ell(\btheta)$ with a test on the slope along the search direction. With it, fit reaches its default tolerance, $\max_i |\pip_i - \piptgt_i| \le 10^{-10}$, even at $N = 1000$.✓

from conditional_poisson import ConditionalPoissonNumPy

# Target: the inclusion probabilities from our earlier example
pi_star = cp.incl_prob.numpy()

# Fit from scratch
cp_fit = ConditionalPoissonNumPy.fit(pi_star, n=3, verbose=True)

max|pi-pi| = 3.890e-02
  max|pi-pi| = 8.576e-03
  max|pi-pi| = 2.961e-04
  max|pi-pi| = 3.307e-05
  max|pi-pi| = 3.495e-06
  max|pi-pi| = 2.310e-07
  max|pi-pi| = 2.619e-08
  max|pi-pi| = 7.287e-10
  max|pi-pi*| = 8.304e-11

Comparison with R. The R sampling package solves the same problem with UPMEpiktildefrompik. It works with Poisson inclusion probabilities $\tilde{p}_i$ (weights $\w_i = \tilde{p}_i / (1 - \tilde{p}_i)$), starts from $\tilde{p}_i = \piptgt_i$, and iterates $\tilde{p}_i \leftarrow \tilde{p}_i + \piptgt_i - \pip_i$, where $\pip_i$ is the inclusion probability under the current weights. Its documentation calls this Newton's method, but the update uses the identity matrix in place of the Jacobian, so it is a fixed-point iteration. Each iteration recomputes $\bpip$ with an $\mathcal{O}(Nn)$ DP.

The identity is what the Jacobian would be if the Poisson approximation were exact: the update keeps $\sum_{i=1}^{N} \tilde{p}_i = n$, so the contour radius stays $r = 1$ and $\tilde{p}_i = p_i(r)$,✓ and the approximation reads $\pip_i \approx \tilde{p}_i$. Our L-BFGS starts from the same approximation, but it refines its estimate of the curvature from the gradients it sees, while the fixed-point iteration keeps the approximation throughout. To compare the two iterations, we run a copy of R's update with our exact $\bpip$ (R's own code takes the same number of steps on the targets we checked) and count evaluations of $\bpip$ until $\max_i |\pip_i - \piptgt_i| \le 10^{-10}$.✓ On the earlier example and on $\text{Exp}(1)$ weights with $N = 10, 50, 200, 1000$ and $n = 0.4N$, the counts are close: L-BFGS needs 3 to 9 evaluations and the fixed-point iteration 4 to 15. When the weights are very uneven, the fixed-point iteration slows down: on 100 random targets with $\log \w_i \sim \mathcal{N}(0, 10^2)$ and $N \le 300$, it needed a median of 32.5 evaluations, and 3 targets (all with $N = 3$) were still short of the tolerance after 2,000, while L-BFGS never needed more than 23. The raw data is in fit_comparison.json, and the script is bench_fit.py.

Timing

Do the asymptotic claims hold up on a real machine? This section times the product tree against the $\mathcal{O}(Nn)$ Pascal DP and against the existing R implementation.

Setup. The product tree is ConditionalPoissonTorch (PyTorch, FFT, $\bpip$ by autograd). It batches the multiplications at each level of the tree into one operation, so a call makes $\mathcal{O}(\log N)$ tensor operations rather than $\mathcal{O}(N)$. The DP is ConditionalPoissonSequentialNumPy: a loop over the $N$ items whose steps are vectorized over $k$, with $\bpip$ computed by a hand-written reverse pass. Both rescale the weights by the contour radius (Appendix A). The R baseline is the sampling package (version 2.11). Weights are drawn from $\text{Exp}(1)$, and every method gets the same weights. Measurements run one at a time on a single thread of an Intel Core i7-8650U laptop CPU, and each time is the fastest of up to five repeats after a warm-up call.✓ Every method computes exact quantities, and the benchmark checks each output against an independent method:✓ the tree and the DP agree on $\bpip$ to within $2 \times 10^{-15}$ and on $\log \Z$ to within a relative $4 \times 10^{-16}$ at every size, and R agrees with both wherever it runs. The raw data is in timing.json, and the benchmark script is bench_timing.py.

The plots below fix $n = 0.4N$; Appendix B varies $N$ and $n$ separately. Hover over a plot for exact times, drag to zoom, double-click to zoom back out, and click a legend entry to hide or show that method.

Computing $\Z$

✓ product tree $\mathcal{O}(N \log^2 n)$ · ✓ DP $\mathcal{O}(Nn)$

The DP is faster for small $N$. The tree's PyTorch implementation has a fixed cost of about 1.3 ms per call (it dispatches about a hundred small tensor operations), and that cost dominates until $N \approx 500$. Beyond that the tree wins, and its lead grows with $N$: at $N = 20{,}000$ it takes 9.0 ms, versus 234 ms for the DP. (When it computes $\Z$ alone, the DP keeps only one row of its table, so it is not slowed down by memory it does not need.)

Computing $\bpip$

This is the main event. Without backpropagation, the natural way to get $\bpip$ is leave-one-out, $\pip_i = \w_i \llbracket \prod_{j \neq i}(1 + \w_j \z) \rrbracket(\z^{n-1}) / \Zw{\bw}{n}$,✓ which takes one forward pass per item.✓ Backpropagation replaces those $N$ forward passes with a single backward pass.

✓ product tree + autograd $\mathcal{O}(N \log^2 n)$ · ✓ DP + reverse pass $\mathcal{O}(Nn)$ · ✓ leave-one-out with either · ✓ R sampling package

Backpropagation costs a small constant factor more than the forward pass, as the Baur-Strassen theorem promises: computing $\bpip$ takes 2.0 to 2.3 times as long as computing $\Z$ with the tree, and 2.1 to 3.1 times as long with the DP. Leave-one-out costs a factor of about $N$ more instead. At $N = 2000$, it takes 4.7 s with the tree (880 times as long as backpropagation) and 18 s with the DP (790 times as long). Between the two backpropagation methods, the tree overtakes the DP by $N = 500$ and is 35 times faster at $N = 20{,}000$ (21 ms versus 734 ms).

The R package also uses an $\mathcal{O}(Nn)$ DP, but it runs as interpreted R code: at $N = 1000$ it takes 242 ms, 54 times as long as the tree. Its DP recurrence also works in linear space without rescaling, so at $N = 2000$ ($n = 800$) the table overflows and it returns NaN. (Our DP rescales by the contour radius, which keeps every entry of its table a probability; see Appendix A.) The R package does not expose $\Z$.

Drawing Samples

✓ product tree $\mathcal{O}(n \log N)$ per sample · ✓ DP table $\mathcal{O}(N)$ per sample · ✓ R sampling package $\mathcal{O}(N)$ per sample

Both samplers first build a data structure, which costs about one forward pass: the tree reuses its product tree, and the DP fills its $(N+1) \times (n+1)$ table. The plot shows the cost of each sample after that. With $n = 0.4N$, the DP's sampler is faster: its $\mathcal{O}(N)$ scan beats the tree's $\mathcal{O}(n \log N)$ walk when $n$ is a constant fraction of $N$ (at $N = 20{,}000$, 8.5 ms versus 55 ms per sample). The tree's sampler wins when the sample is a small fraction of the population: in the grid of Appendix B, it is faster whenever $n$ is at most about 2% of $N$, for example 0.13 ms versus 5.4 ms at $N = 20{,}000$ and $n = 5$. R's sampler is also an $\mathcal{O}(N)$ scan, 2.4 to 2.5 times slower than ours.

In short. For $\Z$ and $\bpip$, the product tree is faster than the DP once $N$ reaches a few hundred, and its lead grows with $N$. For sampling, the DP's scan is faster unless the sample is a small fraction of the population. Appendix B varies $N$ and $n$ separately and checks the measured growth rates against the complexity bounds.

Summary

The takeaway. A single data structure—the polynomial product tree—unifies normalizing constant computation, marginal inference, sampling, and parameter fitting for the conditional Poisson distribution. Each capability is a mechanical program transformation of the previous one:

There are no problem-specific derivations here—each row follows from a general theorem in automatic differentiation or computer algebra.

Acknowledgment. This post was written with extensive help from Claude Code—especially the interactive widgets, which would not have been possible without it.

Found an error or have a suggestion? Please open an issue on GitHub.

Appendix A: Numerical Stability

Floating-point numbers carry about 16 significant digits and range from about $10^{-308}$ to $10^{308}$. Both algorithms compute coefficients of $\prod_{i=1}^{N}(1 + \w_i \z)$, which can fall outside that range and whose sizes can differ by hundreds of orders of magnitude. This appendix shows how one rescaling, by the contour radius, keeps both algorithms accurate.

Rescaling by the Contour Radius

The key identity. For any $r > 0$:✓

$$\Zw{\bw}{n} = r^{-n} \cdot \llbracket \textstyle\prod_{i=1}^{N}(1 + \w_i r\, \z) \rrbracket(\z^n)$$

This is immediate: each $\z^n$ term in the product picks up a factor of $r$ per item, giving $r^n \cdot \Zw{\bw}{n}$. In exact arithmetic, $r$ doesn't matter. In floating point, it matters enormously, because it changes which coefficient is the largest.

The coefficients are probabilities. Let $K$ be the number of heads when each coin $i$ comes up heads, independently, with probability $p_i(r) = \w_i r / (1 + \w_i r)$. Since $(1 + \w_i r \z)/(1 + \w_i r) = (1 - p_i(r)) + p_i(r)\, \z$ is the generating function of coin $i$,

$$\llbracket \textstyle\prod_{i=1}^{N}(1 + \w_i r\, \z) \rrbracket(\z^k) = \Pr[K = k] \cdot \prod_{i=1}^{N}(1 + \w_i r)$$

for every $k$.✓ So, up to a factor that does not depend on $k$, the rescaled coefficients are the distribution of $K$.

Choosing $r$. The contour radius sets the mean of $K$ to $n$:

$$\mathbb{E}[K] = \sum_{i=1}^{N} \frac{\w_i r}{1 + \w_i r} = \sum_{i=1}^{N} p_i(r) = n$$

For a sum of independent Bernoullis, the mode is within 1 of the mean, and when the mean is an integer, the mode equals it (Darroch, 1964, Theorem 4). So $\Pr[K = n]$ is the largest of the $N + 1$ probabilities $\Pr[K = 0], \ldots, \Pr[K = N]$, and in particular $\Pr[K = n] \ge 1/(N+1)$.✓ The coefficient we need is never small next to the others, and both algorithms below rely on this. It is also why the contour radius is the $r$ of the Poisson approximation: both ask for the rescaling that centers the distribution of $K$ at $n$.

The left side increases with $r$, from 0 to $N$, so Newton's method in $\log r$, safeguarded by bisection, finds the root in a few iterations. A simpler heuristic, $r = n / \W$ (where $\W \defeq \sum_{i=1}^{N} \w_i$), linearizes $\w_i r / (1 + \w_i r) \approx \w_i r$.

The Product Tree

An FFT rounds every coefficient of a product to within about $\varepsilon \cdot \max_k|c_k|$, where $c_k$ denotes the $k$^th coefficient of the product and $\varepsilon \approx 10^{-16}$ is machine epsilon. For $\prod_{i=1}^{N}(1 + \w_i \z) = \sum_{k=0}^{N} c_k \z^k$ without rescaling, the largest coefficient can be $\approx 10^{300}$ times larger than $c_n$, the coefficient we need—so FFT noise drowns the signal.

After rescaling by $r$, the $k$^th coefficient becomes $c_k(r) = \Zw{\bw}{k} \cdot r^k$, which is proportional to $\Pr[K = k]$. With the contour radius, the mode of $K$ is $n$, so among the coefficients the tree keeps (degrees up to $n$), $c_n(r)$ is the largest: the dynamic range $\max_{k \le n} c_k(r) / c_n(r)$ is exactly 1,✓ and the rounding error of the FFT at the root is relative to $c_n(r)$ itself.

Rescaling does not prevent overflow: $\prod_{i=1}^{N}(1 + \w_i r)$ can still exceed $10^{308}$. So the tree also divides each node's polynomial by its largest coefficient and adds up the logs of the divisors.✓

Why the name. An FFT of the rescaled coefficients evaluates $f(\z) = \prod_{i=1}^{N}(1 + \w_i \z)$ at equally spaced points on the circle $|\z| = r$, so $r$ is the radius of the circle—the contour—on which the tree's multiplications happen.

Autograd compatibility. The rescaling $\w_i \mapsto \w_i \cdot r$ is on the autograd graph (it's just a scalar multiply); the root-finding for $r$ is not (it's a numerical conditioning choice, not part of the mathematical function). Gradients flow through $\w_i \cdot r$ as if $r$ were a constant—which is correct, since $\log \Z$ does not depend on $r$ (every $r$ gives the same answer in exact arithmetic). The per-node divisors are treated the same way. Griewank & Walther (2008) take "what's good for function values is good for their derivatives" as a working hypothesis (Rule 2) and back it with a backward error analysis (Rule 3, Section 3.4): if nothing overflows or underflows, reverse mode returns the exact gradient of the same computation with each elementary operation perturbed at the level of machine precision. So a stable forward pass does most of the work.

With the contour radius, the tree matches exact rational arithmetic for weights spanning $e^{\pm 300}$.✓ Code: conditional_poisson/tree_torch.py

The Dynamic Program

The DP needs no FFT, but its table holds $\Zw{\w_1, \ldots, \w_m}{k}$, which can be too large or too small for floating-point numbers: R's sampling package runs an unscaled $\mathcal{O}(Nn)$ recurrence and returns NaN at $N = 2000$, $n = 800$ in the Timing section. Our DP rescales the weights by the contour radius and divides row $m$ by $\prod_{i=1}^{m}(1 + \w_i r)$. Row $m$ then holds the distribution of $K_m$, the number of heads among the first $m$ coins, for $k \le n$:

$$\Pr[K_m = k] = \frac{\Zw{\w_1 r, \ldots, \w_m r}{k}}{\prod_{i=1}^{m}(1 + \w_i r)}$$

and the weighted Pascal recurrence becomes

$$\Pr[K_m = k] = \big(1 - p_m(r)\big) \Pr[K_{m-1} = k] + p_m(r) \Pr[K_{m-1} = k-1]$$

with base cases $\Pr[K_0 = 0] = 1$, $\Pr[K_0 = k] = 0$ for $k > 0$, and $\Pr[K_{m-1} = -1] = 0$: coin $m$ comes up tails (first term) or heads (second term).✓ At the end,

$$\log \Zw{\bw}{n} = \log \Pr[K_N = n] + \sum_{i=1}^{N} \log(1 + \w_i r) - n \log r.$$

This form of the recurrence is safe in three ways:✓

• No overflow. Each entry is a weighted average of two entries of the previous row, so every entry is a probability, between 0 and 1.
• No harmful underflow. The answer, $\Pr[K_N = n]$, is the mode of $K_N$, so it is at least $1/(N+1)$. Entries far from the mean can underflow, but each one loses less than $10^{-307}$, and it reaches the answer only through multiplications by probabilities, so it changes the answer by less than that.
• No cancellation. The recurrence only adds nonnegative numbers, and $1 - p_m(r)$ is computed as $1/(1 + \w_m r)$ rather than by subtraction, so rounding errors stay small relative to each entry.

The sampler and the reverse pass that computes $\bpip$ read the same table. Both use only ratios of entries, in which the rescaling cancels. For example, the sampler includes item $m$ with probability $p_m(r) \Pr[K_{m-1} = k-1] / \Pr[K_m = k]$, which equals the probability in the DP's sampler. The reverse pass also carries a second vector, the distribution of the number of heads among the remaining coins, which stays between 0 and 1 for the same reason.✓

With weights spanning $e^{\pm 300}$, the unscaled recurrence overflows, while the scaled DP agrees with 50-digit arithmetic on $\log \Z$ and $\bpip$ to within $10^{-13}$, even though more than a quarter of its table underflows to zero.✓

The same trick—dividing each step of a recurrence by a scale factor and adding up the logs—keeps the forward algorithm for hidden Markov models from underflowing (Rabiner, 1989, Section V.A). There, each row is divided by its sum, which must be computed along the way; here, the divisors $1 + \w_m r$ are known in advance, and the contour radius guarantees that the entry we need is the largest in the last row.

Appendix B: Timing in $N$ and $n$

The Timing section fixes $n = 0.4N$. The 3D plots below vary $N$ from 100 to 20,000 and $n$ from 5 to $N/2$, with the same setup. Drag to rotate, scroll to zoom, and hover for values. For $\Z$ and $\bpip$, the DP is faster at every $n$ when $N \le 200$, and the tree is faster at every $n$ when $N \ge 1000$. The crossover depends on $N$ more than on $n$ because of constant factors: the DP pays a few microseconds of NumPy overhead per item, and the tree pays about 1.3 ms per call. For sampling, the tree's sampler is faster only when $n$ is at most about 2% of $N$.

To check the claimed complexities, the table below fits the log-log slope of time against $N$ (at fixed $n$) and against $n$ (at fixed $N$), each over the four largest sizes, and compares it with the slope of the complexity bound over the same points. It is computed from timing.json by the script that draws the plots.✓ Cells where the two differ by more than 0.2 are flagged in red.

In this run the flags have three causes, none of which contradicts the bounds. First, the tree's fixed cost of about 1.3 ms per call is still about half of its time at $N = 2000$, which flattens its slope in $N$. Second, each step of the DP is a few NumPy calls on vectors of length $n$, and until $n$ reaches the thousands, their fixed overhead outweighs the $n$ multiply-adds, so the DP's time grows more slowly than $n$. Third, the tree's sampler visits about $n \log_2(N/n) + 2n$ nodes, which is $\mathcal{O}(n \log N)$ but has different local slopes than $n \log N$. Away from these effects, the measurements agree with the bounds: the DP's time grows linearly in $N$, and the tree's grows no faster than $\log^2 n$ in $n$.

Appendix C: Identities for $\Z$ and Its Relatives

The normalizing constant $\Zw{\bw}{n}$ is the $n$^th elementary symmetric polynomial $e_n(\bw)$. Here are some useful identities.

Differential Identities

The exponential family structure gives the inclusion probability as the gradient of the log-normalizer:

$$\pip_i = \frac{\partial \log \Z}{\partial \theta_i} = \frac{\w_i \cdot \Zw{\bw^{(-i)}}{n-1}}{\Zw{\bw}{n}}$$ ✓

The leave-one-out formula generalizes to higher-order inclusion probabilities: $\pip(X) = P(X \subseteq S) = \prod_{i \in X} \w_i \cdot \Zw{\bw^{(-X)}}{n-|X|} / \Zw{\bw}{n}$.✓

Recurrences and Algorithms

Proposition (Weighted Vandermonde identity). For disjoint groups $A, B$ with weight vectors $\ba$ and $\bb$:

$$\Zw{(\ba;\, \bb)}{k} = \sum_{j=0}^{k} \Zw{\ba}{j} \cdot \Zw{\bb}{k-j}$$

This is a weighted generalization of Vandermonde's identity $\binom{a+b}{k} = \sum_{j=0}^{k} \binom{a}{j}\binom{b}{k-j}$.✓ The proof is immediate from the generating function: $\Zw{\bw}{k}$ is the $k$^th coefficient of $\prod_{i=1}^{N}(1 + \w_i \z)$, and factoring this product over $A \cup B$ turns coefficient extraction into a convolution.

With $B$ a single item, the sum has only two nonzero terms, and the identity reduces to the weighted Pascal recurrence.✓ This is why polynomial multiplication computes $\Z$: the product tree exploits this identity at every node. It is also what makes the sampling algorithm correct—splitting a quota of $k$ items between two subtrees according to $\Zw{\ba}{j} \cdot \Zw{\bb}{k-j}$ produces the exact conditional distribution.

Newton's identities. The elementary symmetric polynomials can be computed from power sums $g_k \defeq \sum_{i=1}^{N} \w_i^k$ (see Stanley (1999)):

$$\Zw{\bw}{k} = \sum_{i=1}^{k} \frac{(-1)^{i-1}}{k}\, \Zw{\bw}{k-i} \cdot g_i$$

✓ This is an $\mathcal{O}(Nn)$ algorithm that only needs the power sums, not the individual weights—useful when the universe is implicitly defined (e.g., paths in a weighted finite-state automaton, where $g_k$ can be computed via matrix methods).

Connection to $k$-DPPs

References

• Baur & Strassen (1983). "The Complexity of Partial Derivatives." Theoretical Computer Science, 22(3).

• Chen, Dempster & Liu (1994). "Weighted Finite Population Sampling to Maximize Entropy." Biometrika, 81(3).

• Chen & Liu (1997). "Statistical Applications of the Poisson-Binomial and Conditional Bernoulli Distributions." Statistica Sinica, 7(4).

• Darroch (1964). "On the Distribution of the Number of Successes in Independent Trials." The Annals of Mathematical Statistics, 35(3).

• Griewank & Walther (2008). Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, 2nd edition. SIAM.

• Hager & Zhang (2005). "A New Conjugate Gradient Method with Guaranteed Descent and an Efficient Line Search." SIAM Journal on Optimization, 16(1).

• Hájek (1964). "Asymptotic Theory of Rejective Sampling with Varying Probabilities from a Finite Population." The Annals of Mathematical Statistics, 35(4).

• Jaynes (1957). "Information Theory and Statistical Mechanics." Physical Review, 106(4).

• Kulesza & Taskar (2012). "Determinantal Point Processes for Machine Learning." Foundations and Trends in Machine Learning, 5(2–3).

• Meister, Amini, Vieira & Cotterell (2021). "Conditional Poisson Stochastic Beams." Proceedings of EMNLP 2021.

• Rabiner (1989). "A Tutorial on Hidden Markov Models and Selected Applications in Speech Recognition." Proceedings of the IEEE, 77(2).

• Stanley (1999). Enumerative Combinatorics, Volume 2. Cambridge University Press.

• Tillé (2006). Sampling Algorithms. Springer.

• von zur Gathen & Gerhard (2013). Modern Computer Algebra, 3rd edition. Cambridge University Press.

Citing this post:

@misc{vieira2026cps,
  author = {Vieira, Tim},
  title  = {Conditional Poisson Sampling},
  year   = {2026},
  url    = {https://timvieira.github.io/conditional-poisson-sampling/}
}