mixle.enumeration.model_enumeration module

Enumerate the outputs of arbitrary scoring models (neural nets, transformers, …) by score.

The mixle distribution enumerators (dist.enumerator()) need a mixle distribution. These utilities instead work with any model that can score candidates, supplied as a plain Python callable – so a transformer, RNN, n-gram language model, or any classifier can be enumerated in descending log-probability order without being a mixle distribution. Nothing here imports a deep-learning framework; the caller’s callable bridges to their model.

Four entry points, from most general to most specific:

  • best_first(...) – the generic engine: best-first / A* search that yields goal states in descending score, given successors, an is_goal test, a score, and an optional admissible heuristic.

  • best_first_decode(next_logprobs, ...) – EXACT descending-probability enumeration of sequences from an autoregressive model. next_logprobs(prefix) returns (token, log_prob) continuations (e.g. a transformer’s log_softmax of the next-token logits). Yields (sequence, total_log_prob) lazily, best first. Exact because each step’s log-prob is <= 0, so a prefix’s score upper-bounds every completion.

  • beam_search(next_logprobs, beam_width, ...) – the classic approximate top-k decode (fixed beam), for when exact best-first explores too much.

  • top_k_scored(candidates, score, k) – top-k over a finite candidate set scored by a callable (e.g. a classifier’s class log-probabilities).

Example (transformer-style next-token decoding):

import numpy as np
def next_logprobs(prefix):
    logits = my_transformer(prefix)            # (vocab,) numpy/torch -> numpy
    lp = logits - logsumexp(logits)            # log_softmax
    return list(enumerate(lp))                 # [(token_id, log_prob), ...]
for seq, total_lp in best_first_decode(next_logprobs, eos=EOS, max_len=20, max_results=5):
    ...                                        # the 5 highest-probability sequences, best first
best_first(start, successors, is_goal, score, heuristic=None, max_results=None)[source]

Best-first / A* search yielding goal states in descending score (log-probability) order.

Yields (goal_state, score(goal_state)) lazily, highest first. The ordering is exact when f(state) = score(state) + heuristic(state) is an admissible upper bound on the score of every goal reachable from state – in particular when heuristic is omitted (treated as 0) and score never increases along a path (the usual case for cumulative log-probabilities, which add terms <= 0).

Parameters:

  • start (Any) – the initial (typically partial) state.

  • successors (Callable[[Any], Iterable[Any]]) – expand a state into its child states.

  • is_goal (Callable[[Any], bool]) – True when a state is a complete output to yield (and not expanded further).

  • score (Callable[[Any], float]) – the (partial) log-score of a state.

  • heuristic (Callable[[Any], float] | None) – optional admissible upper bound on the best completion score reachable from a state.

  • max_results (int | None) – stop after yielding this many goals; None enumerates until the frontier is empty.

Yields:

(goal_state, score) in nonincreasing score.

Return type:

Iterator[tuple[Any, float]]

best_first_decode(next_logprobs, eos=None, max_len=None, start=(), heuristic=None, max_results=None)[source]

Exactly enumerate an autoregressive model’s sequences in descending total log-probability.

Parameters:

  • next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]]) – next_logprobs(prefix) returns an iterable of (token, log_prob) continuations of prefix (e.g. the log-softmax of a transformer’s next-token logits). Log-probs must be <= 0.

  • eos (Any) – end-of-sequence token; a prefix whose last token is eos is complete (and not extended).

  • max_len (int | None) – maximum sequence length; a prefix of this length is complete. At least one of eos / max_len should be given or enumeration may not terminate.

  • start (tuple) – the initial prefix (default empty).

  • heuristic (Callable[[tuple], float] | None) – optional admissible upper bound on the remaining log-probability from a prefix (e.g. remaining_steps * max_step_logprob); tightens the search. Omit for the exact h=0 search.

  • max_results (int | None) – stop after this many complete sequences.

Yields:

(sequence_tuple, total_log_prob) in nonincreasing total log-probability.

Return type:

Iterator[tuple[tuple, float]]

beam_search(next_logprobs, beam_width, eos=None, max_len=None, start=(), num_results=None)[source]

Approximate top sequences of an autoregressive model by beam search.

Keeps at most beam_width live prefixes per step (the highest-scoring ones); a prefix that emits eos or reaches max_len is finalized. Returns the finalized sequences sorted by total log-probability. This is the standard heuristic decode – faster than exact best-first but not guaranteed to return the true top-k.

Parameters:

  • next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]]) – next_logprobs(prefix) -> [(token, log_prob), ...] (see best_first_decode).

  • beam_width (int) – number of prefixes kept per step.

  • eos (Any) – end-of-sequence token (optional).

  • max_len (int | None) – maximum length (optional, but recommended to bound the search).

  • start (tuple) – the initial prefix.

  • num_results (int | None) – number of sequences to return (default beam_width).

Returns:

A list of (sequence_tuple, total_log_prob) sorted by nonincreasing log-probability.

Return type:

list[tuple[tuple, float]]

top_k_scored(candidates, score, k=None)[source]

Return a finite candidate set scored by score, sorted in descending score.

For a classifier: candidates are the class labels and score is lambda c: model.log_prob(c | x).

Parameters:

  • candidates (Iterable[Any]) – a finite iterable of candidate outputs.

  • score (Callable[[Any], float]) – the (log-)score of a candidate.

  • k (int | None) – keep only the top k (uses a bounded heap); None returns all, sorted.

Returns:

A list of (candidate, score) in nonincreasing score.

Return type:

list[tuple[Any, float]]

quantized_best_first_decode(next_logprobs=None, eos=None, max_len=None, top_k=None, top_p=None, bucket_bits=12, batch_next_logprobs=None, batch_size=64, start=(), max_results=None, min_mass=None)[source]

Fast descending-probability sequence enumeration specialized for neural / transformer decoders.

Three structure-aware accelerations over best_first_decode():

  1. Nucleus / top-k pruning. Neural next-token distributions are sharply peaked, so each step is restricted to its top_k tokens or its top_p nucleus – dropping the long low-probability tail collapses the branching factor (a ~50k vocab down to a handful) at negligible mass loss.

  2. Quantized bucket priority queue. Cumulative log-probs only decrease, so instead of an O(log n) comparison heap the frontier is bucketed by quantized score (bucket = floor(score * 2**bucket_bits)) and drained highest-bucket first – O(1) pushes/pops, and prefixes of near-equal score are grouped. Buckets are disjoint score ranges, so order is exact across buckets and within ~2**-bucket_bits inside one.

  3. Batched scoring. The cost is dominated by model forward passes. Pass batch_next_logprobs to score up to batch_size frontier prefixes in one call (one padded GPU forward) instead of one at a time.

With top_k=top_p=None and a large bucket_bits this reduces to the exact enumeration; pruning is the only approximation (report min_mass to stop once enough probability is covered).

Parameters:

  • next_logprobs (Callable[[tuple], Iterable[tuple[Any, float]]] | None) – next_logprobs(prefix) -> [(token, log_prob), ...] (used when batch_next_logprobs is not given). Log-probs must be <= 0.

  • eos (Any) – end-of-sequence token (a prefix ending in eos is complete).

  • max_len (int | None) – maximum sequence length. Give eos and/or max_len.

  • top_k (int | None) – keep only the top_k highest-probability tokens per step.

  • top_p (float | None) – keep the smallest set of tokens per step whose probability sums to >= top_p (nucleus).

  • bucket_bits (int) – score-quantization resolution; larger = finer ordering, slower bookkeeping.

  • batch_next_logprobs (Callable[[list[tuple]], list[Iterable[tuple[Any, float]]]] | None) – optional batch_next_logprobs([prefix, ...]) -> [[(token, log_prob), ...], ...] scoring a batch of prefixes in one forward pass.

  • batch_size (int) – number of frontier prefixes expanded per (batched) scoring call.

  • start (tuple) – initial prefix.

  • max_results (int | None) – stop after this many complete sequences.

  • min_mass (float | None) – stop once the yielded sequences cover at least this much probability mass.

Yields:

(sequence_tuple, total_log_prob), highest probability first (exact across score buckets).

Return type:

Iterator[tuple[tuple, float]]