mixle.models.dpo_leaf module

Direct Preference Optimization (DPO) as a mixle leaf – alignment as a supervised preference likelihood.

Observation = (x, chosen, rejected): a context and a preferred vs dispreferred action/completion. The leaf carries a POLICY module and a FROZEN REFERENCE module; seq_log_density returns the DPO log-sigmoid reward

log sigma( beta * [ (log pi(chosen|x) - log pi_ref(chosen|x)) - (log pi(rejected|x) - log pi_ref(rejected|x)) ] )

(higher = the policy prefers chosen over rejected, relative to the reference). The M-step gradient-steps the policy; the reference stays frozen. No reward model, no RL – the alignment stage of the LLM pipeline as a likelihood, on the same substrate as pretrain/CPT/SFT.

This is the genuinely-new paired leaf the design flagged: it couples two forward passes plus a frozen reference, so it does not reduce to a single Categorical (the log_density contract is over a pair, not a single token). It composes through the same estimate() driver; the M-step owns the policy optimizer.

class DPOModel(policy, ref, beta=0.1, m_steps=100, lr=1e-3, device='cpu')

Bases: SequenceEncodableProbabilityDistribution

DPO over (x, chosen, rejected) preference triples. policy is trained, ref is frozen.

Parameters:

• policy (Any)

• ref (Any)

• beta (float)

• m_steps (int)

• lr (float)

• device (str)

seq_log_density(enc)

Return vectorized log-density values for sequence-encoded observations.

Parameters:

enc (Any)

Return type:

ndarray

log_density(xcr)

Return the log-density or log-mass at a single observation.

Parameters:

xcr (Any)

Return type:

float

prefers(x)

The policy’s argmax action at x – what the aligned policy now picks.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return a sampler for drawing observations from this distribution.

Parameters:

seed (int | None)

Return type:

DPOModelSampler

estimator(pseudo_count=None)

Return an estimator for fitting this distribution from data.

Parameters:

pseudo_count (float | None)

Return type:

DPOModelEstimator

dist_to_encoder()

Return the data encoder used by this distribution for vectorized methods.

Return type:

DPOEncoder

to_dict()

Return a safe JSON-compatible representation of this distribution.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Reconstruct a distribution from to_dict output.

Parameters:

payload (dict[str, Any])

Return type:

DPOModel

class DPOModelSampler(dist, seed=None)

Bases: DistributionSampler

Parameters:

• dist (DPOModel)

• seed (int | None)

sample(size=None, *, batched=True)

Draw observations.

Combinator samplers (mixture/sequence/…) accept batched. With batched=True (the default) each child stream is drawn in one vectorized call instead of a per-draw Python loop – far faster. Because every child sampler owns an independent RandomState, batching consumes each stream in the same order as the loop, so the draws are identical to the legacy path. batched=False forces that legacy per-draw loop as a guaranteed- stable reference. Leaf samplers are already vectorized and ignore the flag.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class DPOEncoder

Bases: DataSequenceEncoder

seq_encode(data)

Encode the iid observation sequence x for vectorized evaluation.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray, ndarray]

class DPOAccumulator

Bases: SequenceEncodableStatisticAccumulator

update(xcr, weight, estimate)

Parameters:

• xcr (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Parameters:

• enc (Any)

• weights (Any)

• estimate (Any)

Return type:

None

initialize(xcr, weight, rng)

Parameters:

• xcr (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Parameters:

• enc (Any)

• weights (Any)

• rng (Any)

Return type:

None

combine(other)

Parameters:

other (Any)

Return type:

DPOAccumulator

value()

Return type:

tuple

from_value(v)

Parameters:

v (tuple)

Return type:

DPOAccumulator

acc_to_encoder()

Return type:

DPOEncoder

class DPOAccumulatorFactory

Bases: StatisticAccumulatorFactory

make()

Return type:

DPOAccumulator

class DPOModelEstimator(policy, ref, beta, m_steps, lr, device)

Bases: ParameterEstimator

DPO M-step: m_steps of gradient on the POLICY minimizing -log sigmoid(beta * margin); ref frozen.

Parameters:

• policy (Any)

• ref (Any)

• beta (float)

• m_steps (int)

• lr (float)

• device (str)

accumulator_factory()

Return type:

DPOAccumulatorFactory

estimate(nobs, suff_stat)

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

DPOModel

DPOLeaf

alias of DPOModel

DPOLeafEstimator

alias of DPOModelEstimator