Ideas for Testing Randomized Algorithms A new blog post outlines testing strategies for randomized algorithms, using the AWRS sampling algorithm from Lipkin et al. (2025, arXiv:2504.05410) as a worked example for constrained generation in sequential Monte Carlo (SMC) with language models. The post presents four Python implementations of AWRS — variants A, B and C among them — at least one of which contains a bug, and develops progressively stronger tests, each catching a bug the previous test missed, before plugging the block into a full SMC pipeline. The author exploits a swappable `draw` parameter, which defaults to categorical sampling, to enumerate every possible execution during testing. Randomized algorithms are hard to test. Their outputs are, by design, unpredictable—and that unpredictability is a feature , not a bug. It is what gives them useful guarantees like unbiasedness and computational efficiency. But it also makes them slippery: how do you check that a stochastic program is correct when it gives different answers every time you run it? This post sketches some ideas for how to think about testing randomized algorithms—especially the kinds of algorithms that we use for probabilistic inference. I'll present a concrete sampling algorithm with several plausible implementations—some correct, some subtly broken—and develop progressively more powerful testing strategies to tell them apart. Each strategy catches a bug that the previous one missed. The algorithm is a small building block. Its job is to propose the next token inside sequential Monte Carlo SMC for constrained generation from language models. Testing the block in isolation is only half the job, so at the end we'll plug it into SMC and test the whole pipeline. A Tricky Little Sampling Algorithm Suppose we have a distribution $p$ over a finite set $\mathcal{X} = \{1, \ldots, N\}$ and a condition $C \subseteq \mathcal{X}$ the set of "good" elements . We want to sample from the conditional distribution $p \cdot \mid C $. But rather than returning a plain sample, we want a properly weighted sample $ x, w $—a sample-weight pair that can be used for unbiased estimation. The following algorithm, which we call AWRS from Lipkin et al., 2025 https://arxiv.org/abs/2504.05410 , does exactly this. Here is the pseudocode from the paper: The idea: run rejection sampling twice. The first pass finds a conditioned sample $x c$ while tracking the total rejected probability mass $\psi 0$. The second pass generates additional rejections to get a count $n$. The weight $ 1 - \psi 0 / n + 1 $ combines these two signals—the rejected mass and the rejection count—into an unbiased estimate of the normalizing constant $Z$. The number of rejections before finding a valid token contains signal about the magnitude of $Z$, and weighting by the actual probability mass removed rather than just counting reduces variance. Four Implementations, at Least One Bug Did we implement it correctly? That's the question we'd like our tests to answer. To make the exercise concrete, here are four Python implementations. They all look reasonable, but at least one has a bug. Which ones are correct? Each variant takes a draw parameter that defaults to sampling from a categorical distribution. This is a hook we'll exploit later: in production it calls the RNG, but in testing we'll swap it for something that lets us enumerate every possible execution. Variant A python def awrs a p, cond, draw=sample categorical : p cond = p.copy n, psi = 0, 0 xc0 = None for i in range 2 : while True: q = p cond / p cond.sum x = draw q, context= 'awrs', i, xc0 if cond x : if i == 0: xc0 = x break else: if i == 0: psi += p cond x p cond x = 0 n += 1 return xc0, 1 - psi / n + 1 Variant B python def awrs b p, cond, draw=sample categorical : p cond = p.copy n, psi = 0, 0 xc0 = None for i in range 2 : while True: q = p cond / p cond.sum x = draw q, context= 'awrs', i, xc0 if cond x : if i == 0: xc0 = x break else: if i == 0: psi += p cond x p cond x = 0 n += 1 return xc0, 1 + psi / n + 1 Variant C python def awrs c p, cond, draw=sample categorical : p cond = p.copy n, psi = 0, 0 xc0 = None for i in range 2 : while True: q = p cond / p cond.sum x = draw q, context= 'awrs', i, xc0 if cond x : if i == 0: xc0 = x break else: psi += p cond x p cond x = 0 n += 1 return xc0, 1 - psi / n + 1 Variant D python def awrs d p, cond, draw=sample categorical : p cond = p.copy n, psi = 0, 0 xc0 = None for i in range 2 : while True: q = p cond / p cond.sum x = draw q, context= 'awrs', i, xc0 if cond x : if i == 0: xc0 = x break else: if i == 0: psi += q x p cond x = 0 n += 1 return xc0, 1 - psi / n + 1 The variants differ from one another by a line or two. I encourage you to read them carefully and try to spot the differences before scrolling further. Hint Compare the else branch and the return statement across the four variants. Our job is to build testing strategies that catch the buggy ones—ideally without needing to know which line is wrong. But before we can test anything , we need to answer the most important question in all of software testing: What Should This Code Do? You cannot test code without a specification. For randomized algorithms, the specification comes from theory—and translating theory into a checkable property is half the battle. We'll keep coming back to this question for every algorithm in this post. For AWRS, the specification is: it should return a sample-weight pair $ x, w $ such that $w \cdot \mathbf{1} x = x' $ is an unbiased estimator of the unnormalized conditional distribution for every $x' \in \mathcal{X}$: $$ \mathop{\mathbb{E}} { x,w \sim \text{AWRS}} \left w \cdot \mathbf{1} x = x' \right = p x' \cdot \mathbf{1} x' \in C $$ The right-hand side is the unnormalized target—$p$ restricted to the condition $C$, with normalizing constant $Z = \sum {x \in C} p x $. That AWRS satisfies this is a theorem in Lipkin et al., 2025 https://arxiv.org/abs/2504.05410 , and we'll take it as given. Our question is whether the code does what the paper proves. Note that you may often need to tweak definitions, narrow or broaden the scope of theory in order to suit your code's intended use cases. This is an important skill for turning theory into tests. Setting Up a Test Case Now that we have a testable definition, we need to think about which things in the equation we can compute and which things we need to estimate. Let's make a concrete test case: p = np.array 0.5, 0.3, 0.15, 0.04, 0.01 a distribution with a long tail cond = np.array True, False, True, False, False the "good" elements N = len p Z = p cond .sum target = p cond unnormalized target not a distribution Strategy 1: Monte Carlo Averaging The obvious first testing strategy: run the code "a lot of times" and compute an average. np.random.seed 0 reps = 1 000 for name, awrs in variants.items : have = np.zeros N for r in range 1, 1 + reps : x, w = awrs p, lambda x: cond x have x += w err = abs have / reps - target .sum print f'Variant {name}: error = {err:.6f}' Variant A: error = 0.023367 Variant B: error = 0.098657 Variant C: error = 0.030520 Variant D: error = 0.012054 Is this error small enough? How do you know? Variant D has the smallest error of the four. This is the central problem with Monte Carlo testing: a single number tells you almost nothing without a reference frame. We need to know what should happen as we increase the number of samples. Convergence Rates A better approach is to watch the error decrease as a function of the number of samples $M$. It should keep going down forever , modulo numerical precision. And crucially, we often know the rate at which it should decrease. The general picture: suppose we have a sequence of estimators $\widehat{\mu}^{ M }$ targeting some quantity $\mu$. The mean squared error decomposes as $$ \mathop{\mathbb{E}} \| \widehat{\mu}^{ M } - \mu \|^2 = \| \mathop{\mathbb{E}} \widehat{\mu}^{ M } - \mu \|^2 + \mathop{\mathbb{E}} \| \widehat{\mu}^{ M } - \mathop{\mathbb{E}} \widehat{\mu}^{ M } \|^2 $$ The first term is the squared bias; the second is the variance more precisely, the trace of the covariance of $\widehat{\mu}^{ M }$ . Different classes of estimators make different predictions about how these terms behave: - I.I.D. unbiased e.g., our AWRS test . The bias is zero for all $M$. The variance is $V/M$ where $V = \text{tr} \text{Cov} X $—the sum of the componentwise variances of a single sample. So the mean squared error is exactly $V/M$.If you prefer root mean squared error, the slope is $-1/2$ and the intercept is $\frac{1}{2} \log V$. - Correlated but unbiased e.g., sampling without replacement . The bias is still zero. Samples drawn without replacement are negatively correlated, which shrinks the variance: the rate is still $O 1/M $, but the constant is smaller. - Biased but consistent e.g., self-normalized importance sampling, MCMC . The bias is $O 1/M $, so the squared bias is $O 1/M^2 $. The variance is $O 1/M $. The MSE is still $O 1/M $ overall, so on a log–log plot the slope looks the same—the bias hides in the noise at large $M$. In all three cases, $\text{MSE} \sim c / M$ for some constant $c$, which means $\log \text{MSE} = -\log M + \log c$: a line with slope $-1$ on a log–log plot. The intercept $\log c$ tells us the constant. Deriving $\,V\,$ for vector-valued estimators Our estimator is vector-valued a distribution over $N$ elements , so we need a scalar measure of error. The natural choice is the squared $L 2$ norm $\| \widehat{\mu}^{ M } - \mu \|^2$ where $\widehat{\mu}^{ M } = \frac{1}{M} \sum {m=1}^{M} X^{ m }$ and $\mu = \mathop{\mathbb{E}} X^{ m } $: $$ \mathop{\mathbb{E}} \| \widehat{\mu}^{ M } - \mu \|^2 = \mathop{\mathbb{E}} \sum {j=1}^{N} \widehat{\mu}^{ M } j - \mu j ^2 = \sum {j=1}^{N} \text{Var} \widehat{\mu}^{ M } j = \frac{1}{M} \sum {j=1}^{N} \text{Var} X^{ m } j = \frac{V}{M} $$ where $V = \text{tr} \text{Cov} X^{ m } $ is the trace of the covariance matrix of a single sample. The last step uses independence; for correlated samples, the $1/M$ factor becomes $1/M$ times an autocorrelation correction. results = {} for name, awrs in variants.items : np.random.seed 0 have = np.zeros N xs, ys = , for r in range 1, 1 + 1 000 : item, w = awrs p, lambda x: cond x have item += w if r = 10: xs.append r ys.append abs have / r - target .sum slope, intercept = np.polyfit np.log xs , np.log ys , deg=1 results name = xs, ys, slope, intercept plot single runs results A note on reading this plot: when we say "log–log scale" we mean the axes are log-transformed but the tick labels show the original values e.g., $10^2$, $10^3$ . This is what a log axis does in a plotting library. This is different from literally plotting $\log y$ vs $\log x$, where the axes are linear but the values are logs. The line-fitting code does the latter np.polyfit on $\log x$, $\log y$ , but for visual inspection, log-scaled axes are usually clearer because you can read off the original quantities directly. For a correct estimator, the fitted slope should be close to $-0.5$ we're plotting error, not squared error . Variant B fails this test plainly: its error stops going down after a few dozen samples and the slope is positive. B is out. The other three look fine. C's slope is a little shallow, but the slope of a single run is itself a noisy statistic—rerunning the correct variant with other seeds gives slopes anywhere from about $-0.25$ to $-0.75$. D's curve is indistinguishable from A's. That's not a coincidence, and we'll see why later. What a Single Run Can and Can't Tell Us Watching a single long run tells us the rate of convergence, but it can't distinguish the three cases above. An unbiased i.i.d. estimator, an estimator built from samples drawn without replacement, and an MCMC chain can all produce very similar log–log plots—the slopes are the same and the constants differ in ways that are hard to eyeball. A pedagogical example: suppose we accidentally computed the average as $\frac{1}{M+1} \sum {m=1}^{M} X^{ m }$ instead of $\frac{1}{M} \sum {m=1}^{M} X^{ m }$. This estimator is biased —it systematically underestimates the target—but consistent , because the bias factor $M/ M+1 \to 1$. The single-run convergence plot would look nearly identical. The same thing is happening to C. Its bias is real, but at a thousand samples it is still smaller than the noise. A longer run would expose it eventually, but how long is long enough? A single run has no way to tell us. Averaging Over Multiple Runs To separate bias from variance, we need multiple independent runs . The idea: run the estimator many times at each sample size $M$, and compare the average behavior to the theoretical prediction. Recall the decomposition from earlier: $$ \mathop{\mathbb{E}} \| \widehat{\mu}^{ M } - \mu \|^2 = \| \mathop{\mathbb{E}} \widehat{\mu}^{ M } - \mu \|^2 + \mathop{\mathbb{E}} \| \widehat{\mu}^{ M } - \mathop{\mathbb{E}} \widehat{\mu}^{ M } \|^2 $$ For the i.i.d. unbiased case, the variance term is $V/M$. So: - If the observed MSE matches $V/M$, the squared bias must be near zero—evidence of unbiasedness. - If the estimator has persistent bias, $\| \mathop{\mathbb{E}} \widehat{\mu}^{ M } - \mu \|^2$ does not vanish with $M$. At large $M$, the variance $V/M$ shrinks but the bias dominates, causing the observed curve to flatten out above the theoretical line. - For biased-but-consistent estimators like self-normalized IS , the squared bias is $O 1/M^2 $—it vanishes faster than the variance. This means the MSE curve will eventually match $V/M$, but may sit slightly above it for small $M$. results = {} for name, awrs in variants.items : np.random.seed 0 reps, M = 1000, 200 errors = np.zeros reps, M squared error of the running average, per run samples = single-sample estimates, for V for r in range reps : have = np.zeros N for m in range 1, M + 1 : item, w = awrs p, lambda x: cond x have item += w onehot = np.zeros N onehot item = w samples.append onehot errors r, m - 1 = np.linalg.norm have / m - target 2 V = np.trace np.cov np.array samples .T results name = errors, V plot mse results A and D hug the theoretical $V/M$ line, which is strong evidence that they are properly weighted—not just consistent, but unbiased, with the right variance. B flattens out, as before. C is the interesting one: it tracks the line for small $M$ and then lifts off as the variance shrinks below its squared bias. Averaging over a thousand runs makes the gap far larger than the confidence band. C is out. D, on the other hand, has now passed every test we've thrown at it. But Monte Carlo testing is still a blunt instrument . Small biases hide easily in the noise, and you need many runs to have any statistical power. For correlated estimators, the effective variance may be hard to compute, making the theoretical line itself uncertain. In software—unlike the real world—we can almost always do better. Strategy 2: Program Tracing Here's the key insight: if the state space is small enough, we can enumerate all possible executions of the randomized program and verify the correctness condition exactly up to floating point . The idea is to hijack the random number generation inside the program to make it systematically explore every possible branch. We track each execution trace's probability, and combine all traces to compute the exact output distribution. The Tracer The tracing strategy comes from Shi et al., 2020 https://arxiv.org/abs/2002.09067 , who use it to sample without replacement from programs such as sequence models. The same idea underlies exact inference by enumeration in probabilistic programming languages: WebPPL explores every execution of a program by capturing a continuation at each random choice Goodman & Stuhlmüller, 2014 http://dippl.org/chapters/03-enumeration.html . The implementation uses a trie data structure that mirrors the branching structure of the program's random choices. Each node in the trie corresponds to a decision point, and its children correspond to the possible outcomes. Each node stores the probability mass that hasn't been explored below it. A node is created the first time the program reaches it, holding the probability of the path that leads there, so at the end of a trace tracer.cur.mass is that trace's probability; this is how the checks below weight each output. When a trace finishes, its mass is subtracted from every node on the way back up to the root. Choosing children in proportion to their unexplored mass then never revisits a finished trace, and the root's mass is the total probability of the traces not yet run. Here is the tracer at work on variant A with our test case. Each step runs the program once more and adds one trace to the trie. Node colors mark which pass of AWRS made the choice, and a node fades as the probability mass below it is explored. The order menu is explained below which-trace-next . The beautiful thing about TraceSWOR is that it works as a context manager: each with tracer: block explores one trace, and when it exits, that trace's mass is zeroed out so we never revisit it. We loop until the root's mass hits zero—meaning we've exhausted every possible execution. The tracer is also available as a small, tested library, program tracing https://github.com/timvieira/testing-randomized-algorithms , in this post's repository. It includes a traces program helper that runs this loop for you. Tracer implementation python class Node: """A node in the execution trie.""" def init self, idx, mass, parent, token=None : self.idx = idx self.mass = mass remaining probability mass self.parent = parent self.token = token label for visualization self.child masses = None set on first visit self.keys = None outcomes of the choice made here self.active children = {} self. mass = mass remember original mass self.context = None def sample self : """Sample a child proportional to remaining mass.""" p = self.child masses / self.child masses.sum return np.random.choice len p , p=p def update self : """Propagate mass changes up to the root.""" if self.parent is not None: self.parent.child masses self.idx = self.mass self.parent.mass = np.sum self.parent.child masses self.parent.update class Tracer: """ Lazily materializes the probability tree of a generative process by program tracing. Works with both array-valued and dict-valued distributions. """ def init self : self.root = Node idx=-1, mass=1.0, parent=None self.cur = None def call self, p, context=None : """Drop-in replacement for any sampling function.""" if isinstance p, dict : keys = list p.keys vals = np.array p k for k in keys , dtype=float else: p = np.asarray p, dtype=float keys = list range len p vals = p vals = vals / vals.sum cur = self.cur if cur.child masses is None: cur.child masses = cur.mass vals cur.keys = keys cur.context = context elif keys = cur.keys: raise ValueError 'the program made a different choice at the same point ' 'of two traces; it must be deterministic given its draws' a = self.choose cur if a not in cur.active children: cur.active children a = Node idx=a, mass=cur.child masses a , parent=cur, token=keys a , self.cur = cur.active children a return keys a def choose self, node : """Pick which child to explore next.""" return node.sample class TraceSWOR Tracer : """ Sampling without replacement meets program tracing. Each execution trace is explored exactly once. After a trace completes, its probability mass is zeroed out and propagated up the trie, ensuring the next trace is fresh. """ def enter self : self.cur = self.root return self def exit self, args : self.cur.mass = 0 self.cur.update Exact Verification Now we can verify all four AWRS variants exactly : python def exact moments awrs : """Exact first and second moments of w onehot x , over all traces.""" tracer = TraceSWOR m1, m2 = np.zeros N , np.zeros N while tracer.root.mass 0: with tracer: item, w = awrs p, lambda x: cond x , draw=tracer m1 item += w tracer.cur.mass m2 item += w 2 tracer.cur.mass return m1, m2 for name, awrs in variants.items : prob, = exact moments awrs err = abs target - prob .sum status = '\u2713' if np.allclose target, prob else '\u2717' print f'Variant {name}: {status} error = {err:.2e} ' Variant A: ✓ error = 3.61e-16 Variant B: ✗ error = 9.08e-02 Variant C: ✗ error = 3.15e-02 Variant D: ✗ error = 1.61e-04 This is not a statistical test. It is an exact check modulo floating point that the proper weighting condition holds.The difference between target and prob is not zero—it rarely will be with floating point numbers. But np.allclose checks that they agree to within a reasonable tolerance $\approx 10^{-8}$ . No confidence intervals, no squinting at log–log plots. The tracer gives a definitive yes or no for each variant. D is out , and A is the only survivor. D's bug is in the else branch: it accumulates the normalized probability q x instead of the remaining mass p cond x . The two agree on the first rejection, because p cond still sums to one. They differ only when the first pass rejects two or more elements. With our long-tailed p , that happens in about one run in twenty, and even then the weight is off by only a couple of percent ✓ https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test article.py L85 —which is why D's single run looked just like A's. How Many Samples Would Monte Carlo Need? The tracer computes more than a pass/fail bit. Collecting the second moment alongside the first gives the exact squared bias and the exact variance $V$ of each variant. The bias of the average stays fixed as $M$ grows while its variance shrinks like $V/M$, so the two cross at $M^ = V / \text{bias}^2$. Below $M^ $, the bug is buried in the noise. for name, awrs in variants.items : m1, m2 = exact moments awrs bias2 = np.sum m1 - target 2 V = np.sum m2 - m1 2 if np.isclose bias2, 0, atol=1e-20 : print f'Variant {name}: unbiased' else: print f'Variant {name}: bias^2 = {bias2:.1e}, V = {V:.2f}, ' f'M = {V / bias2:,.0f}' Variant A: unbiased Variant B: bias^2 = 5.3e-03, V = 0.25, M = 47 Variant C: bias^2 = 6.4e-04, V = 0.24, M = 382 Variant D: bias^2 = 1.7e-08, V = 0.24, M = 14,544,122 This is the ladder we just climbed. B surfaces after a few dozen samples, so a single run catches it. C needs a few hundred, so it takes careful averaging. D needs about fifteen million, and no reasonable Monte Carlo test would ever see it. The tracer took a fraction of a second. Comparing Variances The same moments let us compare estimators that are all correct. Unbiasedness is only half the story: AWRS exists because it has lower variance than simpler alternatives. The simplest properly weighted sampler draws once from $p$ and keeps the draw with weight 1 if it lands in $C$, or weight 0 otherwise. It is unbiased, so no correctness test can tell it apart from AWRS. The tracer can compare their variances exactly: python def naive p, cond, draw=sample categorical : """Draw once from p; weight 1 if the draw satisfies the condition, else 0.""" x = draw p return x, float cond x for name, sampler in 'AWRS', awrs a , 'naive', naive : m1, m2 = exact moments sampler bias2 = np.sum m1 - target 2 V = np.sum m2 - m1 2 var w = m2.sum - m1.sum 2 variance of w as an estimate of Z print f'{name: 5}: bias^2 = {bias2:.1e}, V = {V:.3f}, Var w = {var w:.3f}' AWRS: bias^2 = 6.5e-32, V = 0.244, Var w = 0.094 naive: bias^2 = 0.0e+00, V = 0.378, Var w = 0.227 Both are unbiased, but the naive weight has about 2.4 times the variance of the AWRS weight as an estimate of $Z$, so it needs about 2.4 times as many samples for the same accuracy. ✓ https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test article.py L121 This turns a performance claim into a test. If a refactor leaves an estimator correct but noisier, an exact variance check catches it, with no line to fit through a noisy plot. The same goes for the SMC pipeline below: the tracer gives the exact variance of $\widehat{Z}$, so we can check whether a better proposal or more particles actually help. Which Trace Next? TraceSWOR makes each choice at random, in proportion to the unexplored mass below it. That isn't the only reasonable order. A deterministic alternative is greedy : always descend into the child with the most unexplored mass. It's a one-method subclass: class TraceGreedy TraceSWOR : """Always descend into the child with the most unexplored mass.""" def choose self, node : return int np.argmax node.child masses If we run every trace, the order doesn't matter: both give the same exact answer. It matters when we stop early. After exploring traces with total probability $c$, the unexplored traces have total probability $1 - c$, so if every weight is at most $B$, the answer can be off by at most $B 1 - c $. AWRS's weights are at most 1. ✓ https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test article.py L131 The bound holds for any order, and greedy tightens it faster because it heads for the heaviest unexplored part of the trie: program = lambda draw: awrs a p, lambda x: cond x , draw=draw def trace masses tracer : masses = while tracer.root.mass 0: with tracer: program tracer masses.append tracer.cur.mass return masses np.random.seed 0 orders = { 'sampling': trace masses TraceSWOR , 'greedy': trace masses TraceGreedy , } plot unexplored orders Sampling covers mass more slowly, but it offers something greedy doesn't. The traces it has seen are a sample without replacement, which supports an unbiased estimate of the part it hasn't seen Shi et al., 2020 https://arxiv.org/abs/2002.09067 . A deterministic order gives a guaranteed bound on the error; sampling gives an unbiased estimate. The replay above can run in either order. Why not go further, and always run the most probable unexplored trace next? Greedy heads for the heaviest subtree, not the heaviest single trace, so it isn't that. True best-first search needs to resume the program at an arbitrary unexplored branch, and the tracer can't: it remembers probabilities, not the program's state at each choice point. The workaround is to rerun the program from the start, forcing the choices along the path to that branch, and stop it there. That adds assumptions on top of the ones tracing already makes—rerunning must have no side effects, and the program must be safe to interrupt partway—and it costs a rerun for every node of the trie rather than one per trace. Probabilistic programming languages avoid the rerun by capturing a continuation at each choice point. WebPPL keeps them in a priority queue, which gives exactly this most-likely-first order Goodman & Stuhlmüller, 2014 http://dippl.org/chapters/03-enumeration.html . Without continuations, it's not something to bolt onto a tracer casually. Strategy 3: Testing the Whole Pipeline AWRS is not meant to be used on its own. In Lipkin et al., 2025 https://arxiv.org/abs/2504.05410 , it is the proposal step inside sequential Monte Carlo for generating text from a language model subject to a constraint. Each particle asks AWRS for a next token that satisfies the constraint, and multiplies its weight by the weight that AWRS returns. Program tracing is not limited to a single call of a sampling routine. It works on anything that makes discrete random choices through a hookable interface, so we can trace the whole pipeline at once—every AWRS call in every particle, plus the resampling steps. What Should This Code Do? Let $p$ be a language model over strings, and let $C$ be a constraint that can be checked one token at a time. The unnormalized target is $$ \gamma \mathbf{x} = p \mathbf{x} \cdot \mathbf{1} \mathbf{x} \in C $$ and its normalizing constant $Z = \sum {\mathbf{x}} \gamma \mathbf{x} $ is the probability that a string drawn from $p$ satisfies the constraint. SMC returns $M$ weighted particles $ \mathbf{x} m, w m $, and the specification is the proper weighting condition from before, lifted from tokens to strings: $$ \mathop{\mathbb{E}} \left \frac{1}{M} \sum {m=1}^{M} w m \cdot \mathbf{1} \mathbf{x} m = \mathbf{x}' \right = \gamma \mathbf{x}' $$ for every string $\mathbf{x}'$. Summing over $\mathbf{x}'$ gives the more familiar statement that $\frac{1}{M} \sum m w m$ is an unbiased estimate of $Z$. Checking every string is the stronger test: a bug that moves mass from one string to another without changing the total would slip past a check on $Z$ alone. Why should this hold? Proper weighting composes. At step $t$, AWRS returns a token and a weight that are properly weighted for $p \cdot \mid \mathbf{x} {\lt t} $ restricted to the allowed tokens. Multiplying these weights along a string gives a properly weighted sample for the product of the per-step targets, which is $\gamma$. Resampling preserves proper weighting as long as every resampled particle carries the average weight. Setting Up a Test Case The language model is a small, explicit distribution over strings. The constraint says that a b must be followed by an a or by the end of the string, so bb and bc are ruled out: lm = SimpleLM { : 0.1, 'a', : 0.1, 'a', 'a' : 0.1, 'a', 'b' : 0.1, 'b', : 0.1, 'b', 'a' : 0.15, 'b', 'b' : 0.25, 'b', 'c' : 0.1, } VOCAB = 'a', 'b', 'c', EOS def allowed prefix, token : """After a 'b', the next token must be 'a' or EOS.""" return not prefix -1: == 'b', and token not in 'a', EOS target lm = {s: prob is valid s for s, prob in lm.sequences.items } Z lm = sum target lm.values = 0.65, not 1 Utilities: SimpleLM , is valid EOS = '