mixle.experimental.program module

Declarative optimization programs (differentiable games) over parameter groups and objectives.

The whole zoo of “fit this model” – supervised, LoRA fine-tuning, multi-objective, GANs, constrained optimization, policy-gradient RL, continual learning, and classical EM – is one idea at the optimization level: a program built from MOVES and COMBINATORS.

• A move is minimize / maximize an objective over scoped parameters (minimize(), maximize()). An em(...) step on a mixle estimator is also a move, so probabilistic and neural models compose in one program.

• Scoped parameters decide which tensors move: trainable(), freeze(), subset(), and lora() (low-rank adapters – the base stays frozen).

• Combinators schedule the moves: weighted() (cooperative / multi-objective), alternate() (adversarial / coordinate / EM). constrain() adds a Lagrange-multiplier dual player (constrained optimization is the same min-max game as a GAN). reinforce() turns a sampled reward into a score-function objective (RL).

• fit() runs the program.

Examples:

fit(minimize(nll, over=trainable(net)))                                  # supervised
fit(minimize(lm_loss, over=lora(model, rank=8)))                         # LoRA fine-tune
fit(weighted([(recon, 1.0), (kl, beta)], over=trainable([enc, dec])))    # multi-objective (VAE)
fit(alternate(minimize(d_loss, over=D), minimize(g_loss, over=G)))       # GAN
fit(minimize(f, over=th), constraints=[constrain(g, 0.0, "<=")])         # constrained (primal-dual)
fit(maximize(reinforce(sample_reward), over=policy))                     # policy-gradient RL
fit(weighted([(new_loss, 1.0), (replay_loss, 1.0)], over=net))           # continual learning (replay)

Torch is imported lazily, so this module imports without it; the gradient moves require torch, the em move requires only mixle estimators.

trainable(module)

All requires_grad parameters of a module (or a list of modules).

Parameters:

module (Any)

Return type:

list

freeze(module)

Freeze a module in place (requires_grad = False) and return it – e.g. a teacher / old checkpoint.

Parameters:

module (Any)

Return type:

Any

subset(module, *name_substrings)

Trainable parameters whose name contains any of name_substrings (partial fine-tuning).

Parameters:

• module (Any)

• name_substrings (str)

Return type:

list

class LoRALinear(base, rank, alpha)

Bases: object

A torch.nn.Module wrapping a frozen Linear with a trainable low-rank adapter B @ A.

Parameters:

• base (Any)

• rank (int)

• alpha (float)

Return type:

Any

lora(module, rank=8, alpha=16.0)

Replace every Linear under module with a LoRA adapter (base frozen) and return the adapter params.

The module’s forward now routes through low-rank adapters; only the A/B matrices are trainable, so minimize(loss, over=lora(model, rank=8)) fine-tunes a large model cheaply.

Parameters:

• module (Any)

• rank (int)

• alpha (float)

Return type:

list

class Move(objective, params, sign, lr=None)

Bases: object

Minimize (sign=+1) or maximize (sign=-1) objective() over params (a list of tensors).

Parameters:

• objective (Callable[[], Any])

• params (Iterable)

• sign (float)

• lr (float | None)

minimize(objective, over, lr=None)

Parameters:

• objective (Callable[[], Any])

• over (Iterable)

• lr (float | None)

Return type:

Move

maximize(objective, over, lr=None)

Parameters:

• objective (Callable[[], Any])

• over (Iterable)

• lr (float | None)

Return type:

Move

class EMMove(estimator, data, init)

Bases: object

An EM step on a mixle estimator – a first-class move, so stats models join the program.

move.model is the current fitted distribution; other moves’ objectives may read it (e.g. a gating network reading the mixture’s responsibilities), making neural<->stats coupling one program.

Parameters:

• estimator (Any)

• data (Sequence)

• init (Any)

property model: Any

em(estimator, data, init)

A mixle EM step as a move. init is the starting distribution (the E-step needs a current model).

Parameters:

• estimator (Any)

• data (Sequence)

• init (Any)

Return type:

EMMove

class Program(moves)

Bases: object

Parameters:

moves (Sequence)

alternate(*items)

Run each move (or sub-program) once per round, in order – GANs, EM, coordinate ascent.

Parameters:

items (Any)

Return type:

Program

weighted(terms, over)

A single move minimizing sum(w * objective() for objective, w in terms) – cooperative multi-objective.

Parameters:

• terms (Sequence[tuple])

• over (Iterable)

Return type:

Program

class Constraint(g, bound=0.0, kind='<=')

Bases: object

Parameters:

• g (Callable[[], Any])

• bound (float)

• kind (str)

violation()

Return type:

Any

constrain(g, bound=0.0, kind='<=')

g() <= bound (or >=). Passed to fit() as constraints=[...]; enforced by dual ascent.

Parameters:

• g (Callable[[], Any])

• bound (float)

• kind (str)

Return type:

Constraint

reinforce(sample_and_reward)

Wrap sample_and_reward() -> (log_probs, rewards) into the score-function surrogate E[r·logπ].

Maximizing it gives the policy gradient E[r·∇logπ]. log_probs are the log-probabilities of the sampled actions (carry grad); rewards are detached returns.

Parameters:

sample_and_reward (Callable[[], tuple])

Return type:

Callable[[], Any]

fit(program, steps=1000, lr=1e-3, constraints=None, callback=None, data=None, steps_per_chunk=1)

Run an optimization program; each round runs every move once, in order.

Gradient moves take one optimizer step (Adam, per-move learning rate move.lr or the global lr); em moves take one EM step. constraints add dual-ascent multiplier moves (primal-dual).

Fixed mode (data is None): run steps rounds. Streaming mode (data is a Stream): advance through the data chunks – the parameters/optimizers persist across chunks (warm-started), running steps_per_chunk rounds per chunk; objectives read stream.current. This is the continuous-pretraining loop (combine the task loss with an anti-forget term via weighted()). Returns the program updated in place.

Parameters:

• program (Any)

• steps (int)

• lr (float)

• constraints (Sequence[Constraint] | None)

• callback (Callable[[int, Program], None] | None)

• data (Stream | None)

• steps_per_chunk (int)

Return type:

Program

class Stream(chunks)

Bases: object

A holder over an iterable of data chunks. fit(data=stream) advances it each round; objectives read stream.current (the active chunk). The model’s parameters persist across chunks (warm-started).

Parameters:

chunks (Iterable)

advance()

Return type:

None

class ReplayBuffer(capacity=16)

Bases: object

Fixed-capacity FIFO of past chunks, for replay-based anti-forgetting.

Parameters:

capacity (int)

add(item)

Parameters:

item (Any)

Return type:

ReplayBuffer

all()

Return type:

list

snapshot(params)

Detached clones of params – the anchor for ewc() / L2-SP regularization.

Parameters:

params (Iterable)

Return type:

list

replay(loss_fn, buffer)

An objective averaging loss_fn(chunk) over the replay buffer – a term for weighted().

Parameters:

• loss_fn (Callable[[Any], Any])

• buffer (ReplayBuffer)

Return type:

Callable[[], Any]

distill(student_out, teacher_out)

MSE distillation: keep student outputs near a frozen teacher’s (logit/feature matching anti-forget).

Parameters:

• student_out (Callable[[], Any])

• teacher_out (Callable[[], Any])

Return type:

Callable[[], Any]

ewc(params, fisher, anchor, weight=1.0)

Elastic Weight Consolidation penalty weight · Σ Fᵢ (θᵢ - anchorᵢ)² – anchors params important to the old task. Pair with fisher_diagonal() (a torch net) or a mixle leaf’s to_fisher.

Parameters:

• params (Iterable)

• fisher (Sequence)

• anchor (Sequence)

• weight (float)

Return type:

Callable[[], Any]

fisher_diagonal(net, batches, kind='classification')

Diagonal Fisher information for EWC, using MODEL-sampled labels so it does NOT vanish at convergence.

kind='classification' (net returns logits) or 'regression' (net returns a Gaussian mean). Returns one tensor per trainable parameter. (The naive data-label Fisher is ~0 at a converged optimum – sampling the label from the model is what makes EWC actually anchor.)

Parameters:

• net (Any)

• batches (Iterable)

• kind (str)

Return type:

list

bilevel(model, inner_loss, outer_loss, sample_tasks, inner_steps=1, inner_lr=0.01)

Meta-learning (MAML): meta-learn model’s params so a few inner gradient steps adapt to a task.

sample_tasks() yields (support, query) batches; inner_loss(forward, support) and outer_loss(forward, query) return scalars, where forward(x) runs the model with the current (adapted) parameters. The returned move’s objective differentiates through the inner adaptation (second-order), so fit(bilevel(...), steps=N) is MAML; fit minimizes the query loss over the meta-parameters.

Parameters:

• model (Any)

• inner_loss (Callable[[Callable, Any], Any])

• outer_loss (Callable[[Callable, Any], Any])

• sample_tasks (Callable[[], Iterable[tuple]])

• inner_steps (int)

• inner_lr (float)

Return type:

Move

class ParetoMove(objectives, params, lr=None)

Bases: Move

A move that steps along the MGDA common-descent direction – decreases every objective at once.

Parameters:

• objectives (Sequence[Callable[[], Any]])

• params (Iterable)

• lr (float | None)

pareto(objectives, over, lr=None)

Multi-objective optimization with NO fixed weights: each step descends all objectives at once (MGDA).

Parameters:

• objectives (Sequence[Callable[[], Any]])

• over (Iterable)

• lr (float | None)

Return type:

Program

class StreamingEMMove(estimator, stream, init, iters_per_chunk=1)

Bases: object

Online EM over a Stream: each round, warm-started EM iterations on the current chunk.

Composes in fit(data=stream) next to gradient moves, so a stats model and a neural net adapt to the same stream together (the LLM<->stats continual-coupling case). move.model is the current fit.

Parameters:

• estimator (Any)

• stream (Stream)

• init (Any)

• iters_per_chunk (int)

property model: Any

streaming_em(estimator, stream, init, iters_per_chunk=1)

A stats EM move that continually adapts over a chunk stream (use inside fit(data=stream)).

Parameters:

• estimator (Any)

• stream (Stream)

• init (Any)

• iters_per_chunk (int)

Return type:

StreamingEMMove

gail(discriminator, sample_expert, sample_policy, disc_params, policy_params)

GAIL / adversarial inverse RL = alternate(minimize(disc_loss), maximize(reinforce(policy))).

Recover an expert’s behavior (and a reward) from demonstrations alone. discriminator(features) -> logits (high = expert; this logit is the recovered reward). sample_expert() -> features is a batch of expert transition features; sample_policy() -> (features, action_logprobs) is a policy rollout. The discriminator separates expert from policy transitions while the policy is reinforced to fool it.

Parameters:

• discriminator (Callable[[Any], Any])

• sample_expert (Callable[[], Any])

• sample_policy (Callable[[], tuple])

• disc_params (Iterable)

• policy_params (Iterable)

Return type:

Program

maxent_irl(reward, reward_params, expert_features, policy_features)

Maximum-entropy inverse RL by feature matching (Ziebart et al.).

reward(features) -> scalar (e.g. w·φ). expert_features is the expert’s expected feature vector (from demonstrations). policy_features() -> features returns the expected features under the maxent-optimal policy for the current reward – the inner forward / soft-value solve, recomputed each step. For structured dynamics that inner solve is exactly mixle’s forward (soft-value) pass, so the partition function is computed without sampling. The move’s gradient matches expert to policy features.

Parameters:

• reward (Callable[[Any], Any])

• reward_params (Iterable)

• expert_features (Any)

• policy_features (Callable[[], Any])

Return type:

Move