mixle.enumeration.rescore module

Speculative enumeration: build the index with a low-cost DRAFT model, score results with the TARGET.

Speculative decoding’s economics applied to enumeration. Building any autoregressive index costs one forward per live prefix – prohibitive when the model is a large transformer. But the ordering work (which sequences are near a rank/threshold) tolerates approximation, while the scores must be the real model’s. So: let a low-cost draft (an n-gram, a distilled student, a quantized twin) pay for the tree or envelope build, and touch the target only for the sequences a query actually returns – one batched teacher-forcing forward for all of them (AutoregressiveEnumerable.score_sequences()).

Contract: every returned log_prob is the target’s exact score. The order is draft-approximate, repaired locally by window reranking: top_k(k) / slice pull k + rerank_window draft-ordered candidates, rescore them all with the target in one batch, and sort by target score. That is exact whenever no unpulled sequence out-scores the returned ones.

Soundness (the certified verdict on RescoreResult): the draft index’s rank order is only guaranteed up to its own quantized fine-bucket width (SeekIndex and AREnvelopeIndex are both built over a Quantizer) – so the last PULLED candidate’s own exact draft score does NOT bound every unpulled one: an unpulled candidate can share the edge item’s quantized bucket, or even score higher, while still being ranked later (MXR-080-0233). The sound bound instead comes from the draft index’s OWN quantizer: every candidate ranked at or after the pulled window has a fine bucket at least as deep as the edge item’s, so that bucket’s own best-case (least negative) log-prob edge – strictly looser than the edge item’s exact score, never tighter – is a certified upper bound on any unpulled candidate’s draft score. certified is True only when the worst target score in the returned set provably clears bound + assumed_gap; False when certification was requested but could not be established (no sound bound available – e.g. the draft index exposes no quantizer – or the bound did not clear the gap); None when no assumed_gap was supplied at all, so certification was never attempted. The observed gap is always reported as a diagnostic, certified or not.

class RescoredIndex(draft_index, target, *, rerank_window=64, assumed_gap=None)[source]

Bases: object

Draft-ordered, target-scored enumeration with window reranking.

Parameters:

  • draft_index (Any) – any index with unrank(i) -> (sequence, draft_log_prob) – a SeekIndex over a low-cost AutoregressiveEnumerable, an AREnvelopeIndex, or anything equivalent. Sound certification additionally requires a .quantizer attribute exposing the Quantizer the index was built with – both accepted index types have one; without it, queries still work but never certify.

  • target (Any) – the expensive model – an AutoregressiveEnumerable (its score_sequences() batch scorer is used) or a bare callable [seqs] -> log_probs.

  • rerank_window (int) – extra draft candidates pulled around a query and reranked by target score. Larger = more robust to draft/target disagreement, one batched forward either way. Must be a non-negative integer – a negative window used to silently shrink or empty the pulled set while still reporting a falsely certified result (MXR-080-0233).

  • assumed_gap (float | None) – optional global bound on |target_lp - draft_lp| (nats). Must be finite and non-negative. When supplied, results carry a genuine certified verdict (see RescoreResult); otherwise certified is None and the observed gap is reported as a diagnostic only.

top_k(k)[source]

The k best sequences by TARGET score among the k + rerank_window draft head.

Returns a RescoreResult – target-exact scores, draft+window-approximate completeness (see the class/module docstrings for exactly what certified proves).

Parameters:

k (int)

Return type:

RescoreResult

slice(start, k)[source]

Target-reranked [start, start + k) slice of the pulled start + k + rerank_window head.

Same semantics as top_k(): order within the pulled set is target-exact; the certificate covers whether an unpulled sequence could belong in (or before) the slice.

Parameters:
Return type:

RescoreResult

unrank(i)[source]

The draft’s rank-i sequence with the TARGET’s exact log-probability.

The rank coordinate is the draft’s (no reranking): the low-cost random-access primitive. Use top_k() / slice() when local target-order matters.

Parameters:

i (int)

Return type:

tuple[tuple, float]

class RescoreResult(items, certified, bound, gap)[source]

Bases: object

Outcome of a RescoredIndex query: target-exact items plus an explicit certificate.

top_k/slice are speculative: the returned items are always target-exact and target-ordered, but completeness (whether some unpulled candidate could outscore them) is only provable under an assumed_gap AND a draft index that can supply a sound bound – a caller that needs the proof, not just a plausible answer, must check certified is True rather than trusting items alone (MXR-080-0233: a bare boolean bolted onto an unqualified value cannot make that distinction).

Parameters:
items

(sequence, target_log_prob) pairs, target-exact and sorted by target score.

Type:

list[tuple[tuple, float]]

certified

None when no assumed_gap was supplied to the index – certification was never attempted and items is a plausible, uncertified answer. True only when bound was established and the returned set’s worst target score provably clears every unpulled candidate (bound + assumed_gap). False when certification was requested but could not be proven – either the draft index offered no sound bound, or the bound did not clear the gap; items is then still the best-effort draft/window answer, never a proof.

Type:

bool | None

bound

the certified upper bound on any unpulled candidate’s draft log-probability that the certificate was checked against, or None when no bound applies – nothing was left unpulled (the draft support was exhausted), no assumed_gap was supplied, or the draft index could not supply a sound bound.

Type:

float | None

gap

running max observed |target_lp - draft_lp| over everything rescored on this index so far – a diagnostic of what the gap has actually been, never a certificate.

Type:

float