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 per-row DPO preference log likelihoods for encoded triples.

Parameters:

enc (Any)

Return type:

ndarray

log_density(xcr)

Return the DPO log likelihood for one (x, chosen, rejected) triple.

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 the sampler for the preference-scoring leaf.

Parameters:

seed (int | None)

Return type:

DPOModelSampler

estimator(pseudo_count=None)

Return the DPO estimator that trains the policy while keeping the reference fixed.

Parameters:

pseudo_count (float | None)

Return type:

DPOModelEstimator

dist_to_encoder()

Return the encoder for preference triples.

Return type:

DPOEncoder

to_dict()

Serialize policy/reference modules and DPO hyperparameters.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a DPOModel from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

DPOModel

class DPOModelSampler(dist, seed=None)

Bases: DistributionSampler

Sampler facade for DPO leaves, which score preference pairs rather than generating.

Parameters:

• dist (DPOModel)

• seed (int | None)

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

Raise because DPO is a preference-scoring likelihood, not a generator.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class DPOEncoder

Bases: DataSequenceEncoder

Encode (context, chosen, rejected) preference triples for DPO.

seq_encode(data)

Convert preference triples into batched contexts and integer action arrays.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray, ndarray]

class DPOAccumulator

Bases: SequenceEncodableStatisticAccumulator

Buffer weighted preference triples for the DPO M-step.

update(xcr, weight, estimate)

Add one weighted preference triple to the accumulator.

Parameters:

• xcr (Any)

• weight (float)

• estimate (Any)

Return type:

None

seq_update(enc, weights, estimate)

Add an encoded batch of preference triples and optional weights.

Parameters:

• enc (Any)

• weights (Any)

• estimate (Any)

Return type:

None

initialize(xcr, weight, rng)

Initialize from one preference triple using the ordinary update path.

Parameters:

• xcr (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Initialize from an encoded batch using the ordinary batch update path.

Parameters:

• enc (Any)

• weights (Any)

• rng (Any)

Return type:

None

combine(other)

Merge the value tuple from another DPO accumulator.

Parameters:

other (Any)

Return type:

DPOAccumulator

value()

Return buffered contexts, chosen actions, rejected actions, and weights.

Return type:

tuple

from_value(v)

Restore accumulator buffers from a value tuple.

Parameters:

v (tuple)

Return type:

DPOAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

DPOEncoder

class DPOAccumulatorFactory

Bases: StatisticAccumulatorFactory

Factory for DPO accumulators.

make()

Create a fresh accumulator.

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 an accumulator factory for weighted preference triples.

Return type:

DPOAccumulatorFactory

estimate(nobs, suff_stat)

Run the weighted DPO M-step and return the updated policy leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

DPOModel

DPOLeaf

alias of DPOModel

DPOLeafEstimator

alias of DPOModelEstimator