From 05ec7a0aa73ee37a38a5612c0c986f7e207c8e74 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 22 Oct 2026 07:16:10 -0400 Subject: [PATCH 1/2] eoc-route-learned: recommend the stage cheapest in expectation, not the likeliest The crate says a learned router "predicts the cheapest stage that is likely to satisfy the query". Three of its four routers recommended the stage likeliest to succeed: MF and both bandits took an argmax over success, and the bandits ignored the joule cost `observe` passed them. Neural succeeds on everything and costs the most, so they learned to send every query there. The synthetic LinUCB test pinned exactly that as the goal. - router: `expected_joules` prices a skip the way LearnedCascade dispatches it (Kv or Graph falling through to Neural, Neural alone, Cache as the full walk), and `cheapest_in_expectation` picks the least. A broken probability prices a stage dearer, never cheaper. `observe` is the outcome of one stage that ran, with that stage's own joules. - MF picks by that rule over its per-stage chances and learned costs. - LinUCB and Thompson learn each stage's cost and pick by the rule over optimistic and sampled chances. Confidence is the mean belief, not the bound or sample used to explore. Older saved states import with the default costs. - train_logreg labels each query with its cheapest successful stage; the classifier's online update no longer reinforces a correct stage costlier than the one it predicts. - evaluate replays each recommendation over per-stage outcomes. It charged every row's own cost whatever the router chose and priced the walk at 4x that, so every router showed a lift of 3x its spend. A query it cannot price is counted unscored, not charged zero, and joules_per_correct is None when nothing was correct. Co-Authored-By: Claude Opus 5.5 --- .../crates/eoc-route-learned/src/bandit.rs | 178 +++++++++--- .../eoc-route-learned/src/classifier.rs | 55 +++- .../eoc-route-learned/src/evaluation.rs | 256 +++++++++++++++--- .../crates/eoc-route-learned/src/lib.rs | 7 +- .../src/matrix_factorization.rs | 62 +---- .../crates/eoc-route-learned/src/router.rs | 197 +++++++++++++- .../crates/eoc-route-learned/src/training.rs | 87 +++++- .../eoc-route-learned/tests/linucb_bandit.rs | 229 ++++++++++++---- .../eoc-route-learned/tests/mf_synthetic.rs | 221 ++++++++++++--- .../tests/state_roundtrip.rs | 42 +++ 10 files changed, 1084 insertions(+), 250 deletions(-) diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs index 4d3c093d6..cae5f7214 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/bandit.rs @@ -3,10 +3,20 @@ //! Two algorithms ship here: //! //! - **LinUCB** (Li et al. 2010) — contextual bandit with linear reward -//! models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`. Pick -//! the arm that maximises `θ_a · x + α · √(x^T A^{-1} x)`. +//! models. Each arm (stage) keeps `A = X^T X + I` and `b = X^T r`, where +//! `r` is 1 when the stage resolved the query and 0 when it did not. //! - **Thompson Sampling** — Beta-Bernoulli on per-arm success rates. Cheap //! and surprisingly competitive when contexts are noisy. +//! +//! Both learn each stage's chance of resolving a query, and the cost of each +//! stage from the joules `observe` reports. Both pick the stage that is +//! cheapest to answer from in expectation +//! ([`cheapest_in_expectation`](crate::router::cheapest_in_expectation)), +//! not the arm with the best success rate: the neural stage always has the +//! best rate, and picking it skips every cheaper stage that would have +//! answered. Exploration enters through the success estimate, optimistic for +//! LinUCB (`θ_a · x + α · √(x^T A^{-1} x)`) and sampled for Thompson, so an +//! untried cheap stage looks likely to resolve and gets tried. use std::collections::{BTreeMap, HashMap}; @@ -18,7 +28,10 @@ use rand_distr::{Beta, Distribution}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ + ALL_STAGES, LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation, + default_stage_joules, stage_index, +}; /// Per-arm state for [`LinUcbRouter`]. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -52,7 +65,8 @@ impl LinUcbArm { } } - fn ucb(&self, x: &DVector, alpha: f32) -> f32 { + /// The mean reward estimate `θ · x` and its upper confidence bound. + fn estimate(&self, x: &DVector, alpha: f32) -> (f32, f32) { let theta = self.theta(); let mean = theta.dot(x); let a_inv = self @@ -61,7 +75,7 @@ impl LinUcbArm { .try_inverse() .unwrap_or_else(|| DMatrix::::identity(self.dim, self.dim)); let bonus_sq = (x.transpose() * &a_inv * x)[(0, 0)].max(0.0); - mean + alpha * bonus_sq.sqrt() + (mean, mean + alpha * bonus_sq.sqrt()) } fn update(&mut self, x: &DVector, reward: f32) { @@ -71,6 +85,10 @@ impl LinUcbArm { } /// LinUCB contextual bandit over the four cascade stages. +/// +/// The reward model is linear in the context with no intercept, as in Li et +/// al.: put a constant feature in the context for a stage's base rate to be +/// learnable. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LinUcbRouter { /// Context (embedding) dimension. @@ -81,6 +99,9 @@ pub struct LinUcbRouter { pub ridge: f32, /// Per-stage arm state. pub arms: HashMap, + /// Per-stage cost estimate in microjoules, learned from `observe`. + #[serde(default = "default_stage_joules")] + pub joules: HashMap, } impl LinUcbRouter { @@ -95,11 +116,14 @@ impl LinUcbRouter { alpha, ridge, arms, + joules: default_stage_joules(), } } - /// Argmax UCB across stages. - pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> { + /// Each stage's mean estimate and upper confidence bound that it resolves + /// the query, in [`ALL_STAGES`] order. A missing arm counts as a certain + /// miss. + fn estimates(&self, embedding: &[f32]) -> Result<([f32; 4], [f32; 4])> { if embedding.len() != self.dim { return Err(Error::DimensionMismatch { expected: self.dim, @@ -107,20 +131,27 @@ impl LinUcbRouter { }); } let x = DVector::::from_row_slice(embedding); - let mut best = Stage::Cache; - let mut best_score = f32::NEG_INFINITY; - for s in ALL_STAGES { - let arm = match self.arms.get(&s) { - Some(a) => a, - None => continue, - }; - let u = arm.ucb(&x, self.alpha); - if u > best_score { - best_score = u; - best = s; + let mut means = [0.0_f32; 4]; + let mut bounds = [0.0_f32; 4]; + for (i, s) in ALL_STAGES.iter().enumerate() { + if let Some(arm) = self.arms.get(s) { + (means[i], bounds[i]) = arm.estimate(&x, self.alpha); } } - Ok((best, best_score)) + Ok((means, bounds)) + } + + /// The stage cheapest to answer from in expectation, with every stage's + /// chance taken at its upper confidence bound, and that stage's bound. + pub fn pick(&self, embedding: &[f32]) -> Result<(Stage, f32)> { + let (_, bounds) = self.estimates(embedding)?; + let stage = cheapest_in_expectation(&bounds, &self.joules); + Ok((stage, bounds[stage_index(stage)])) + } + + /// Record an observed stage cost — EMA with coefficient 0.1. + pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { + blend_joules(&mut self.joules, stage, microjoules); } /// Update the chosen arm with the observed reward. @@ -145,13 +176,14 @@ impl LearnedRouter for LinUcbRouter { let Some(emb) = query.embedding.as_ref() else { return StagePrediction::new(Stage::Cache, 0.0); }; - match self.pick(emb) { - Ok((stage, score)) => { - // Squash UCB score into [0,1] via logistic for "confidence". - let conf = 1.0 / (1.0 + (-score).exp()); - StagePrediction::new(stage, conf) + match (self.pick(emb), self.estimates(emb)) { + // The confidence is the mean estimate, not the optimistic bound + // the pick used: the bound is how much the router is willing to + // explore, not how likely the stage is to resolve. + (Ok((stage, _)), Ok((means, _))) => { + StagePrediction::new(stage, means[stage_index(stage)]) } - Err(_) => StagePrediction::new(Stage::Cache, 0.0), + _ => StagePrediction::new(Stage::Cache, 0.0), } } @@ -160,8 +192,9 @@ impl LearnedRouter for LinUcbRouter { query: &Query, chosen_stage: Stage, was_correct: bool, - _joule_cost: u64, + joule_cost: u64, ) { + self.observe_cost(chosen_stage, joule_cost); if let Some(emb) = query.embedding.as_ref() { let reward = if was_correct { 1.0 } else { 0.0 }; let _ = self.update(emb, chosen_stage, reward); @@ -224,6 +257,9 @@ pub struct ThompsonSamplingRouter { pub arms: HashMap, /// RNG seed for reproducibility — incremented every sample. pub seed: u64, + /// Per-stage cost estimate in microjoules, learned from `observe`. + #[serde(default = "default_stage_joules")] + pub joules: HashMap, } impl ThompsonSamplingRouter { @@ -233,30 +269,43 @@ impl ThompsonSamplingRouter { for s in ALL_STAGES { arms.insert(s, BetaArm::new()); } - Self { arms, seed } + Self { + arms, + seed, + joules: default_stage_joules(), + } } - /// Sample once from each arm; return the argmax. Iteration order is the - /// fixed cheapest-first stage order, so the result is deterministic for - /// a given `seed`. + /// Sample each stage's chance of resolving once, and return the stage + /// cheapest to answer from in expectation under those samples, with its + /// sample. Sampling runs in the fixed cheapest-first stage order, so the + /// result is deterministic for a given `seed`. A missing arm counts as a + /// certain miss. pub fn pick(&self) -> (Stage, f32) { let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed); - let mut best = Stage::Cache; - let mut best_v = f32::NEG_INFINITY; - for s in ALL_STAGES { - let arm = match self.arms.get(&s) { - Some(a) => a, - None => continue, + let mut samples = [0.0_f32; 4]; + for (i, s) in ALL_STAGES.iter().enumerate() { + let Some(arm) = self.arms.get(s) else { + continue; }; let dist = Beta::new(arm.alpha as f64, arm.beta as f64) .unwrap_or_else(|_| Beta::new(1.0, 1.0).expect("beta(1,1) ok")); - let v = dist.sample(&mut rng) as f32; - if v > best_v { - best_v = v; - best = s; - } + samples[i] = dist.sample(&mut rng) as f32; } - (best, best_v) + let stage = cheapest_in_expectation(&samples, &self.joules); + (stage, samples[stage_index(stage)]) + } + + /// The posterior mean chance that `stage` resolves a query. + pub fn mean(&self, stage: Stage) -> f32 { + self.arms + .get(&stage) + .map_or(0.0, |a| a.alpha / (a.alpha + a.beta)) + } + + /// Record an observed stage cost — EMA with coefficient 0.1. + pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { + blend_joules(&mut self.joules, stage, microjoules); } /// Update Beta posterior for `stage`. @@ -275,8 +324,11 @@ impl ThompsonSamplingRouter { #[async_trait] impl LearnedRouter for ThompsonSamplingRouter { async fn route(&self, _query: &Query) -> StagePrediction { - let (stage, v) = self.pick(); - StagePrediction::new(stage, v.clamp(0.0, 1.0)) + // The confidence is the posterior mean, not the sample the pick used: + // a sample is a draw for exploration, and would make the threshold + // policy's skip decision a coin toss. + let (stage, _) = self.pick(); + StagePrediction::new(stage, self.mean(stage)) } async fn observe( @@ -284,8 +336,9 @@ impl LearnedRouter for ThompsonSamplingRouter { _query: &Query, chosen_stage: Stage, was_correct: bool, - _joule_cost: u64, + joule_cost: u64, ) { + self.observe_cost(chosen_stage, joule_cost); self.update(chosen_stage, was_correct); } @@ -332,4 +385,39 @@ mod tests { assert_eq!(arm.alpha, 2.0); assert_eq!(arm.beta, 1.0); } + + /// `observe` is how a stage's cost reaches the router; it was ignored. + #[tokio::test] + async fn observing_a_stage_learns_its_cost() { + let q = Query::new("q").with_embedding(vec![1.0, 0.0]); + let mut lin = LinUcbRouter::new(2, 0.1, 1.0); + let mut ts = ThompsonSamplingRouter::new(0); + lin.observe(&q, Stage::Kv, false, 1_050_000).await; + ts.observe(&q, Stage::Kv, false, 1_050_000).await; + // A tenth of the way from the default 50 000. + assert_eq!(lin.joules[&Stage::Kv], 150_000); + assert_eq!(ts.joules[&Stage::Kv], 150_000); + assert_eq!(ts.arms[&Stage::Kv].beta, 2.0); + } + + /// The confidence handed to the threshold policy is the router's belief + /// that the stage resolves, not the bound or sample it explored with: an + /// untried router believes nothing yet. + #[tokio::test] + async fn confidence_is_the_mean_belief_not_the_exploration_bound() { + let q = Query::new("q").with_embedding(vec![1.0, 0.0]); + let lin = LinUcbRouter::new(2, 1.0, 1.0); + let (stage, bound) = lin.pick(&[1.0, 0.0]).unwrap(); + assert!(bound > 0.5, "an untried arm explores optimistically"); + let pred = lin.route(&q).await; + assert_eq!((pred.recommended, pred.confidence), (stage, 0.0)); + + let mut ts = ThompsonSamplingRouter::new(3); + for _ in 0..3 { + ts.update(Stage::Cache, true); + } + let pred = ts.route(&q).await; + assert_eq!(pred.confidence, ts.mean(pred.recommended)); + assert_eq!(ts.mean(Stage::Cache), 0.8); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs index 48e432fe4..f1c3bfb17 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs @@ -21,19 +21,10 @@ use rand_distr::{Distribution, Normal}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction, stage_index}; const NUM_STAGES: usize = 4; -fn stage_index(s: Stage) -> usize { - match s { - Stage::Cache => 0, - Stage::Kv => 1, - Stage::Graph => 2, - Stage::Neural => 3, - } -} - fn stage_from_index(i: usize) -> Stage { ALL_STAGES[i.min(3)] } @@ -179,8 +170,16 @@ impl LearnedRouter for LogRegRouter { was_correct: bool, _joule_cost: u64, ) { - // Only train on successful resolutions — we want P(stage = correct | x). - if was_correct && let Some(emb) = query.embedding.as_ref() { + // The classifier predicts the cheapest stage that resolves a query. + // A stage that resolved it shows that stage is enough, not that it is + // the cheapest that is, so it is evidence only when it is no costlier + // than the stage already predicted. Otherwise a correct Neural answer, + // which every query gets, would teach it to predict Neural for all. + if was_correct + && let Some(emb) = query.embedding.as_ref() + && let Ok((predicted, _, _)) = self.predict(emb) + && stage_index(chosen_stage) <= stage_index(predicted) + { let _ = self.sgd_step(emb, chosen_stage, 0.01); } } @@ -228,4 +227,36 @@ mod tests { let r = LogRegRouter::new(8, 1e-3, 1); assert!(r.predict_probs(&[0.0; 7]).is_err()); } + + /// Neural answers every query correctly. Reinforcing it whenever it did + /// drifted the classifier to predict Neural for everything; only a + /// correct stage no costlier than the prediction may move it. + #[tokio::test] + async fn a_correct_costlier_stage_does_not_move_the_prediction() { + let emb = vec![1.0, -0.5, 0.25]; + let q = Query::new("q").with_embedding(emb.clone()); + let mut r = LogRegRouter::new(3, 0.0, 9); + for _ in 0..200 { + r.sgd_step(&emb, Stage::Kv, 0.1).unwrap(); + } + assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv); + + for _ in 0..500 { + r.observe(&q, Stage::Neural, true, 50_000_000).await; + } + assert_eq!( + r.predict(&emb).unwrap().0, + Stage::Kv, + "Neural being enough is no evidence Kv is not" + ); + + for _ in 0..500 { + r.observe(&q, Stage::Cache, true, 1_000).await; + } + assert_eq!( + r.predict(&emb).unwrap().0, + Stage::Cache, + "a cheaper stage that resolved it is" + ); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs index ba31b8485..7d3f27681 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/evaluation.rs @@ -1,13 +1,25 @@ //! Routing-quality metrics. //! -//! `evaluate` consumes a `LearnedRouter` and a held-out test set, returning -//! `RouterMetrics`: +//! `evaluate` replays a router's recommendations over a held-out set of +//! per-stage outcomes, returning `RouterMetrics`: //! -//! - **accuracy** — fraction of examples where the predicted stage matched -//! the recorded successful stage. -//! - **joules_per_correct** — total observed joules / correct predictions. -//! - **lift** — joules saved vs. always running the full cascade (sum of all -//! stage costs). +//! - **accuracy** — fraction of scored queries the dispatch answered +//! correctly. +//! - **total_microjoules** — what the dispatches the router chose spend. +//! - **full_cascade_microjoules** — what the plain cascade walk spends on the +//! same queries. +//! - **lift** — the walk's joules minus the router's: what routing saved, or +//! cost when negative. +//! - **joules_per_correct** — total joules / correct answers, absent when +//! none was correct. +//! +//! Every joule figure is priced from the recorded per-stage costs, run the way +//! [`LearnedCascade`](crate::inference::LearnedCascade) dispatches: Kv or +//! Graph, falling through to Neural on a miss; Neural alone; or, for Cache, +//! the full walk. A row that did not succeed is a miss. A query whose +//! dispatch or walk would run a stage the set has no row for is not scored: +//! it is counted in `unscored` and enters no figure, rather than being priced +//! at zero. //! //! Latency metrics are deferred to the caller (we don't run the cascade //! here, just predictions). @@ -17,76 +29,238 @@ use std::collections::HashMap; use eoc_core::{Query, Stage}; use serde::{Deserialize, Serialize}; -use crate::router::LearnedRouter; -use crate::training::Example; +use crate::router::{LearnedRouter, stage_index}; +use crate::training::{Example, query_key}; /// Summary statistics produced by [`evaluate`]. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RouterMetrics { - /// Number of test examples scored. + /// Number of queries scored. pub n: usize, - /// Number of correct (predicted == actual successful) predictions. + /// Queries left out because their dispatch or the walk would run a stage + /// the test set has no row for. + pub unscored: usize, + /// Scored queries whose dispatch ended on a correct answer. pub correct: usize, /// `correct / n`. pub accuracy: f32, - /// Sum of observed joule cost. + /// Joules the router's dispatches spend on the scored queries. pub total_microjoules: u64, - /// `total_microjoules / max(correct, 1)`. - pub joules_per_correct: u64, - /// Per-stage prediction count. + /// `total_microjoules / correct`; `None` when no answer was correct. + pub joules_per_correct: Option, + /// Per-stage recommendation count, over every query routed. pub stage_counts: HashMap, - /// Joules saved vs. running the full cascade for every example - /// (sum of all stage costs in the dataset). + /// Joules the plain cascade walk spends on the scored queries. + pub full_cascade_microjoules: u64, + /// `full_cascade_microjoules − total_microjoules`: what routing saved + /// over the walk, negative when it cost more. pub lift_microjoules: i64, } -/// Evaluate a router on a labelled test set. +/// One stage's recorded outcome on a query: whether it resolved it, and what +/// it spent. +type Outcome = Option<(bool, u64)>; + +/// What skipping to `stage` spends on a query and whether it ends correct, +/// or `None` when it would run a stage with no row. +fn dispatch(stage: Stage, rows: &[Outcome; 4]) -> Option<(bool, u64)> { + let neural = rows[stage_index(Stage::Neural)]; + match stage { + Stage::Cache => walk(rows), + Stage::Neural => neural, + Stage::Kv | Stage::Graph => match rows[stage_index(stage)]? { + (true, uj) => Some((true, uj)), + (false, uj) => neural.map(|(ok, n)| (ok, uj.saturating_add(n))), + }, + } +} + +/// The plain cascade: every stage in order until one resolves the query, +/// Neural last. +fn walk(rows: &[Outcome; 4]) -> Option<(bool, u64)> { + let mut spent = 0u64; + for (i, row) in rows.iter().enumerate() { + let (ok, uj) = (*row)?; + spent = spent.saturating_add(uj); + if ok || i == rows.len() - 1 { + return Some((ok, spent)); + } + } + unreachable!("the walk ends at Neural") +} + +/// Evaluate a router on a labelled test set of per-stage outcomes. /// -/// Each example is treated as a (embedding → routed_stage) prediction. If the -/// example's `success` is true and the router picks the same stage, we count -/// it correct. If the example was a failure, we ignore correctness but still -/// account for joules. +/// Rows with the same embedding are one query, and each row is one stage's +/// outcome on it; when a stage has more than one row for a query, the first +/// counts. See the module docs for how each recommendation is priced. pub async fn evaluate(router: &R, test_set: &[Example]) -> RouterMetrics { + let mut index: HashMap, usize> = HashMap::new(); + let mut queries: Vec<(&[f32], [Outcome; 4])> = Vec::new(); + for ex in test_set { + let i = *index.entry(query_key(&ex.embedding)).or_insert_with(|| { + queries.push((&ex.embedding, [None; 4])); + queries.len() - 1 + }); + let slot = &mut queries[i].1[stage_index(ex.stage)]; + if slot.is_none() { + *slot = Some((ex.success, ex.cost_microjoules)); + } + } + + let mut n = 0usize; + let mut unscored = 0usize; let mut correct = 0usize; let mut total_microjoules: u64 = 0; let mut full_cascade_microjoules: u64 = 0; let mut stage_counts: HashMap = HashMap::new(); - for ex in test_set { - let q = Query::new("").with_embedding(ex.embedding.clone()); - let prediction = router.route(&q).await; - *stage_counts.entry(prediction.recommended).or_insert(0) += 1; - if ex.success && prediction.recommended == ex.stage { - correct += 1; - } - // Saved joules: only pay for the recommended stage, not all four. - total_microjoules = total_microjoules.saturating_add(ex.cost_microjoules); - // "Full-cascade" baseline pays every stage's average cost — we use - // 4x the recorded cost as a conservative proxy when per-stage costs - // aren't broken out in the test set. - full_cascade_microjoules = full_cascade_microjoules.saturating_add(ex.cost_microjoules * 4); + for (embedding, rows) in &queries { + let q = Query::new("").with_embedding(embedding.to_vec()); + let recommended = router.route(&q).await.recommended; + *stage_counts.entry(recommended).or_insert(0) += 1; + let (Some((ok, uj)), Some((_, walk_uj))) = (dispatch(recommended, rows), walk(rows)) else { + unscored += 1; + continue; + }; + n += 1; + correct += usize::from(ok); + total_microjoules = total_microjoules.saturating_add(uj); + full_cascade_microjoules = full_cascade_microjoules.saturating_add(walk_uj); } - let n = test_set.len(); let accuracy = if n == 0 { 0.0 } else { correct as f32 / n as f32 }; - let joules_per_correct = if correct == 0 { - 0 - } else { - total_microjoules / correct as u64 - }; + let joules_per_correct = (correct > 0).then(|| total_microjoules / correct as u64); let lift_microjoules = full_cascade_microjoules as i64 - total_microjoules as i64; RouterMetrics { n, + unscored, correct, accuracy, total_microjoules, joules_per_correct, stage_counts, + full_cascade_microjoules, lift_microjoules, } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::router::{RouterState, StagePrediction}; + use async_trait::async_trait; + + /// Recommends one stage for every query. + struct Always(Stage); + + #[async_trait] + impl LearnedRouter for Always { + async fn route(&self, _q: &Query) -> StagePrediction { + StagePrediction::new(self.0, 1.0) + } + async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {} + fn export_state(&self) -> RouterState { + unimplemented!() + } + fn import_state(_s: RouterState) -> crate::Result { + unimplemented!() + } + } + + fn row(q: f32, stage: Stage, success: bool, uj: u64) -> Example { + Example::new(vec![q], stage, success, uj) + } + + /// Query 1 is resolved first by Graph. Query 2 has rows only for Cache + /// (a hit) and Neural. + fn set() -> Vec { + vec![ + row(1.0, Stage::Cache, false, 1), + row(1.0, Stage::Kv, false, 10), + row(1.0, Stage::Graph, true, 100), + row(1.0, Stage::Neural, true, 1_000), + row(2.0, Stage::Cache, true, 1), + row(2.0, Stage::Neural, true, 1_000), + ] + } + + #[tokio::test] + async fn the_routers_choice_decides_the_joules() { + let m = evaluate(&Always(Stage::Graph), &set()).await; + // Query 1: Graph resolves it for 100; the walk spends 1 + 10 + 100. + // Query 2: no Graph row, so it is not scored. + assert_eq!((m.n, m.unscored, m.correct), (1, 1, 1)); + assert_eq!( + (m.total_microjoules, m.full_cascade_microjoules), + (100, 111) + ); + assert_eq!(m.lift_microjoules, 11); + assert_eq!(m.joules_per_correct, Some(100)); + assert_eq!(m.stage_counts[&Stage::Graph], 2); + + let m = evaluate(&Always(Stage::Kv), &set()).await; + // Query 1: Kv misses for 10 and Neural answers for 1 000. + assert_eq!((m.n, m.total_microjoules), (1, 1_010)); + assert_eq!(m.lift_microjoules, 111 - 1_010); + + let m = evaluate(&Always(Stage::Cache), &set()).await; + // Both walks: 111, and 1 for query 2's cache hit. + assert_eq!((m.n, m.total_microjoules, m.lift_microjoules), (2, 112, 0)); + + let m = evaluate(&Always(Stage::Neural), &set()).await; + assert_eq!((m.n, m.total_microjoules), (2, 2_000)); + } + + /// The old metric charged each row's own cost whatever the router chose, + /// and priced the walk at four times it, so every router showed a lift of + /// three times its spend. One row per query prices no walk, and no + /// dispatch except to the stage the row is for. + #[tokio::test] + async fn one_row_per_query_scores_only_what_it_can_price() { + let rows = vec![ + row(1.0, Stage::Kv, true, 50_000), + row(2.0, Stage::Neural, true, 50_000_000), + ]; + let m = evaluate(&Always(Stage::Kv), &rows).await; + assert_eq!((m.n, m.unscored), (0, 2)); + assert_eq!( + ( + m.total_microjoules, + m.full_cascade_microjoules, + m.lift_microjoules + ), + (0, 0, 0) + ); + assert_eq!(m.joules_per_correct, None); + } + + #[tokio::test] + async fn no_correct_answer_has_no_joules_per_correct() { + let rows = vec![ + row(1.0, Stage::Cache, false, 1), + row(1.0, Stage::Kv, false, 10), + row(1.0, Stage::Graph, false, 100), + row(1.0, Stage::Neural, false, 1_000), + ]; + let m = evaluate(&Always(Stage::Neural), &rows).await; + assert_eq!((m.n, m.correct, m.total_microjoules), (1, 0, 1_000)); + assert_eq!( + m.joules_per_correct, None, + "zero would read as the best possible figure" + ); + } + + #[tokio::test] + async fn the_first_row_for_a_stage_counts() { + let mut rows = set(); + rows.push(row(1.0, Stage::Graph, false, 7)); + let m = evaluate(&Always(Stage::Graph), &rows).await; + assert_eq!((m.correct, m.total_microjoules), (1, 100)); + } +} diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs index 71c3c6b2e..7ca105129 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/lib.rs @@ -4,8 +4,11 @@ //! strict order. That's optimal when most queries land in the cheap stages. //! For workloads where many queries unavoidably need the neural stage, walking //! the cheaper stages first is wasted joules. A *learned router* predicts the -//! cheapest stage that is likely to satisfy the query and lets us **skip -//! ahead** when the prediction is confident. +//! stage the cascade answers from most cheaply in expectation +//! ([`router::expected_joules`]) and lets us **skip ahead** when the +//! prediction is confident. Cheapest in expectation is not likeliest to +//! succeed: Neural succeeds on everything and costs the most, and a router +//! that ranked stages by success alone learned to send every query there. //! //! This crate ships three orthogonal router families plus a threshold policy //! and an end-to-end cascade extension: diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs index 3667a5ddf..a673a842e 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs @@ -30,24 +30,14 @@ use rand_distr::{Distribution, Normal}; use serde::{Deserialize, Serialize}; use crate::error::{Error, Result}; -use crate::router::{ALL_STAGES, LearnedRouter, RouterState, StagePrediction}; +use crate::router::{ + LearnedRouter, RouterState, StagePrediction, blend_joules, cheapest_in_expectation, + default_stage_joules, stage_index, +}; /// Default latent dimension. pub const DEFAULT_LATENT_DIM: usize = 32; -fn stage_index(s: Stage) -> usize { - match s { - Stage::Cache => 0, - Stage::Kv => 1, - Stage::Graph => 2, - Stage::Neural => 3, - } -} - -fn stage_from_index(i: usize) -> Stage { - ALL_STAGES[i.min(3)] -} - /// Matrix-factorisation router (embedding-conditioned). #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MfRouter { @@ -91,7 +81,7 @@ impl MfRouter { stage_factors, stage_bias, global_bias: 0.0, - joule_ema: default_joule_ema(), + joule_ema: default_stage_joules(), } } @@ -105,11 +95,11 @@ impl MfRouter { self.embedding_dim } - /// Forward pass: predict success probability per stage. + /// Forward pass: the probability that each stage resolves the query. /// - /// Returns logits *and* softmax-normalised probabilities. The softmax is - /// taken over per-stage sigmoid outputs so the recommended-stage - /// probability stays in [0, 1] when the router is uncertain. + /// One sigmoid per stage, not a softmax: the stages are not alternatives + /// to one another (a query Kv resolves, Neural resolves too), so the + /// probabilities need not sum to one. pub fn predict_probs(&self, embedding: &[f32]) -> Result<[f32; 4]> { if embedding.len() != self.embedding_dim { return Err(Error::DimensionMismatch { @@ -129,11 +119,13 @@ impl MfRouter { Ok(probs) } - /// Pick the recommended stage and its confidence. + /// Pick the stage cheapest to answer from in expectation + /// ([`cheapest_in_expectation`], priced at [`Self::joule_estimates`]). + /// The confidence is that stage's probability of resolving the query. pub fn predict(&self, embedding: &[f32]) -> Result<(Stage, f32, [f32; 4])> { let probs = self.predict_probs(embedding)?; - let (idx, conf) = argmax(&probs); - Ok((stage_from_index(idx), conf, probs)) + let stage = cheapest_in_expectation(&probs, &self.joule_ema); + Ok((stage, probs[stage_index(stage)], probs)) } /// One SGD step on a single example. @@ -196,9 +188,7 @@ impl MfRouter { /// Record an observed joule cost — EMA with coefficient 0.1. pub fn observe_cost(&mut self, stage: Stage, microjoules: u64) { - let prev = self.joule_ema.get(&stage).copied().unwrap_or(microjoules) as f64; - let blended = 0.9 * prev + 0.1 * microjoules as f64; - self.joule_ema.insert(stage, blended as u64); + blend_joules(&mut self.joule_ema, stage, microjoules); } /// Current per-stage joule estimate (EMA). @@ -207,16 +197,6 @@ impl MfRouter { } } -fn default_joule_ema() -> HashMap { - let mut m = HashMap::new(); - // Conservative defaults, ordered cheap → expensive. - m.insert(Stage::Cache, 1_000); - m.insert(Stage::Kv, 50_000); - m.insert(Stage::Graph, 500_000); - m.insert(Stage::Neural, 50_000_000); - m -} - fn sigmoid(x: f32) -> f32 { if x >= 0.0 { let z = (-x).exp(); @@ -227,18 +207,6 @@ fn sigmoid(x: f32) -> f32 { } } -fn argmax(xs: &[f32; 4]) -> (usize, f32) { - let mut best = 0; - let mut bv = xs[0]; - for (i, v) in xs.iter().enumerate().skip(1) { - if *v > bv { - best = i; - bv = *v; - } - } - (best, bv) -} - #[async_trait] impl LearnedRouter for MfRouter { async fn route(&self, query: &Query) -> StagePrediction { diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs index 9ea875328..b856eccb2 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/router.rs @@ -63,10 +63,20 @@ pub struct RouterState { /// A learned router predicts the best cascade stage from a query. #[async_trait] pub trait LearnedRouter: Send + Sync { - /// Predict the best stage for `query`. + /// Predict the best stage for `query`: the one [`LearnedCascade`] answers + /// from most cheaply in expectation (see [`expected_joules`]), not the one + /// most likely to succeed. The neural stage is the most likely to succeed + /// and the most expensive, so a router ranking by success alone learns to + /// skip every cheap stage. + /// + /// [`LearnedCascade`]: crate::inference::LearnedCascade async fn route(&self, query: &Query) -> StagePrediction; - /// Online update — observe what actually happened. + /// Online update with the outcome of one stage that ran: whether + /// `chosen_stage` resolved `query` correctly, and the microjoules that + /// stage itself spent. When a dispatch ran several stages (a full walk, + /// or a miss that fell through to Neural), report each of them: a stage + /// that missed is evidence too. async fn observe( &mut self, query: &Query, @@ -86,3 +96,186 @@ pub trait LearnedRouter: Send + Sync { /// The four cascade stages, in cheapest-first order, useful for iteration. pub const ALL_STAGES: [Stage; 4] = [Stage::Cache, Stage::Kv, Stage::Graph, Stage::Neural]; + +/// A stage's position in [`ALL_STAGES`]. +pub(crate) fn stage_index(s: Stage) -> usize { + match s { + Stage::Cache => 0, + Stage::Kv => 1, + Stage::Graph => 2, + Stage::Neural => 3, + } +} + +/// Per-stage cost assumed before any is observed, in microjoules. Conservative +/// defaults, ordered cheap → expensive. +pub fn default_stage_joules() -> HashMap { + ALL_STAGES.iter().map(|s| (*s, default_cost(*s))).collect() +} + +fn default_cost(stage: Stage) -> u64 { + match stage { + Stage::Cache => 1_000, + Stage::Kv => 50_000, + Stage::Graph => 500_000, + Stage::Neural => 50_000_000, + } +} + +/// Fold one observed stage cost into a per-stage estimate — EMA with +/// coefficient 0.1. +pub fn blend_joules(joules: &mut HashMap, stage: Stage, microjoules: u64) { + let prev = joules.get(&stage).copied().unwrap_or(microjoules) as f64; + let blended = 0.9 * prev + 0.1 * microjoules as f64; + joules.insert(stage, blended as u64); +} + +/// Expected microjoules for [`LearnedCascade`] to answer a query it was told +/// to skip to `stage`, given each stage's probability of resolving the query +/// (`resolves`, in [`ALL_STAGES`] order) and each stage's cost (`joules`). +/// +/// It follows the dispatch exactly: +/// +/// ```text +/// Neural c_N +/// Kv, Graph c_s + (1 − p_s) · c_N a miss falls through to Neural +/// Cache c_C + (1 − p_C)(c_K + (1 − p_K)(c_G + (1 − p_G) · c_N)) +/// the full walk +/// ``` +/// +/// Neural always answers, so its own probability does not enter. The walk +/// treats the stages' misses as independent, because the routers estimate +/// each stage's chance on its own rather than given that a cheaper one +/// missed. A probability outside [0, 1] is clamped and NaN counts as a miss, +/// so a broken estimate prices a stage dearer, never cheaper. A stage with no +/// cost estimate is priced at [`default_stage_joules`]. +/// +/// [`LearnedCascade`]: crate::inference::LearnedCascade +pub fn expected_joules(stage: Stage, resolves: &[f32; 4], joules: &HashMap) -> f64 { + let cost = |s: Stage| joules.get(&s).copied().unwrap_or_else(|| default_cost(s)) as f64; + let miss = |i: usize| { + let p = resolves[i]; + if p.is_nan() { + 1.0 + } else { + 1.0 - f64::from(p.clamp(0.0, 1.0)) + } + }; + let neural = cost(Stage::Neural); + match stage { + Stage::Neural => neural, + Stage::Kv => cost(Stage::Kv) + miss(1) * neural, + Stage::Graph => cost(Stage::Graph) + miss(2) * neural, + Stage::Cache => { + let graph_on = cost(Stage::Graph) + miss(2) * neural; + let kv_on = cost(Stage::Kv) + miss(1) * graph_on; + cost(Stage::Cache) + miss(0) * kv_on + } + } +} + +/// The stage [`LearnedCascade`] answers from most cheaply in expectation +/// ([`expected_joules`]). A tie goes to the cheaper stage. +/// +/// [`LearnedCascade`]: crate::inference::LearnedCascade +pub fn cheapest_in_expectation(resolves: &[f32; 4], joules: &HashMap) -> Stage { + let mut best = ALL_STAGES[0]; + let mut best_uj = expected_joules(best, resolves, joules); + for s in &ALL_STAGES[1..] { + let uj = expected_joules(*s, resolves, joules); + if uj < best_uj { + best = *s; + best_uj = uj; + } + } + best +} + +#[cfg(test)] +mod tests { + use super::*; + + fn costs() -> HashMap { + let mut m = HashMap::new(); + m.insert(Stage::Cache, 1); + m.insert(Stage::Kv, 10); + m.insert(Stage::Graph, 100); + m.insert(Stage::Neural, 1_000); + m + } + + #[test] + fn each_dispatch_is_priced_as_the_cascade_runs_it() { + let p = [0.5, 0.25, 0.75, 0.0]; + let c = costs(); + assert_eq!(expected_joules(Stage::Neural, &p, &c), 1_000.0); + assert_eq!(expected_joules(Stage::Kv, &p, &c), 10.0 + 0.75 * 1_000.0); + assert_eq!( + expected_joules(Stage::Graph, &p, &c), + 100.0 + 0.25 * 1_000.0 + ); + let walk = 1.0 + 0.5 * (10.0 + 0.75 * (100.0 + 0.25 * 1_000.0)); + assert_eq!(expected_joules(Stage::Cache, &p, &c), walk); + } + + /// The defect this rule replaces: Neural is the likeliest to succeed, so + /// an argmax over success picked it, at a thousand times Graph's cost. + #[test] + fn the_likeliest_stage_is_not_the_cheapest_to_answer_from() { + let c = costs(); + let p = [0.0, 0.0, 0.95, 1.0]; + assert_eq!(cheapest_in_expectation(&p, &c), Stage::Graph); + // Hopeless cheap stages: skipping them all is cheapest. + assert_eq!( + cheapest_in_expectation(&[0.0, 0.0, 0.0, 1.0], &c), + Stage::Neural + ); + // A likely cache hit: walk from the top. + assert_eq!( + cheapest_in_expectation(&[0.9, 0.0, 0.0, 1.0], &c), + Stage::Cache + ); + } + + #[test] + fn a_broken_estimate_never_prices_a_stage_cheaper() { + let c = costs(); + let nan = [f32::NAN; 4]; + assert_eq!(expected_joules(Stage::Kv, &nan, &c), 1_010.0); + assert_eq!(expected_joules(Stage::Kv, &[0.0, 7.0, 0.0, 0.0], &c), 10.0); + assert_eq!( + expected_joules(Stage::Kv, &[0.0, -3.0, 0.0, 0.0], &c), + 1_010.0 + ); + } + + #[test] + fn a_tie_goes_to_the_cheaper_stage() { + let mut c = costs(); + c.insert(Stage::Graph, 10); + // Kv and Graph both cost 10 and both certainly resolve. + assert_eq!( + cheapest_in_expectation(&[0.0, 1.0, 1.0, 1.0], &c), + Stage::Kv + ); + } + + #[test] + fn a_stage_without_an_estimate_is_priced_at_the_default() { + let p = [0.0; 4]; + assert_eq!( + expected_joules(Stage::Neural, &p, &HashMap::new()), + 50_000_000.0 + ); + } + + #[test] + fn blending_moves_a_tenth_of_the_way() { + let mut j = costs(); + blend_joules(&mut j, Stage::Neural, 2_000); + assert_eq!(j[&Stage::Neural], 1_100); + let mut empty = HashMap::new(); + blend_joules(&mut empty, Stage::Kv, 40); + assert_eq!(empty[&Stage::Kv], 40); + } +} diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs index c0d8dfd46..3ac727f28 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/training.rs @@ -5,6 +5,8 @@ //! `train_logreg` (logistic regression). Both consume `&[Example]` and emit a //! ready-to-serve router. +use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::path::Path; use eoc_core::Stage; @@ -15,6 +17,7 @@ use serde::{Deserialize, Serialize}; use crate::classifier::LogRegRouter; use crate::error::{Error, Result}; use crate::matrix_factorization::{DEFAULT_LATENT_DIM, MfRouter}; +use crate::router::stage_index; /// One labeled training example. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -141,16 +144,46 @@ pub fn train_mf_with_dim( router } -/// Train a logistic-regression router. Only successful examples shape the -/// classifier — failed-stage examples are skipped because we want -/// `P(stage = correct | x)`. +/// The key that makes rows one query: their embedding, bit for bit. +pub(crate) fn query_key(embedding: &[f32]) -> Vec { + embedding.iter().map(|v| v.to_bits()).collect() +} + +/// Each query's cheapest successful row, in the order queries first appear. +/// Rows with the same embedding are one query; a query no stage resolved has +/// no row here. +pub(crate) fn cheapest_successes(examples: &[Example]) -> Vec<&Example> { + let mut index: HashMap, usize> = HashMap::new(); + let mut out: Vec<&Example> = Vec::new(); + for e in examples.iter().filter(|e| e.success) { + match index.entry(query_key(&e.embedding)) { + Entry::Occupied(slot) => { + let kept = &mut out[*slot.get()]; + if stage_index(e.stage) < stage_index(kept.stage) { + *kept = e; + } + } + Entry::Vacant(slot) => { + slot.insert(out.len()); + out.push(e); + } + } + } + out +} + +/// Train a logistic-regression router to predict the cheapest stage that +/// resolves a query. Each query is labelled with its cheapest successful row +/// ([`Example`]s with the same embedding are one query); failed rows label +/// nothing. Labelling every successful row instead would teach it the stage +/// that most often succeeds, and Neural succeeds on everything. pub fn train_logreg(examples: &[Example], config: TrainingConfig) -> LogRegRouter { let dim = first_embedding_dim(examples).unwrap_or(1); let mut router = LogRegRouter::new(dim, config.l2, config.seed); if examples.is_empty() { return router; } - let positive: Vec<&Example> = examples.iter().filter(|e| e.success).collect(); + let positive = cheapest_successes(examples); if positive.is_empty() { return router; } @@ -182,4 +215,50 @@ mod tests { let r = train_logreg(&[], TrainingConfig::default()); assert_eq!(r.embedding_dim(), 1); } + + #[test] + fn a_query_is_labelled_with_its_cheapest_success() { + let a = vec![1.0, 0.0]; + let b = vec![0.0, 1.0]; + let rows = vec![ + Example::new(a.clone(), Stage::Neural, true, 50_000_000), + Example::new(a.clone(), Stage::Cache, false, 1_000), + Example::new(b.clone(), Stage::Graph, false, 500_000), + Example::new(a.clone(), Stage::Graph, true, 500_000), + Example::new(a.clone(), Stage::Kv, true, 50_000), + Example::new(b.clone(), Stage::Neural, true, 50_000_000), + Example::new(vec![9.0, 9.0], Stage::Kv, false, 50_000), + ]; + let kept: Vec<(Vec, Stage)> = cheapest_successes(&rows) + .into_iter() + .map(|e| (e.embedding.clone(), e.stage)) + .collect(); + assert_eq!(kept, vec![(a, Stage::Kv), (b, Stage::Neural)]); + } + + /// Every query here is resolved by Kv, its cheapest stage, and by Graph + /// and Neural; Neural ran twice. Trained on every successful row, the + /// classifier learned Neural, the stage that most often succeeds. + #[test] + fn logreg_learns_the_cheapest_stage_not_the_most_frequent_success() { + let mut rows = Vec::new(); + for i in 0..40 { + let emb = vec![1.0, i as f32 / 40.0]; + rows.push(Example::new(emb.clone(), Stage::Kv, true, 50_000)); + rows.push(Example::new(emb.clone(), Stage::Graph, true, 500_000)); + rows.push(Example::new(emb.clone(), Stage::Neural, true, 50_000_000)); + rows.push(Example::new(emb, Stage::Neural, true, 50_000_000)); + } + let cfg = TrainingConfig { + epochs: 30, + learning_rate: 0.1, + l2: 0.0, + ..TrainingConfig::default() + }; + let r = train_logreg(&rows, cfg); + for i in 0..40 { + let emb = [1.0, i as f32 / 40.0]; + assert_eq!(r.predict(&emb).unwrap().0, Stage::Kv); + } + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs index 76b5a94d3..5118306c3 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/linucb_bandit.rs @@ -1,68 +1,193 @@ -//! 100-round LinUCB synthetic bandit test. +//! Synthetic bandit tests: LinUCB and Thompson sampling. //! -//! Four arms with known mean rewards. Context is a fixed-length embedding -//! that *informs* but doesn't determine reward (Gaussian noise). After ~50 -//! rounds the router should pull the best arm more than any other. +//! Each stage resolves a query with a known probability. Each round the +//! router picks a stage, the pick is run the way `LearnedCascade` dispatches +//! it (Kv or Graph falling through to Neural on a miss, Neural alone, or for +//! Cache the full walk), and every stage that ran reports its outcome and +//! cost. The router must learn the dispatch that spends least, not the stage +//! likeliest to succeed: before, both bandits converged on Neural. use eoc_core::Stage; -use eoc_route_learned::bandit::LinUcbRouter; -use rand::SeedableRng; +use eoc_route_learned::bandit::{LinUcbRouter, ThompsonSamplingRouter}; +use eoc_route_learned::router::ALL_STAGES; +use rand::{RngExt, SeedableRng}; use rand_distr::{Distribution, Normal}; use std::collections::HashMap; const DIM: usize = 4; -fn arm_mean(stage: Stage) -> f32 { +fn cost(stage: Stage) -> u64 { match stage { - Stage::Cache => 0.10, - Stage::Kv => 0.20, - Stage::Graph => 0.30, - // Best arm: Neural. - Stage::Neural => 0.80, + Stage::Cache => 1_000, + Stage::Kv => 50_000, + Stage::Graph => 500_000, + Stage::Neural => 50_000_000, } } -#[test] -fn linucb_converges_to_best_arm() { - // Lower alpha → less exploration, more exploitation. Ridge=0.1 so the - // estimated theta moves quickly. - let mut router = LinUcbRouter::new(DIM, 0.1, 0.1); - let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF); - let ctx_dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); - let reward_noise = Normal::new(0.0_f32, 0.02).expect("valid normal"); - - let mut pulls: HashMap = HashMap::new(); - let mut tail_pulls: HashMap = HashMap::new(); - - for round in 0..100 { - let ctx: Vec = (0..DIM).map(|_| ctx_dist.sample(&mut rng)).collect(); - let (pick, _score) = router.pick(&ctx).expect("pick"); - *pulls.entry(pick).or_insert(0) += 1; - if round >= 50 { - *tail_pulls.entry(pick).or_insert(0) += 1; +/// The stages a dispatch to `stage` runs, each with whether it resolved. +fn dispatch(stage: Stage, resolves: &HashMap) -> Vec<(Stage, bool)> { + let order: &[Stage] = match stage { + Stage::Cache => &ALL_STAGES, + Stage::Kv => &[Stage::Kv, Stage::Neural], + Stage::Graph => &[Stage::Graph, Stage::Neural], + Stage::Neural => &[Stage::Neural], + }; + let mut ran = Vec::new(); + for s in order { + ran.push((*s, resolves[s])); + if resolves[s] { + break; } - let reward = (arm_mean(pick) + reward_noise.sample(&mut rng)).clamp(0.0, 1.0); - router.update(&ctx, pick, reward).expect("update"); } + ran +} - println!("LinUCB pulls (all 100 rounds): {pulls:?}"); - println!("LinUCB pulls (rounds 50-99): {tail_pulls:?}"); - - let neural_tail = *tail_pulls.get(&Stage::Neural).unwrap_or(&0); - let best_other = [Stage::Cache, Stage::Kv, Stage::Graph] - .iter() - .map(|s| *tail_pulls.get(s).unwrap_or(&0)) - .max() - .unwrap_or(0); - assert!( - neural_tail >= best_other, - "neural ({neural_tail}) should be pulled at least as often as any other ({best_other}) after round 50" - ); - // And neural should dominate over the full 100 rounds, since its mean - // reward is 0.8 vs ≤0.3 for every other arm. - let neural_total = *pulls.get(&Stage::Neural).unwrap_or(&0); - assert!( - neural_total >= 40, - "neural should be pulled ≥40 times over 100 rounds, got {neural_total}" - ); +/// Pick counts over the last `tail` of `rounds`, and the joules spent there. +struct Run { + tail_picks: HashMap, + tail_microjoules: u64, +} + +fn run( + rounds: usize, + tail: usize, + chance: [f64; 4], + mut pick: impl FnMut(&[f32]) -> Stage, + mut learn: impl FnMut(&[f32], Stage, bool, u64), +) -> Run { + let mut rng = rand::rngs::StdRng::seed_from_u64(0xBEEF); + let noise = Normal::new(0.0_f32, 1.0).expect("valid normal"); + let mut out = Run { + tail_picks: HashMap::new(), + tail_microjoules: 0, + }; + for round in 0..rounds { + // A constant feature first, so LinUCB can learn each stage's base rate. + let mut ctx = vec![1.0_f32]; + ctx.extend((1..DIM).map(|_| 0.1 * noise.sample(&mut rng))); + let resolves: HashMap = ALL_STAGES + .iter() + .zip(chance) + .map(|(s, p)| (*s, rng.random_bool(p))) + .collect(); + let stage = pick(&ctx); + for (s, ok) in dispatch(stage, &resolves) { + learn(&ctx, s, ok, cost(s)); + if round >= rounds - tail { + out.tail_microjoules += cost(s); + } + } + if round >= rounds - tail { + *out.tail_picks.entry(stage).or_insert(0) += 1; + } + } + out +} + +fn linucb(rounds: usize, tail: usize, chance: [f64; 4]) -> Run { + let router = std::cell::RefCell::new(LinUcbRouter::new(DIM, 0.1, 1.0)); + run( + rounds, + tail, + chance, + |ctx| router.borrow().pick(ctx).expect("pick").0, + |ctx, s, ok, uj| { + let mut r = router.borrow_mut(); + r.update(ctx, s, if ok { 1.0 } else { 0.0 }) + .expect("update"); + r.observe_cost(s, uj); + }, + ) +} + +fn thompson(rounds: usize, tail: usize, chance: [f64; 4]) -> Run { + let router = std::cell::RefCell::new(ThompsonSamplingRouter::new(7)); + run( + rounds, + tail, + chance, + |_| router.borrow().pick().0, + |_, s, ok, uj| { + let mut r = router.borrow_mut(); + r.update(s, ok); + r.observe_cost(s, uj); + }, + ) +} + +fn share(run: &Run, stage: Stage, tail: usize) -> f32 { + *run.tail_picks.get(&stage).unwrap_or(&0) as f32 / tail as f32 +} + +/// Graph resolves 95% of queries and nothing cheaper resolves any: skip to +/// Graph, at about 3 million µJ a query, not Neural at 50 million. +#[test] +fn bandits_skip_to_the_stage_that_resolves() { + let chance = [0.0, 0.0, 0.95, 1.0]; + for (name, r) in [ + ("linucb", linucb(3000, 500, chance)), + ("thompson", thompson(3000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + assert!( + share(&r, Stage::Graph, 500) > 0.8, + "{name} picks Graph only {}", + share(&r, Stage::Graph, 500) + ); + assert!( + share(&r, Stage::Neural, 500) < 0.05, + "{name} still skips to Neural" + ); + assert!( + r.tail_microjoules / 500 < 5_000_000, + "{name} spends {} µJ a query", + r.tail_microjoules / 500 + ); + } +} + +/// Nothing short of Neural resolves anything, so going straight to Neural, +/// at 50 million µJ, is cheapest; the full walk costs 50.551 million. The +/// bandits come within 1% of Neural. They still probe Kv now and then: a +/// probe costs a thousandth of a Neural call and would save the whole call +/// if it hit, so it is worth trying until the bound on Kv's chance falls +/// below 0.1%. +#[test] +fn bandits_come_within_a_percent_of_neural_when_nothing_cheaper_resolves() { + let chance = [0.0, 0.0, 0.0, 1.0]; + for (name, r) in [ + ("linucb", linucb(4000, 500, chance)), + ("thompson", thompson(4000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + let per_query = r.tail_microjoules / 500; + assert!( + per_query < 50_500_000, + "{name} spends {per_query} µJ a query" + ); + } +} + +/// The old test's arms. Neural is likeliest to succeed (0.8), and the old +/// bandits learned to pick it every time. But the cheap stages resolve often +/// enough that walking them first costs least, so the walk is the pick. +#[test] +fn bandits_walk_when_cheap_stages_often_resolve() { + let chance = [0.1, 0.2, 0.3, 0.8]; + for (name, r) in [ + ("linucb", linucb(1000, 500, chance)), + ("thompson", thompson(1000, 500, chance)), + ] { + println!("{name}: {:?}, {} µJ", r.tail_picks, r.tail_microjoules); + assert!( + share(&r, Stage::Cache, 500) > 0.8, + "{name} walks only {}", + share(&r, Stage::Cache, 500) + ); + assert!( + share(&r, Stage::Neural, 500) < 0.05, + "{name} skips to Neural" + ); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs index 1338e996e..c6a9b3efb 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/mf_synthetic.rs @@ -1,73 +1,204 @@ -//! Synthetic-data accuracy test for the MF router. +//! Synthetic test for the MF router. //! -//! Generates 1000 examples where the "best stage" is a deterministic -//! function of the embedding (first two dims define quadrants → stage). The -//! MF router should learn this assignment with >80% accuracy. +//! Each query has one row per stage. Each cheap stage has its own coverage: +//! the chance it resolves a query is a logistic function of one feature of +//! the embedding (Cache's centred at 0.5, Kv's at 0, Graph's at −0.5), and +//! Neural resolves everything. The router must learn to skip to the stage +//! that is cheapest to answer from in expectation, not to the one likeliest +//! to succeed, which is Neural. +//! +//! The coverage is soft, as real coverage is, and of the family the router +//! models, so its learned chances can be right. They have to be: with the +//! default costs a wrong skip pays a Neural call, a thousand times Graph, +//! while a right one saves at most a hundredth of one. A router sure of +//! itself on noise-free boundaries skips wrongly at their edges and spends +//! more than the plain walk. +//! +//! With the default costs the plain walk is itself within a fraction of a +//! percent of the best choices, so there is little for routing to win, and +//! the router lands within 1% of the best. Where Kv and Graph cost a real +//! share of a Neural call, the walk is 4–5% above the best, but at this +//! training budget MF's chances are off by 5–9% a stage, and the router lands +//! at about the walk's cost rather than capturing that. Training longer made +//! it beat the walk by 2% while its chances got worse, so that is a property +//! of the estimator, not of the rule; the test holds it to not losing. -use eoc_core::Stage; +use std::collections::HashMap; + +use async_trait::async_trait; +use eoc_core::{Query, Stage}; +use eoc_route_learned::evaluation::evaluate; use eoc_route_learned::matrix_factorization::MfRouter; +use eoc_route_learned::router::{ + ALL_STAGES, LearnedRouter, RouterState, StagePrediction, cheapest_in_expectation, + default_stage_joules, expected_joules, +}; use eoc_route_learned::training::{Example, TrainingConfig, train_mf}; -use rand::SeedableRng; +use rand::{RngExt, SeedableRng}; use rand_distr::{Distribution, Normal}; const DIM: usize = 8; -fn quadrant_stage(emb: &[f32]) -> Stage { - let (a, b) = (emb[0], emb[1]); - match (a >= 0.0, b >= 0.0) { - (true, true) => Stage::Cache, - (true, false) => Stage::Kv, - (false, true) => Stage::Graph, - (false, false) => Stage::Neural, - } +/// Kv and Graph at a real share of a Neural call. +fn dear_cheap_stages() -> HashMap { + let mut c = default_stage_joules(); + c.insert(Stage::Kv, 2_000_000); + c.insert(Stage::Graph, 10_000_000); + c } -fn synthesise(n: usize, seed: u64) -> Vec { +/// The true chance that each stage resolves the query with this embedding, +/// in `ALL_STAGES` order. +fn chances(emb: &[f32]) -> [f32; 4] { + let s = |x: f32| 1.0 / (1.0 + (-4.0 * x).exp()); + [s(emb[0] - 0.5), s(emb[1]), s(emb[2] + 0.5), 1.0] +} + +fn synthesise(n: usize, seed: u64, costs: &HashMap) -> Vec { let mut rng = rand::rngs::StdRng::seed_from_u64(seed); let dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); - let mut out = Vec::with_capacity(n); + let mut out = Vec::with_capacity(n * 4); for _ in 0..n { - let mut emb = vec![0.0_f32; DIM]; - for v in emb.iter_mut() { - *v = dist.sample(&mut rng); + let emb: Vec = (0..DIM).map(|_| dist.sample(&mut rng)).collect(); + let p = chances(&emb); + for (i, stage) in ALL_STAGES.into_iter().enumerate() { + let success = rng.random_bool(f64::from(p[i])); + out.push(Example::new(emb.clone(), stage, success, costs[&stage])); } - let stage = quadrant_stage(&emb); - // Pure labels — the synthetic test focuses on classification accuracy, - // not noise-handling. - let success = true; - // cost varies per stage so cost estimates are meaningful. - let cost = match stage { - Stage::Cache => 1_000, - Stage::Kv => 50_000, - Stage::Graph => 500_000, - Stage::Neural => 50_000_000, - }; - out.push(Example::new(emb, stage, success, cost)); } out } -#[test] -fn mf_router_accuracy_above_80() { - let train = synthesise(1000, 17); - let test = synthesise(200, 42); - +fn router(costs: &HashMap) -> MfRouter { let cfg = TrainingConfig { epochs: 30, learning_rate: 0.1, l2: 1e-5, ..TrainingConfig::default() }; - let router: MfRouter = train_mf(&train, cfg); + train_mf(&synthesise(1000, 17, costs), cfg) +} - let mut correct = 0usize; - for ex in &test { - let (stage, _conf, _probs) = router.predict(&ex.embedding).expect("predict"); - if stage == ex.stage { - correct += 1; +/// The rule the routers used before: the stage likeliest to succeed. +struct LikeliestToSucceed<'a>(&'a MfRouter); + +#[async_trait] +impl LearnedRouter for LikeliestToSucceed<'_> { + async fn route(&self, q: &Query) -> StagePrediction { + let probs = self.0.predict_probs(q.embedding.as_ref().unwrap()).unwrap(); + let best = (0..4).fold(0, |b, i| if probs[i] > probs[b] { i } else { b }); + StagePrediction::new(ALL_STAGES[best], probs[best]) + } + async fn observe(&mut self, _q: &Query, _s: Stage, _ok: bool, _j: u64) {} + fn export_state(&self) -> RouterState { + unimplemented!() + } + fn import_state(_s: RouterState) -> eoc_route_learned::Result { + unimplemented!() + } +} + +/// Expected µJ over held-out queries, priced at their true chances: the best +/// choices, the router's, the plain walk's and the likeliest-stage rule's; +/// and the share of queries where the router's choice is the best one. +struct Priced { + best: f64, + routed: f64, + walk: f64, + likeliest: f64, + agree: f64, +} + +fn priced(costs: &HashMap) -> Priced { + let router = router(costs); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let dist = Normal::new(0.0_f32, 1.0).expect("valid normal"); + let mut out = Priced { + best: 0.0, + routed: 0.0, + walk: 0.0, + likeliest: 0.0, + agree: 0.0, + }; + let mut fit = [0.0_f64; 4]; + let n = 2000; + for _ in 0..n { + let emb: Vec = (0..DIM).map(|_| dist.sample(&mut rng)).collect(); + let p = chances(&emb); + let (choice, _, learned) = router.predict(&emb).expect("predict"); + let ideal = cheapest_in_expectation(&p, costs); + let top = (0..4).fold(0, |b, i| if learned[i] > learned[b] { i } else { b }); + for i in 0..4 { + fit[i] += f64::from((learned[i] - p[i]).abs()) / f64::from(n); } + out.agree += f64::from(u8::from(choice == ideal)) / f64::from(n); + out.routed += expected_joules(choice, &p, costs); + out.walk += expected_joules(Stage::Cache, &p, costs); + out.best += expected_joules(ideal, &p, costs); + out.likeliest += expected_joules(ALL_STAGES[top], &p, costs); } - let acc = correct as f32 / test.len() as f32; - println!("MF synthetic accuracy: {:.3}", acc); - assert!(acc > 0.8, "MF accuracy {acc} below 0.8 threshold"); + println!( + "expected µJ: best {:.0}, routed {:.0}, walk {:.0}, likeliest {:.0}; agrees with best {:.3}", + out.best, out.routed, out.walk, out.likeliest, out.agree + ); + println!("mean |learned − true| chance per stage: {fit:.3?}"); + out +} + +#[test] +fn with_the_default_costs_the_router_lands_within_a_percent_of_the_best() { + let r = priced(&default_stage_joules()); + assert!( + r.agree > 0.9, + "the router's choice is the best for only {} of queries", + r.agree + ); + assert!( + r.routed <= r.best * 1.01, + "routed {:.0} µJ is over 1% above the best {:.0}", + r.routed, + r.best + ); + assert!( + r.routed * 2.0 < r.likeliest, + "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}", + r.routed, + r.likeliest + ); +} + +#[test] +fn where_cheap_stages_are_dear_the_router_does_not_lose_to_the_walk() { + let r = priced(&dear_cheap_stages()); + assert!( + r.routed <= r.walk * 1.01, + "routed {:.0} µJ is over 1% above the walk's {:.0}", + r.routed, + r.walk + ); + assert!( + r.routed * 2.0 < r.likeliest, + "routed {:.0} µJ is not far below the likeliest-stage rule's {:.0}", + r.routed, + r.likeliest + ); +} + +/// The same, replayed over sampled outcomes by `evaluate`. +#[tokio::test] +async fn mf_router_spends_far_less_than_the_likeliest_stage_rule() { + let costs = default_stage_joules(); + let router = router(&costs); + let test = synthesise(500, 42, &costs); + let routed = evaluate(&router, &test).await; + let likeliest = evaluate(&LikeliestToSucceed(&router), &test).await; + println!("routed: {routed:?}"); + println!("likeliest: {likeliest:?}"); + assert_eq!((routed.n, routed.unscored), (500, 0)); + assert!( + routed.total_microjoules * 2 < likeliest.total_microjoules, + "routed {} µJ is not far below the likeliest-stage rule's {}", + routed.total_microjoules, + likeliest.total_microjoules + ); } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs b/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs index fe0b56677..04e840d35 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/tests/state_roundtrip.rs @@ -46,3 +46,45 @@ fn rejects_wrong_algorithm_tag() { s.algorithm = "wrong".into(); assert!(MfRouter::import_state(s).is_err()); } + +/// The bandits learn each stage's cost. A learned cost survives export, and +/// a state saved before they learned costs imports with the defaults rather +/// than failing. +#[test] +fn bandit_stage_costs_roundtrip_and_older_states_import() { + use eoc_core::Stage; + use eoc_route_learned::router::default_stage_joules; + + fn without_joules(state: &mut eoc_route_learned::router::RouterState) { + let mut v: serde_json::Value = serde_json::from_slice(&state.weights_blob).unwrap(); + v.as_object_mut() + .unwrap() + .remove("joules") + .expect("the field is exported"); + state.weights_blob = serde_json::to_vec(&v).unwrap(); + } + + let mut lin = LinUcbRouter::new(4, 1.0, 1.0); + lin.observe_cost(Stage::Neural, 10_000); + let restored = LinUcbRouter::import_state(lin.export_state()).expect("import"); + assert_eq!(restored.joules, lin.joules); + let mut old = lin.export_state(); + without_joules(&mut old); + assert_eq!( + LinUcbRouter::import_state(old).expect("import").joules, + default_stage_joules() + ); + + let mut ts = ThompsonSamplingRouter::new(3); + ts.observe_cost(Stage::Kv, 7); + let restored = ThompsonSamplingRouter::import_state(ts.export_state()).expect("import"); + assert_eq!(restored.joules, ts.joules); + let mut old = ts.export_state(); + without_joules(&mut old); + assert_eq!( + ThompsonSamplingRouter::import_state(old) + .expect("import") + .joules, + default_stage_joules() + ); +} -- 2.54.0 (Apple Git-157) From 690e5c2d1d4122d7dbd320c0ebd556df3d68d4ba Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 22 Oct 2026 07:19:34 -0400 Subject: [PATCH 2/2] eoc-route-learned: pin the classifier's self-reinforcement and MF's confidence Two mutants of the previous commit survived: a correct result on the predicted stage itself no longer reinforcing it, and MF reporting a fixed confidence. Each now has a test. Co-Authored-By: Claude Opus 5.5 --- .../eoc-route-learned/src/classifier.rs | 21 +++++++++++++++++++ .../src/matrix_factorization.rs | 9 ++++++++ 2 files changed, 30 insertions(+) diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs index f1c3bfb17..8883172ce 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/classifier.rs @@ -259,4 +259,25 @@ mod tests { "a cheaper stage that resolved it is" ); } + + /// The predicted stage resolving the query is evidence for it. + #[tokio::test] + async fn the_predicted_stage_resolving_reinforces_it() { + let emb = vec![1.0, -0.5, 0.25]; + let q = Query::new("q").with_embedding(emb.clone()); + let mut r = LogRegRouter::new(3, 0.0, 9); + for _ in 0..20 { + r.sgd_step(&emb, Stage::Kv, 0.1).unwrap(); + } + let (stage, before, _) = r.predict(&emb).unwrap(); + assert_eq!(stage, Stage::Kv); + for _ in 0..50 { + r.observe(&q, Stage::Kv, true, 50_000).await; + } + let (_, after, _) = r.predict(&emb).unwrap(); + assert!( + after > before, + "confidence in Kv went from {before} to {after}" + ); + } } diff --git a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs index a673a842e..7907db673 100644 --- a/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs +++ b/eoc/eoc-rs/crates/eoc-route-learned/src/matrix_factorization.rs @@ -291,6 +291,15 @@ mod tests { } } + #[test] + fn the_confidence_is_the_chosen_stages_chance() { + let r = MfRouter::new(4, 4, 7); + let emb = [0.5, -0.5, 0.25, -0.25]; + let (stage, conf, probs) = r.predict(&emb).unwrap(); + assert_eq!(conf, probs[stage_index(stage)]); + assert!(conf < 1.0); + } + #[test] fn sgd_reduces_loss() { let mut r = MfRouter::new(4, 4, 7); -- 2.54.0 (Apple Git-157)