{"slug": "ideas-for-testing-randomized-algorithms", "title": "Ideas for Testing Randomized Algorithms", "summary": "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.", "body_md": "Randomized algorithms are hard to test. Their outputs are, by design,\nunpredictable—and that unpredictability is a *feature*, not a bug. It is what\ngives them useful guarantees like unbiasedness and computational efficiency. But\nit also makes them slippery: how do you check that a stochastic program is\ncorrect when it gives different answers every time you run it?\n\nThis 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.\n\nThe 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.\n\n## A Tricky Little Sampling Algorithm\n\nSuppose we have a distribution $p$ over a finite set $\\mathcal{X} = \\{1,\n\\ldots, N\\}$ and a condition $C \\subseteq \\mathcal{X}$ (the set of \"good\"\nelements). We want to sample from the *conditional* distribution $p(\\cdot \\mid\nC)$. But rather than returning a plain sample, we want a *properly weighted*\nsample $(x, w)$—a sample-weight pair that can be used for unbiased estimation.\n\nThe 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:\n\nThe 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.\n\n### Four Implementations, at Least One Bug\n\nDid 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?\n\nEach variant takes a `draw` parameter that defaults to sampling from a\ncategorical distribution. This is a hook we'll exploit later: in production it\ncalls the RNG, but in testing we'll swap it for something that lets us\nenumerate every possible execution.\n\n## Variant A\n\n``` python\ndef awrs_a(p, cond, draw=sample_categorical):\n    p_cond = p.copy()\n    n, psi = 0, 0\n    xc0 = None\n    for i in range(2):\n        while True:\n            q = p_cond / p_cond.sum()\n            x = draw(q, context=('awrs', i, xc0))\n            if cond(x):\n                if i == 0:\n                    xc0 = x\n                break\n            else:\n                if i == 0:\n                    psi += p_cond[x]\n                p_cond[x] = 0\n            n += 1\n    return xc0, (1 - psi) / (n + 1)\n```\n\n## Variant B\n\n``` python\ndef awrs_b(p, cond, draw=sample_categorical):\n    p_cond = p.copy()\n    n, psi = 0, 0\n    xc0 = None\n    for i in range(2):\n        while True:\n            q = p_cond / p_cond.sum()\n            x = draw(q, context=('awrs', i, xc0))\n            if cond(x):\n                if i == 0:\n                    xc0 = x\n                break\n            else:\n                if i == 0:\n                    psi += p_cond[x]\n                p_cond[x] = 0\n            n += 1\n    return xc0, (1 + psi) / (n + 1)\n```\n\n## Variant C\n\n``` python\ndef awrs_c(p, cond, draw=sample_categorical):\n    p_cond = p.copy()\n    n, psi = 0, 0\n    xc0 = None\n    for i in range(2):\n        while True:\n            q = p_cond / p_cond.sum()\n            x = draw(q, context=('awrs', i, xc0))\n            if cond(x):\n                if i == 0:\n                    xc0 = x\n                break\n            else:\n                psi += p_cond[x]\n                p_cond[x] = 0\n            n += 1\n    return xc0, (1 - psi) / (n + 1)\n```\n\n## Variant D\n\n``` python\ndef awrs_d(p, cond, draw=sample_categorical):\n    p_cond = p.copy()\n    n, psi = 0, 0\n    xc0 = None\n    for i in range(2):\n        while True:\n            q = p_cond / p_cond.sum()\n            x = draw(q, context=('awrs', i, xc0))\n            if cond(x):\n                if i == 0:\n                    xc0 = x\n                break\n            else:\n                if i == 0:\n                    psi += q[x]\n                p_cond[x] = 0\n            n += 1\n    return xc0, (1 - psi) / (n + 1)\n```\n\nThe 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.\n\n## Hint\n\nCompare the `else` branch and the `return` statement across the four variants.\n\nOur job is to build testing strategies that catch the buggy ones—ideally without needing to know which line is wrong.\n\nBut before we can test *anything*, we need to answer the most important\nquestion in all of software testing:\n\n### What Should This Code Do?\n\nYou 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.\n\nFor AWRS, the specification is: it should return a sample-weight pair $(x, w)$\nsuch that $w \\cdot \\mathbf{1}[x = x']$ is an **unbiased estimator** of the\nunnormalized conditional distribution for every $x' \\in \\mathcal{X}$:\n\n$$ \\mathop{\\mathbb{E}}_{(x,w) \\sim \\text{AWRS}} \\left[ w \\cdot \\mathbf{1}[x = x'] \\right] = p(x') \\cdot \\mathbf{1}[x' \\in C] $$\n\nThe right-hand side is the unnormalized target—$p$ restricted to the condition $C$, with normalizing constant $Z = \\sum_{x \\in C} p(x)$.\n\nThat AWRS satisfies this is a theorem in\n[Lipkin et al., 2025](https://arxiv.org/abs/2504.05410), and we'll take it as\ngiven. Our question is whether the code does what the paper proves.\n\nNote 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.\n\n### Setting Up a Test Case\n\nNow 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:\n\n```\np = np.array([0.5, 0.3, 0.15, 0.04, 0.01])  # a distribution with a long tail\ncond = np.array([True, False, True, False, False])  # the \"good\" elements\nN = len(p)\n\nZ = p[cond].sum()\ntarget = p * cond  # unnormalized target (not a distribution!)\n```\n\n## Strategy 1: Monte Carlo Averaging\n\nThe obvious first testing strategy: run the code \"a lot of times\" and compute an average.\n\n```\nnp.random.seed(0)\nreps = 1_000\nfor name, awrs in variants.items():\n    have = np.zeros(N)\n    for r in range(1, 1 + reps):\n        (x, w) = awrs(p, lambda x: cond[x])\n        have[x] += w\n    err = abs(have / reps - target).sum()\n    print(f'Variant {name}: error = {err:.6f}')\nVariant A: error = 0.023367\nVariant B: error = 0.098657\nVariant C: error = 0.030520\nVariant D: error = 0.012054\n```\n\nIs this error small enough? How do you know? (Variant D has the smallest error of the four.)\n\nThis is the central problem with Monte Carlo testing: a single number tells you\nalmost nothing without a reference frame. We need to know what *should* happen\nas we increase the number of samples.\n\n### Convergence Rates\n\nA better approach is to watch the error decrease as a function of the number of\nsamples $M$. It should keep going down *forever*, modulo numerical precision.\nAnd crucially, we often know the *rate* at which it should decrease.\n\nThe general picture: suppose we have a sequence of estimators $\\widehat{\\mu}^{(M)}$ targeting some quantity $\\mu$. The mean squared error decomposes as\n\n$$ \\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 $$\n\nThe 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:\n\n- \n**I.I.D. unbiased** (e.g., our AWRS test). The bias is zero for all $M$.\n  The variance is $V/M$ where $V = \\text{tr}(\\text{Cov}(X))$—the sum of\n  the componentwise variances of a single sample. So the mean squared error is\n  exactly $V/M$.If you prefer root mean squared\n  error, the slope is $-1/2$ and the intercept is $\\frac{1}{2} \\log\n  V$.\n- \n**Correlated but unbiased** (e.g., sampling without replacement). The bias\n  is still zero. Samples drawn without replacement are negatively correlated,\n  which shrinks the variance: the rate is still $O(1/M)$, but the constant is\n  smaller.\n- \n**Biased but consistent** (e.g., self-normalized importance sampling, MCMC). The\n  bias is $O(1/M)$, so the squared bias is $O(1/M^2)$. The variance is\n  $O(1/M)$.\n  The MSE is still $O(1/M)$ overall, so on a log–log plot the slope looks\n  the same—the bias hides in the noise at large $M$.\n\nIn 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.\n\n## Deriving $\\,V\\,$ for vector-valued estimators\n\nOur 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)}]$:\n\n$$ \\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} $$\n\nwhere $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.\n\n```\nresults = {}\nfor name, awrs in variants.items():\n    np.random.seed(0)\n    have = np.zeros(N)\n    xs, ys = [], []\n    for r in range(1, 1 + 1_000):\n        (item, w) = awrs(p, lambda x: cond[x])\n        have[item] += w\n        if r >= 10:\n            xs.append(r)\n            ys.append(abs(have / r - target).sum())\n    slope, intercept = np.polyfit(np.log(xs), np.log(ys), deg=1)\n    results[name] = (xs, ys, slope, intercept)\n\nplot_single_runs(results)\n```\n\nA note on reading this plot: when we say \"log–log scale\" we mean the axes are\nlog-transformed but the tick labels show the original values (e.g., $10^2$,\n$10^3$). This is what a log axis does in a plotting library. This is different from\n*literally plotting* $\\log y$ vs $\\log x$, where the axes are linear but the\nvalues are logs. The line-fitting code does the latter (`np.polyfit` on\n$\\log x$, $\\log y$), but for visual inspection, log-scaled axes are usually\nclearer because you can read off the original quantities directly.\n\nFor a correct estimator, the fitted slope should be close to $-0.5$ (we're\nplotting error, not squared error). Variant B fails this test plainly: its\nerror stops going down after a few dozen samples and the slope is positive.\n**B is out.**\n\nThe 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.\n\n### What a Single Run Can and Can't Tell Us\n\nWatching a single long run tells us the *rate* of convergence, but it can't\ndistinguish the three cases above. An unbiased i.i.d. estimator, an estimator\nbuilt from samples drawn without replacement, and an MCMC chain can all produce\nvery similar log–log plots—the slopes are the same and the constants differ in ways\nthat are hard to eyeball.\n\nA pedagogical example: suppose we accidentally computed the average as\n$\\frac{1}{M+1} \\sum_{m=1}^{M} X^{(m)}$ instead of $\\frac{1}{M} \\sum_{m=1}^{M}\nX^{(m)}$. This estimator is *biased*—it systematically underestimates the\ntarget—but *consistent*, because the bias factor $M/(M+1) \\to 1$. The\nsingle-run convergence plot would look nearly identical.\n\nThe 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.\n\n### Averaging Over Multiple Runs\n\nTo separate bias from variance, we need **multiple independent runs**. The idea:\nrun the estimator many times at each sample size $M$, and compare the *average*\nbehavior to the theoretical prediction.\n\nRecall the decomposition from earlier:\n\n$$ \\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 $$\n\nFor the i.i.d. unbiased case, the variance term is $V/M$. So:\n\n- If the observed MSE matches $V/M$, the squared bias must be near zero—evidence of unbiasedness.\n- If the estimator has persistent bias, $\\| \\mathop{\\mathbb{E}}[\\widehat{\\mu}^{(M)}] -\n  \\mu \\|^2$ does not vanish with $M$. At large $M$, the variance $V/M$ shrinks\n  but the bias dominates, causing the observed curve to *flatten out* above the\n  theoretical line.\n- For biased-but-consistent estimators (like self-normalized IS), the squared\n  bias is $O(1/M^2)$—it vanishes\n  faster than the variance. This means the MSE curve will *eventually* match\n  $V/M$, but may sit slightly above it for small $M$.\n\n```\nresults = {}\nfor name, awrs in variants.items():\n    np.random.seed(0)\n    reps, M = 1000, 200\n    errors = np.zeros((reps, M))  # squared error of the running average, per run\n    samples = []                  # single-sample estimates, for V\n    for r in range(reps):\n        have = np.zeros(N)\n        for m in range(1, M + 1):\n            (item, w) = awrs(p, lambda x: cond[x])\n            have[item] += w\n            onehot = np.zeros(N)\n            onehot[item] = w\n            samples.append(onehot)\n            errors[r, m - 1] = np.linalg.norm(have / m - target)**2\n    V = np.trace(np.cov(np.array(samples).T))\n    results[name] = (errors, V)\n\nplot_mse(results)\n```\n\nA and D hug the theoretical $V/M$ line, which is strong evidence that they are\nproperly weighted—not just consistent, but unbiased, with the right\nvariance. B flattens out, as before. C is the interesting one: it tracks the\nline for small $M$ and then lifts off as the variance shrinks below its squared\nbias. Averaging over a thousand runs makes the gap far larger than the\nconfidence band. **C is out.**\n\nD, on the other hand, has now passed every test we've thrown at it.\n\nBut Monte Carlo testing is still a *blunt instrument*. Small biases hide\neasily in the noise, and you need many runs to have any statistical power. For\ncorrelated estimators, the effective variance may be hard to compute, making the theoretical line itself uncertain. In software—unlike\nthe real world—we can almost always do better.\n\n## Strategy 2: Program Tracing\n\nHere's the key insight: if the state space is small enough, we can *enumerate\nall possible executions* of the randomized program and verify the correctness\ncondition *exactly* (up to floating point).\n\nThe 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.\n\n### The Tracer\n\nThe tracing strategy comes from\n[Shi et al., 2020](https://arxiv.org/abs/2002.09067), who use it to sample\nwithout replacement from programs such as sequence models. The same idea\nunderlies exact inference by enumeration in probabilistic programming\nlanguages: WebPPL explores every execution of a program by capturing a\ncontinuation at each random choice\n([Goodman & Stuhlmüller, 2014](http://dippl.org/chapters/03-enumeration.html)).\nThe implementation uses a trie data structure that mirrors the branching structure of the\nprogram's random choices. Each node in the trie corresponds to a decision\npoint, and its children correspond to the possible outcomes.\n\nEach node stores the probability mass that hasn't been explored below it. A\nnode is created the first time the program reaches it, holding the probability\nof the path that leads there, so at the end of a trace `tracer.cur.mass` is\nthat trace's probability; this is how the checks below weight each output. When\na trace finishes, its mass is subtracted from every node on the way back up to\nthe root. Choosing children in proportion to their unexplored mass then never\nrevisits a finished trace, and the root's mass is the total probability of the\ntraces not yet run.\n\nHere is the tracer at work on variant A with our test case. Each step runs the\nprogram once more and adds one trace to the trie. Node colors mark which pass of\nAWRS made the choice, and a node fades as the probability mass below it is\nexplored. The order menu is explained [below](#which-trace-next).\n\nThe beautiful thing about `TraceSWOR` is that it works as a context manager:\neach `with tracer:` block explores one trace, and when it exits, that trace's\nmass is zeroed out so we never revisit it. We loop until the root's mass hits\nzero—meaning we've exhausted every possible execution.\n\nThe tracer is also available as a small, tested library,\n[`program_tracing`](https://github.com/timvieira/testing-randomized-algorithms),\nin this post's repository. It includes a `traces(program)` helper that runs this\nloop for you.\n\n## Tracer implementation\n\n``` python\nclass Node:\n    \"\"\"A node in the execution trie.\"\"\"\n    def __init__(self, idx, mass, parent, token=None):\n        self.idx = idx\n        self.mass = mass          # remaining probability mass\n        self.parent = parent\n        self.token = token        # label for visualization\n        self.child_masses = None  # set on first visit\n        self.keys = None          # outcomes of the choice made here\n        self.active_children = {}\n        self._mass = mass         # remember original mass\n        self.context = None\n\n    def sample(self):\n        \"\"\"Sample a child proportional to remaining mass.\"\"\"\n        p = self.child_masses / self.child_masses.sum()\n        return np.random.choice(len(p), p=p)\n\n    def update(self):\n        \"\"\"Propagate mass changes up to the root.\"\"\"\n        if self.parent is not None:\n            self.parent.child_masses[self.idx] = self.mass\n            self.parent.mass = np.sum(self.parent.child_masses)\n            self.parent.update()\n\nclass Tracer:\n    \"\"\"\n    Lazily materializes the probability tree of a generative process\n    by program tracing. Works with both array-valued and dict-valued\n    distributions.\n    \"\"\"\n    def __init__(self):\n        self.root = Node(idx=-1, mass=1.0, parent=None)\n        self.cur = None\n\n    def __call__(self, p, context=None):\n        \"\"\"Drop-in replacement for any sampling function.\"\"\"\n        if isinstance(p, dict):\n            keys = list(p.keys())\n            vals = np.array([p[k] for k in keys], dtype=float)\n        else:\n            p = np.asarray(p, dtype=float)\n            keys = list(range(len(p)))\n            vals = p\n        vals = vals / vals.sum()\n        cur = self.cur\n\n        if cur.child_masses is None:\n            cur.child_masses = cur.mass * vals\n            cur.keys = keys\n            cur.context = context\n        elif keys != cur.keys:\n            raise ValueError('the program made a different choice at the same point '\n                             'of two traces; it must be deterministic given its draws')\n\n        a = self.choose(cur)\n        if a not in cur.active_children:\n            cur.active_children[a] = Node(\n                idx=a, mass=cur.child_masses[a], parent=cur,\n                token=keys[a],\n            )\n        self.cur = cur.active_children[a]\n        return keys[a]\n\n    def choose(self, node):\n        \"\"\"Pick which child to explore next.\"\"\"\n        return node.sample()\n\nclass TraceSWOR(Tracer):\n    \"\"\"\n    Sampling without replacement meets program tracing.\n\n    Each execution trace is explored exactly once. After a trace\n    completes, its probability mass is zeroed out and propagated\n    up the trie, ensuring the next trace is fresh.\n    \"\"\"\n    def __enter__(self):\n        self.cur = self.root\n        return self\n\n    def __exit__(self, *args):\n        self.cur.mass = 0\n        self.cur.update()\n```\n\n### Exact Verification\n\nNow we can verify all four AWRS variants *exactly*:\n\n``` python\ndef exact_moments(awrs):\n    \"\"\"Exact first and second moments of w * onehot(x), over all traces.\"\"\"\n    tracer = TraceSWOR()\n    m1, m2 = np.zeros(N), np.zeros(N)\n    while tracer.root.mass > 0:\n        with tracer:\n            (item, w) = awrs(p, lambda x: cond[x], draw=tracer)\n            m1[item] += w * tracer.cur.mass\n            m2[item] += w**2 * tracer.cur.mass\n    return m1, m2\n\nfor name, awrs in variants.items():\n    prob, _ = exact_moments(awrs)\n    err = abs(target - prob).sum()\n    status = '\\u2713' if np.allclose(target, prob) else '\\u2717'\n    print(f'Variant {name}: {status}  (error = {err:.2e})')\nVariant A: ✓  (error = 3.61e-16)\nVariant B: ✗  (error = 9.08e-02)\nVariant C: ✗  (error = 3.15e-02)\nVariant D: ✗  (error = 1.61e-04)\n```\n\nThis is not a statistical test. It is an *exact* check (modulo floating point)\nthat 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\nsquinting at log–log plots. The tracer gives a definitive yes or no for each\nvariant. **D is out**, and A is the only survivor.\n\nD's bug is in the `else` branch: it accumulates the *normalized* probability\n`q[x]` instead of the remaining mass `p_cond[x]`. The two agree on the first\nrejection, because `p_cond` still sums to one. They differ only when the first\npass rejects two or more elements. With our long-tailed `p`, that happens in\nabout one run in twenty, and even then the weight is off by only a couple of\npercent[✓](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.\n\n### How Many Samples Would Monte Carlo Need?\n\nThe 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.\n\n```\nfor name, awrs in variants.items():\n    m1, m2 = exact_moments(awrs)\n    bias2 = np.sum((m1 - target)**2)\n    V = np.sum(m2 - m1**2)\n    if np.isclose(bias2, 0, atol=1e-20):\n        print(f'Variant {name}: unbiased')\n    else:\n        print(f'Variant {name}: bias^2 = {bias2:.1e}, V = {V:.2f}, '\n              f'M* = {V / bias2:,.0f}')\nVariant A: unbiased\nVariant B: bias^2 = 5.3e-03, V = 0.25, M* = 47\nVariant C: bias^2 = 6.4e-04, V = 0.24, M* = 382\nVariant D: bias^2 = 1.7e-08, V = 0.24, M* = 14,544,122\n```\n\nThis 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.\n\n### Comparing Variances\n\nThe 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:\n\n``` python\ndef naive(p, cond, draw=sample_categorical):\n    \"\"\"Draw once from p; weight 1 if the draw satisfies the condition, else 0.\"\"\"\n    x = draw(p)\n    return x, float(cond(x))\n\nfor name, sampler in [('AWRS', awrs_a), ('naive', naive)]:\n    m1, m2 = exact_moments(sampler)\n    bias2 = np.sum((m1 - target)**2)\n    V = np.sum(m2 - m1**2)\n    var_w = m2.sum() - m1.sum()**2   # variance of w as an estimate of Z\n    print(f'{name:>5}: bias^2 = {bias2:.1e}, V = {V:.3f}, Var(w) = {var_w:.3f}')\nAWRS: bias^2 = 6.5e-32, V = 0.244, Var(w) = 0.094\nnaive: bias^2 = 0.0e+00, V = 0.378, Var(w) = 0.227\n```\n\nBoth are unbiased, but the naive weight has about 2.4 times the variance of the\nAWRS weight as an estimate of $Z$, so it needs about 2.4 times as many samples\nfor 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\nrefactor leaves an estimator correct but noisier, an exact variance check\ncatches it, with no line to fit through a noisy plot. The same goes for the\nSMC pipeline below: the tracer gives the exact variance of $\\widehat{Z}$, so we\ncan check whether a better proposal or more particles actually help.\n\n### Which Trace Next?\n\n`TraceSWOR` makes each choice at random, in proportion to the unexplored mass\nbelow it. That isn't the only reasonable order. A deterministic alternative is\n**greedy**: always descend into the child with the most unexplored mass. It's a\none-method subclass:\n\n```\nclass TraceGreedy(TraceSWOR):\n    \"\"\"Always descend into the child with the most unexplored mass.\"\"\"\n    def choose(self, node):\n        return int(np.argmax(node.child_masses))\n```\n\nIf we run every trace, the order doesn't matter: both give the same exact\nanswer. It matters when we stop early. After exploring traces with total\nprobability $c$, the unexplored traces have total probability $1 - c$, so if\nevery weight is at most $B$, the answer can be off by at most $B(1 - c)$.\n(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\ntightens it faster because it heads for the heaviest unexplored part of the\ntrie:\n\n```\nprogram = lambda draw: awrs_a(p, lambda x: cond[x], draw=draw)\n\ndef trace_masses(tracer):\n    masses = []\n    while tracer.root.mass > 0:\n        with tracer:\n            program(tracer)\n            masses.append(tracer.cur.mass)\n    return masses\n\nnp.random.seed(0)\norders = {\n    'sampling': trace_masses(TraceSWOR()),\n    'greedy': trace_masses(TraceGreedy()),\n}\n\nplot_unexplored(orders)\n```\n\nSampling covers mass more slowly, but it offers something greedy doesn't. The\ntraces it has seen are a sample without replacement, which supports an unbiased\nestimate of the part it hasn't seen\n([Shi et al., 2020](https://arxiv.org/abs/2002.09067)). A deterministic order\ngives a guaranteed bound on the error; sampling gives an unbiased estimate. The\nreplay above can run in either order.\n\nWhy not go further, and always run the most probable unexplored trace next?\nGreedy heads for the heaviest subtree, not the heaviest single trace, so it\nisn't that. True best-first search\nneeds to resume the program at an arbitrary unexplored branch, and the tracer\ncan't: it remembers probabilities, not the program's state at each choice\npoint. The workaround is to rerun the program from the start, forcing the\nchoices along the path to that branch, and stop it there. That adds\nassumptions on top of the ones tracing already makes—rerunning must have\nno side effects, and the program must be safe to interrupt partway—and it\ncosts a rerun for every node of the trie rather than one per trace.\nProbabilistic programming languages avoid the rerun by capturing a\ncontinuation at each choice point. WebPPL keeps them in a priority queue, which\ngives exactly this most-likely-first order\n([Goodman & Stuhlmüller, 2014](http://dippl.org/chapters/03-enumeration.html)).\nWithout continuations, it's not something to bolt onto a tracer casually.\n\n## Strategy 3: Testing the Whole Pipeline\n\nAWRS is not meant to be used on its own. In [Lipkin et al., 2025](https://arxiv.org/abs/2504.05410),\nit is the proposal step inside sequential Monte Carlo for generating text from\na language model subject to a constraint. Each particle asks AWRS for a next\ntoken that satisfies the constraint, and multiplies its weight by the weight\nthat AWRS returns.\n\nProgram 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.\n\n### What Should This Code Do?\n\nLet $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\n\n$$ \\gamma(\\mathbf{x}) = p(\\mathbf{x}) \\cdot \\mathbf{1}[\\mathbf{x} \\in C] $$\n\nand 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:\n\n$$ \\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}') $$\n\nfor 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.\n\nWhy 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.\n\n### Setting Up a Test Case\n\nThe language model is a small, explicit distribution over strings. The\nconstraint says that a `b` must be followed by an `a` or by the end of the\nstring, so `bb` and `bc` are ruled out:\n\n```\nlm = SimpleLM({\n    (): 0.1, ('a',): 0.1, ('a', 'a'): 0.1, ('a', 'b'): 0.1,\n    ('b',): 0.1, ('b', 'a'): 0.15, ('b', 'b'): 0.25, ('b', 'c'): 0.1,\n})\n\nVOCAB = ['a', 'b', 'c', EOS]\n\ndef allowed(prefix, token):\n    \"\"\"After a 'b', the next token must be 'a' or EOS.\"\"\"\n    return not (prefix[-1:] == ('b',) and token not in ('a', EOS))\n\ntarget_lm = {s: prob * is_valid(s) for s, prob in lm.sequences.items()}\nZ_lm = sum(target_lm.values())  # = 0.65, not 1\n```\n\n## Utilities: `SimpleLM`, `is_valid`\n\n```\nEOS = '<EOS>'  # end-of-sequence marker\n\nclass SimpleLM:\n    \"\"\"\n    A sequence model defined by an explicit, finite distribution over strings.\n    Maps prefixes to next-token distributions (as dicts).\n    \"\"\"\n    def __init__(self, sequences):\n        \"\"\"\n        sequences: dict mapping tuple-of-tokens to probability.\n        E.g., {('a','a'): 0.3, ('b','b','b'): 0.6, (): 0.1}\n        \"\"\"\n        self.sequences = sequences\n\n    def p_next(self, prefix):\n        \"\"\"Return conditional distribution over next token given prefix.\"\"\"\n        dist = {}\n        for seq, prob in self.sequences.items():\n            if len(seq) < len(prefix):\n                continue\n            if seq[:len(prefix)] != prefix:\n                continue\n            if len(seq) == len(prefix):\n                dist[EOS] = dist.get(EOS, 0) + prob\n            else:\n                token = seq[len(prefix)]\n                dist[token] = dist.get(token, 0) + prob\n        total = sum(dist.values())\n        return {k: v / total for k, v in dist.items()}\n\ndef is_valid(string):\n    return all(allowed(string[:t], string[t]) for t in range(len(string)))\n```\n\nHere is the sampler. At each step, every unfinished particle extends itself\nwith AWRS and multiplies in the AWRS weight. When the effective sample size\ndrops below `tau` times the number of live particles, it resamples. With\n`tau=0` it never resamples, and the algorithm reduces to sequential importance\nsampling (SIS); with `tau>0` it is sequential importance resampling (SIR).\n\n``` python\nfrom collections import namedtuple\nParticle = namedtuple('Particle', 'weight, string')\n\ndef smc(lm, M, tau, awrs, draw=sample_categorical):\n    unfinished = [Particle(1, ()) for _ in range(M)]\n    finished = []\n    while unfinished:\n        next_unfinished = []\n        for particle in unfinished:\n            P = lm.p_next(particle.string)\n            p = np.array([P.get(v, 0) for v in VOCAB])\n            i, w = awrs(p, lambda i: allowed(particle.string, VOCAB[i]), draw=draw)\n            x = VOCAB[i]\n            if x == EOS:\n                finished.append(Particle(particle.weight * w, particle.string))\n            else:\n                next_unfinished.append(Particle(particle.weight * w,\n                                                particle.string + (x,)))\n        unfinished = next_unfinished\n        if not unfinished:\n            break\n\n        W = np.array([particle.weight for particle in unfinished])\n        if W.sum()**2 / (W** 2).sum() < tau * len(unfinished):\n            indices = [draw(W / W.sum(), context='resample') for _ in unfinished]\n            unfinished = [Particle(W.sum() / len(unfinished), unfinished[k].string)\n                          for k in indices]\n\n    return finished\n```\n\nThe same `draw` hook is passed down to every AWRS call and to the resampling\nstep, so a single tracer sees every random choice the pipeline makes.\n\n### Sequential Importance Sampling\n\nOne particle is enough to check the proper weighting condition for SIS. We run\nit with each AWRS variant, plus one more contender: *locally constrained\ndecoding*, which samples each token from the language model restricted to the\nallowed tokens and ignores the weight. It is the obvious way to do constrained\ngeneration.\n\n``` python\ndef locally_constrained(p, cond, draw=sample_categorical):\n    \"\"\"AWRS's sample, but with its weight thrown away.\"\"\"\n    x, _ = awrs_a(p, cond, draw)\n    return x, 1\n\ndef check_smc(awrs, M, tau):\n    tracer = TraceSWOR()\n    have = defaultdict(float)\n    n_traces = 0\n    while tracer.root.mass > 0:\n        with tracer:\n            for particle in smc(lm, M, tau, awrs, draw=tracer):\n                have[particle.string] += particle.weight / M * tracer.cur.mass\n        n_traces += 1\n    err = sum(abs(have[s] - target_lm[s]) for s in target_lm)\n    status = '\\u2713' if err < 1e-10 else '\\u2717'\n    return status, err, n_traces\n\nprint(f'Z = {Z_lm:.2f}')\ncontenders = {**variants, 'locally constrained': locally_constrained}\nfor name, awrs in contenders.items():\n    status, err, n_traces = check_smc(awrs, M=1, tau=0)\n    print(f'SIS with {name}: {status}  '\n          f'(error = {err:.2e}, traces = {n_traces})')\nZ = 0.65\nSIS with A: ✓  (error = 2.78e-17, traces = 162)\nSIS with B: ✗  (error = 1.21e-01, traces = 162)\nSIS with C: ✗  (error = 3.34e-02, traces = 162)\nSIS with D: ✗  (error = 4.22e-03, traces = 162)\nSIS with locally constrained: ✗  (error = 3.50e-01, traces = 162)\n```\n\nEvery bug is caught through the pipeline, including D. After a `b`, two of\nthe four tokens are forbidden, so AWRS's first pass can reject twice in one\nstep—exactly the situation that triggers D's bug.\n\nThe last line is the most instructive one. Locally constrained decoding never\nproduces an invalid string, but it gets the probabilities wrong: it\nover-represents strings whose prefixes leave little room under the constraint.\nHere, it puts too much mass on `b` and `ba`[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L163), because it never pays for the\nprobability that `b` would have been followed by a forbidden token. This bias\nis the reason AWRS returns a weight in the first place.\n\n### Sequential Importance Resampling\n\nResampling is just one more random choice, made through the same `draw` hook,\nso the tracer needs no changes. Two particles are enough to exercise it:\n\n```\nfor name in ['A', 'D']:\n    status, err, n_traces = check_smc(variants[name], M=2, tau=0.75)\n    print(f'SIR with {name}: {status}  '\n          f'(error = {err:.2e}, traces = {n_traces:,})')\nSIR with A: ✓  (error = 1.10e-14, traces = 28,836)\nSIR with D: ✗  (error = 4.22e-03, traces = 30,132)\n```\n\nThe correct implementation passes and D is still caught. But notice the\nnumber of traces: 162 with one particle, and nearly 30,000 with two.[✓](https://github.com/timvieira/testing-randomized-algorithms/blob/master/tests/test_article.py#L175)\n\n**Pro tip:** Make sure the total mass decreases quickly, as the loop will halt\nonly when the total mass is zero. This requires the support to be finite. For\nproblems where the support is large but finite, it may be a reasonable heuristic\nto stop early if the remaining mass is small, trading exactness for a bound or\nan unbiased estimate (see [Which Trace Next?](#which-trace-next)).Traces are a combinatorial space, so this strategy is only practical for smaller problems. But \"smaller problems\" is exactly the regime where tests should live—small enough to verify exactly, but complex enough to exercise the algorithm's logic.\n\n## Limitations\n\nProgram tracing is not a universal tool. It asks a few things of the program under test:\n\n- \n**Discrete, finite randomness.** Every random choice must have finitely many\n  outcomes, and every execution must terminate. A Gaussian draw, or a loop that\n  can run forever (rejection sampling*with* replacement, say), has infinitely\n  many traces, and the tracer never finishes. Sometimes you can test a discrete\n  skeleton of the algorithm, or discretize the continuous part; otherwise you\n  are back to Monte Carlo.\n- \n**All randomness goes through the hook.** The program must be a\n  deterministic function of the values`draw` returns. A stray call to the\n  global RNG, iteration over an unordered set, or a race between threads breaks\n  that assumption, and the results are meaningless.\n- \n**Small instances only.** The number of traces grows combinatorially: one\n  particle gave 162 traces, and two gave nearly 30,000. Tests have to live on\n  toy instances, so bugs that only appear at scale—underflow in long\n  sequences, a code path that only large vocabularies reach—need other\n  tests.\n- \n**Only as good as the specification.** The tracer checks the property you\n  wrote down. If the target is wrong, or the property is too weak (checking\n  only $Z$ when $Z = 1$, for instance), it will confirm the wrong thing with\n  perfect confidence.\n- \n**Exact up to floating point.** Comparisons still need a tolerance, so a bug\n  confined to traces whose total probability is below that tolerance goes\n  unnoticed.\n\n## Takeaways\n\nWe climbed a ladder of tests, and each rung caught a bug that the one below it missed:\n\n| Test | What it caught | Cost | \n|---|---|---|\n| One run, watching the error fall | B: large bias, visible after ~50 samples | one run | \n| Many runs, compared to $V/M$ | C: small bias, visible after ~400 samples | a thousand runs | \n| Program tracing | D: tiny bias, visible after ~15 million samples | one run per trace (196) | \n| Tracing the whole SMC pipeline | bugs in how the pieces fit, like dropping the AWRS weight | one run per trace (162) | \n\nTesting is fundamentally about **control**. Some abstractions give us more\ncontrol over a system's behavior than others, and good tests depend on choosing\nabstractions that make critical behaviors observable, reproducible, and\nmanipulable.\n\nThe `draw` parameter in all of the algorithms above is a good example. It's a\ntiny abstraction—just a function that returns a random choice—but it gives us\nenormous leverage. In production, it calls the RNG. In testing, it becomes a\ntracer that enumerates every possible execution. Same algorithm, radically\ndifferent observability.\n\nA few principles I've found useful:\n\n- \n**Theory gives you test cases.** The proper weighting definition, the\n  unbiasedness of importance sampling, the convergence rate of Monte\n  Carlo—these are all*specifications* that can be mechanically checked. If you\n  don't understand the theory behind your algorithm, you don't know what to test.\n- \n**Small exact tests beat large approximate tests.** A tracer that exhaustively\n  checks a 5-element distribution is worth more than a Monte Carlo test with\n  10,000 samples. The tracer gives you a definitive yes/no; the Monte Carlo test\n  gives you a squiggly line that you have to interpret.\n- \n**Inject errors to calibrate your tests.** A test that passes on correct code\n  but also passes on buggy code is useless. Deliberately break the algorithm and\n  verify that your test catches it.\n- \n**Monte Carlo is the blunt instrument of last resort.** It's indispensable\n  when the state space is too large to enumerate, but in software—unlike the\n  real world—we can almost always find ways to do better. Freezing the RNG\n  makes your code deterministic, but it is at best a regression test and\n  extremely brittle.\n\nOf course, some behaviors are only visible from \"1,000 feet\"—via integration tests, stress tests, or long-horizon statistical checks. Not everything can be checked with a microscope. But the more control you have over the randomness in your program, the more you can check exactly rather than approximately.\n\n**How this post was made.** This post was written with extensive help from\n[Claude Code](https://claude.ai/claude-code). The ideas, the framing, and the\nfinal editorial decisions are mine.\n\nEvery numerical claim in the post (which variants pass, the sample sizes, the variance ratio, the trace counts) is checked by a test in the repository, and the code shown here is the code those tests run. Mistakes are mine.\n\nFound an error or have a suggestion? Please\n[open an issue](https://github.com/timvieira/testing-randomized-algorithms/issues)\non GitHub.", "url": "https://wpnews.pro/news/ideas-for-testing-randomized-algorithms", "canonical_source": "https://timvieira.github.io/testing-randomized-algorithms/", "published_at": "2026-11-04 08:57:20+00:00", "updated_at": "2026-11-04 09:25:32.190489+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-research"], "entities": ["AWRS", "Lipkin et al.", "sequential Monte Carlo (SMC)", "arXiv:2504.05410"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/ideas-for-testing-randomized-algorithms", "markdown": "https://wpnews.pro/news/ideas-for-testing-randomized-algorithms.md", "text": "https://wpnews.pro/news/ideas-for-testing-randomized-algorithms.txt", "jsonld": "https://wpnews.pro/news/ideas-for-testing-randomized-algorithms.jsonld"}}