/**
 * moe_trainer.ts — AdamW training loop for {@link SharedExpertMoE}.
 *
 * Makes "train your own Evermind AI" real: given labelled (input → target)
 * samples, runs minibatch AdamW over the flat parameters with the load-balancing
 * auxiliary loss mixed in (so the router spreads load instead of collapsing onto
 * a few experts). Pure CPU, deterministic given a seeded model — the same loop a
 * WebGPU optimiser kernel would accelerate.
 */
import { SharedExpertMoE } from "./moe_model.js";
export interface MoESample {
    input: ArrayLike<number>;
    target: ArrayLike<number>;
}
export interface MoETrainOptions {
    /** Learning rate. Default 0.01. */
    lr?: number;
    /** AdamW β1. Default 0.9. */
    beta1?: number;
    /** AdamW β2. Default 0.999. */
    beta2?: number;
    /** AdamW ε. Default 1e-8. */
    eps?: number;
    /** Decoupled weight decay. Default 0. */
    weightDecay?: number;
    /** Weight of the load-balancing auxiliary loss. Default 0.01. */
    auxWeight?: number;
    /** Minibatch size. Default = all samples (full batch). */
    batchSize?: number;
    /** Passes over the dataset. Default 1. */
    epochs?: number;
}
export interface MoEEpochResult {
    /** Mean per-sample task (MSE·½) loss over the epoch. */
    loss: number;
    /** Load-balancing auxiliary loss at the end of the epoch (≈1 balanced … E collapsed). */
    auxLoss: number;
}
/**
 * AdamW optimiser over a model's flat parameter list. State (m, v) is keyed by
 * parameter index and persists across {@link step} calls.
 */
export declare class MoETrainer {
    private readonly model;
    private readonly adam;
    private readonly opt;
    constructor(model: SharedExpertMoE, options?: MoETrainOptions);
    /** Train for the configured epochs. Returns the per-epoch loss history. */
    fit(samples: MoESample[]): MoEEpochResult[];
    private runEpoch;
    private scaleGradients;
}
//# sourceMappingURL=moe_trainer.d.ts.map