mixle.enumeration.envelope module

Approximate deep enumeration for autoregressive models (LLMs): the per-depth envelope index.

The exact autoregressive count index (autoregressive) is a tree recursion: each prefix has its own next-token distribution, so counting to a bit budget B must expand every live prefix – Theta(count / V) work. That is the right tool for the head (ranks up to ~1e6 on a real LM), and provably the end of the road for deep ranks: reaching rank 1e15 exactly would visit ~1e13 prefixes.

This module trades that wall for a mean-field approximation with an explicit contract:

• Precompute (the envelope): ancestral-sample n_paths contexts per depth and average their next-token fine-bucket histograms into one envelope E_d per depth (depth 0 is the real root context, so it is exact). Suffix-convolve them once – S_d = E_d (*) S_{d+1} – with float64 counts at C speed. S_0 approximates the full count histogram the way a SequenceDistribution computes its own exactly: the approximation is precisely “treat the steps as independent draws from the per-depth aggregate”.

• Query: count / threshold / mass_above read S_0. unrank(i) descends the real model – one forward per step, V real step buckets at each depth – apportioning the target offset among tokens by their envelope-estimated subtree counts. O(L) forwards per query, never Theta(count).

Contract (honest): the returned sequences are real model outputs and every reported log_prob is the exact model log-probability; only the rank coordinate is approximate. For a prefix-independent (iid-step) model the envelope equals the true per-step histogram and everything here is exact (tested). For a context-dependent model the envelope is an estimate averaged over typical (ancestrally sampled) contexts – the same high-probability contexts that dominate the enumeration head; rank_bracket returns the induced bucket bracket so downstream users can carry the uncertainty. Fixed-length models only (a terminating/eos model needs an absorbing channel – raise, do not guess).

class AREnvelopeIndex(model, *, n_paths=64, seed=0, budget_bits=64.0, calibration_sequences=None)

Bases: object

Envelope (mean-field) seek index over a fixed-length AutoregressiveEnumerable.

Parameters:

• model (Any) – the autoregressive adapter (max_len set; terminating/eos models are rejected).

• n_paths (int) – contexts sampled per depth for the envelope calibration (more = sharper envelope; the root depth is always exact). The calibration forwards are memoized on the model, so they are shared with any exact index built later.

• seed (int) – calibration sampling seed (the index is deterministic given it).

• budget_bits (float) – initial depth of the suffix tables; queries deepen geometrically as needed.

• calibration_sequences (list[tuple] | None) – optional typical sequences (a corpus, a provider’s fast generations) to calibrate the envelope from INSTEAD of ancestral sampling. Each is harvested through model.harvest – with the all_position_logprobs contract that is ONE forward per sequence for all its per-depth contexts, ~L-times cheaper than sampling token by token.

ensure_bits(depth_bits)

Deepen the suffix tables to cover depth_bits (geometric, cheap: L capped convolutions).

Parameters:

depth_bits (float)

Return type:

AREnvelopeIndex

total()

Estimated number of sequences within the built depth (exact for an iid-step model).

Return type:

float

count(min_log_prob)

Estimated number of sequences with log_density >= min_log_prob (mean-field; iid-exact).

Parameters:

min_log_prob (float)

Return type:

float

mass_above(min_log_prob)

A mean-field (lower, upper) estimate of the head mass above min_log_prob.

Same bucket arithmetic as the exact index’s mass_above (each bucket of c sequences holds between c * 2**-hi_bits and c * 2**-lo_bits of mass), applied to the envelope histogram – so the bracket carries the envelope’s estimation error on top of the quantization smear.

Parameters:

min_log_prob (float)

Return type:

tuple[float, float]

unrank(i)

The approximately-i-th most probable sequence and its exact log-probability.

Costs one model forward per step (L total, memoized) – never a tree expansion. The rank coordinate inherits the envelope approximation; the returned sequence and its log-probability are exact model quantities. Raises IndexError past the estimated support size.

Parameters:

i (int)

Return type:

tuple[tuple, float]

threshold(rank)

Exact log-probability of the sequence the envelope places at rank (approximate boundary).

Parameters:

rank (int)

Return type:

float

rank_bracket(sequence)

Estimated [lo, hi] rank bracket of sequence – its envelope bucket’s rank span.

lo counts the estimated sequences in strictly shallower buckets; hi adds the sequence’s own bucket. Exact for iid-step models; otherwise a mean-field estimate (floats, not certificates).

Parameters:

sequence (Any)

Return type:

tuple[float, float]

class LatticeEnvelopeIndex(model, *, n_clusters=None, cluster_fn=None, n_paths=128, seed=0, budget_bits=64.0)

Bases: object

Cluster-conditioned envelope index: the Markov refinement between mean-field and exact.

AREnvelopeIndex averages one envelope per depth – exact only for iid steps. This index conditions each depth’s envelope on the cluster of the last token (cluster_fn): the calibration histograms are kept per (depth, cluster) and split by the next token’s cluster, and the suffix tables become the lattice DP S_d[c] = sum_c' H_d[c][c'] (*) S_{d+1}[c'] – the same shape as HMMPathIndex with clusters as states. The result is exact for any order-1 Markov model whose next-token distribution depends only on (depth, cluster_fn(last token)) – with cluster_fn the identity, that is every last-token Markov model, where the mean-field envelope is provably lossy. Coarser clusterings interpolate: m = 1 recovers mean-field, m = V conditions on the full last token.

Queries mirror AREnvelopeIndex: counts and thresholds off the root table, unrank descends the REAL model apportioning by cluster-conditioned subtree estimates (O(L) forwards), every returned log-probability exact. A (depth, cluster) never visited in calibration borrows the depth’s aggregate envelope (mean-field fallback for that pocket) – raise n_paths to shrink the pockets. Fixed-length models only.

Parameters:

• model (Any)

• n_clusters (int | None)

• cluster_fn (Any)

• n_paths (int)

• seed (int)

• budget_bits (float)

ensure_bits(depth_bits)

Parameters:

depth_bits (float)

Return type:

LatticeEnvelopeIndex

total()

Estimated sequences within the built depth (exact for depth+last-cluster Markov models).

Return type:

float

count(min_log_prob)

Estimated number of sequences with log_density >= min_log_prob (Markov-refined estimate).

Parameters:

min_log_prob (float)

Return type:

float

unrank(i)

The approximately-i-th most probable sequence with its exact log-probability.

Parameters:

i (int)

Return type:

tuple[tuple, float]

threshold(rank)

Parameters:

rank (int)

Return type:

float

rank_bracket(sequence)

Estimated [lo, hi] rank bracket – exact for depth+last-cluster Markov models.

Parameters:

sequence (Any)

Return type:

tuple[float, float]