dew.objectives

Name	Summary
Aux	Everything a loss returns besides its statistics. Documented in dew.objectives.base.
DistillationObjective	Mix the student’s own loss with a frozen teacher’s soft targets.
EMASpec	Say which leaves the EMA copy tracks, and how fast it follows them. Documented in dew.objectives.base.
Mean	Carry a scalar sum together with the mass it is averaged over. Documented in dew.objectives.base.
Objective	Define what is being learned: the parameters, the loss, what evaluation produces. Documented in dew.objectives.base.
Prediction	Hold what a token objective scored a batch with, for a teacher to compare. Documented in dew.objectives.base.
Step	What the trainer tells an objective about the current step. Documented in dew.objectives.base.
mean_loss	Reduce a shared-denominator estimator, including empty support. Documented in dew.objectives.base.
scalar_loss	Evaluate and reduce canonical statistics for direct JAX differentiation. Documented in dew.objectives.base.

DistillationObjective

class source

class DistillationObjective(
    student: Objective[Loss, Effects],
    teacher: Objective[Loss, Effects],
    *,
    alpha: Weight = 0.5,
    temperature: Weight = 1.0,
    features: Sequence[tuple[int, int]] = (),
    beta: Weight = 0.0,
    feature_loss: FeatureLoss = 'cosine',
)

Mix the student’s own loss with a frozen teacher’s soft targets.

DistillationObjective.held_variables

def held_variables() -> Variables | None

Return the student’s held tree, if any, with the teacher’s under teacher.

DistillationObjective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

DistillationObjective.student_variables

def student_variables(params: Variables) -> Variables

Cut the student’s own tree out of the whole.

That is what the student’s methods read, and what a distilled checkpoint hands on to a plain student run.

DistillationObjective.loss

def loss(params: Variables, batch: Batch, step: Step) -> tuple[Mean, Aux[Effects]]

DistillationObjective.apply_effects

def apply_effects(variables: Variables, effects: Effects) -> Variables

DistillationObjective.evaluate

def evaluate(params: Variables, batch: Batch, step: Step) -> Artifacts | None

DistillationObjective.preview

def preview(
    params: Variables,
    batch: Batch,
    step: Step,
    *,
    scored: Artifacts | None = None,
) -> Artifacts | None

DistillationObjective.pipeline

def pipeline(state: TrainState, *, ema: bool = True)

The student as its inference task; the teacher stays behind.