…ent (SPEC-DISTILL-001 Phase 2)

Wires Phase 1's teacher provider into a per-batch KD orchestration step
that produces both the combined α·CE + (1-α)·T²·KL scalar loss (for
telemetry) and the KD-aware logit-space gradient (Phase 2b plug point).

What this PR adds
=================

New module `aprender-train-distill::kd_step`:

  pub fn kd_loss(
      student_logits: &[f32],
      teacher_logits: &[f32],
      label: usize,
      temperature: f32,
      alpha: f32,
  ) -> f32;

  pub fn kd_logit_gradient(
      student_logits: &[f32],
      teacher_logits: &[f32],
      label: usize,
      temperature: f32,
      alpha: f32,
  ) -> Vec<f32>;

  pub fn kd_step<F: FnMut(&[u32]) -> Vec<f32>>(
      teacher: &mut dyn TeacherLogitsProvider,
      input_ids: &[Vec<u32>],
      labels: &[usize],
      temperature: f32,
      alpha: f32,
      compute_student_logits: F,
  ) -> Result<(f32, Vec<Vec<f32>>)>;

The gradient is the Hinton et al. 2015 §2 derivation:

    ∂L/∂s = α · (softmax(s) - one_hot(label))
          + (1-α) · T · (softmax(s/T) - softmax(t/T))

(T factor, not T² — one T factor is absorbed by the softmax derivative
chain rule.)

Scope: Phase 2 vs Phase 2b
==========================

Phase 2 (this PR) ships the orchestration math, all in pure Rust on the
CPU. The output `Vec<Vec<f32>>` of per-batch gradients is what Phase 2b
will plumb into `CudaTransformerTrainer.forward_backward_kd_batch` as
the backward-pass seed (replacing the CE-only gradient currently used by
forward_backward_batch).

Splitting Phase 2 into 2a/2b lets us land the orchestration layer + its
tests now, separate from extending the complex GPU trainer code path.

Falsifiers pinned
=================

3 KD-step falsifiers + 6 sanity tests, all passing:

- F-DISTILL-KDSTEP-001 (alpha=1 → pure CE)
- F-DISTILL-KDSTEP-002 (student==teacher → zero KL gradient under alpha=0)
- F-DISTILL-KDSTEP-003 (loss monotone in student-teacher divergence)
- softmax unit-sum + non-negative
- CE gradient correct sign at label vs non-label positions
- kd_step orchestration end-to-end
- kd_step empty-batch sanity
- kd_step vocab-mismatch error path
- kd_loss alpha=1 collapses to pure CE

All 50 aprender-train-distill lib tests pass (was 41 — 9 new).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

…TILL-001 Phase 2b)

Mirrors Phase 1's TeacherLogitsProvider for the student side. The
student has two methods: logits_for_batch (forward) and
apply_kd_gradient (backward + optimizer step). FixtureStudent
implements both for CPU-only unit testing — Phase 2c will add a
CudaStudentProvider that wraps CudaTransformerTrainer.

What this PR adds
=================

  pub trait StudentLogitsProvider {
      fn vocab_size(&self) -> usize;
      fn logits_for_batch(&mut self, input_ids: &[Vec<u32>])
          -> Result<Vec<Vec<f32>>>;
      fn apply_kd_gradient(&mut self, gradient: &[Vec<f32>])
          -> Result<()>;
  }

  pub struct FixtureStudent {
      vocab_size: usize,
      logits: Vec<f32>,           // current student parameters
      learning_rate: f32,
  }

FixtureStudent's apply_kd_gradient averages the gradient across batch
elements (canonical SGD batch averaging) and subtracts the scaled
gradient from its internal logits buffer. This isn't a real model —
it's a logit-space optimization fixture that lets us validate the
KD pipeline's gradient direction is correct without needing CUDA.

Falsifiers pinned
=================

7 student_provider tests + 2 falsifiers, all passing:

- F-DISTILL-STUDENT-001 — one KD step moves student logits toward
  teacher's preferred token. Setup: uniform student, teacher prefers
  token 5, alpha=0 (pure KL signal). After one step, student logit
  at index 5 must be strictly greater than before.

- F-DISTILL-STUDENT-002 — 10 sequential KD steps strictly decrease
  per-step KD loss. With LR=0.5, loss after 10 steps < 90% of initial.
  Validates the gradient direction is correct (descent, not ascent).

Plus 5 sanity tests covering vocab_size reporting, batch broadcast,
shape validation, in-place logit update, and batch averaging math.

Architecture
============

Stacks on top of #1788 (kd_step). Pipeline integration that uses both
TeacherLogitsProvider + StudentLogitsProvider + kd_step is Phase 2c.

Phase 2c (PMAT-696, follow-up): CudaStudentProvider that wraps
CudaTransformerTrainer for production runs. Once it lands, end-to-end
GPU distillation is unblocked.

Tests
=====

All 57 aprender-train-distill lib tests pass (was 50 — 7 new).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>