Modular Adaptive Rejection Sampling

Autoregressive language models are increasingly used in combination: one wants to intersect two models so that outputs are probable under both, take their union, sharpen a model, or contrast it against a weaker model. Existing approaches for sampling from such combinations are either efficient but not faithful to the distribution induced by the model combination, or faithful but not efficient, paying for every disagreement between the models at every sample. We present Modular Adaptive Rejection Sampling (\mars), an exact rejection sampling algorithm that supports complex combinations of distributions and that learns from its rejections. \mars generates a sample token by token, keeping for every prefix an upper bound on the total target weight of the sequences that start with that prefix and choosing each next token in proportion to these bounds. When \mars extends a prefix, the bounds of the prefix's possible extensions reveal how much the prefix's own bound overpromised; \mars rejects with exactly that probability and remembers the tighter bound, so the same rejection never happens again. Such bounds, which we call envelopes, are closed under intersection, union, sharpening, reweighting, and constraining, which yields an algebra of language models: any globally-normalized distribution built from these operations can be sampled exactly by the same algorithm. On three benchmarks, \mars uses up to $8\times$ fewer forward passes than sequential MC at the particle count prior work reports, while sampling exactly; against rejection sampling, \mars needs up to $4.3\times$ fewer descents per sample, and sometimes samples where rejection sampling times out at 400 attempts.