feat(batch16): SOAP optimizer, lookahead decoding, SWA+SWAG

- SoapOptimizer: Adam in Shampoo eigenbasis (arXiv:2409.11321); Jacobi
  eigendecomposition for L/R Kronecker factors; projection G_hat=Q_L^T@G@Q_R,
  bias-corrected Adam, unproject U=Q_L@U_hat@Q_R^T; 1D plain Adam fallback; 19 tests
- LookaheadDecoder: NGramCache (FIFO eviction, count-sorted candidates);
  draft-then-verify loop; auto-cache update on accepted tokens; LookaheadStats
  with avg_tokens_per_step; 22 tests
- SwaTrainer+SwagBuffer: cyclic cosine LR schedule; online incremental mean
  (SwaBuffer); E[θ²]-E[θ]² diagonal variance + low-rank deviation columns;
  Box-Muller SWAG sample; 29 tests

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

crates/production/rtx-inference/src/lib.rs

@@ -83,7 +83,7 @@ pub use speculative::{
     EagleConfig,
     EagleDraftModel,
     FusionMethod,
-    LookaheadConfig,
+    LookaheadConfig as SpeculativeLookaheadConfig,
     MedusaConfig,
     MedusaDraftModel,
     NgramPool,
@@ -105,6 +105,9 @@ pub use speculative::{
     StreamedToken,
 };
 
+pub mod lookahead;
+pub use lookahead::{LookaheadDecoder, LookaheadConfig, NGramCache, NGram, LookaheadStats};
+
 // Re-export ONNX Runtime types when feature is enabled
 #[cfg(feature = "onnx-runtime")]
 pub use model_loader::OnnxRuntimeModel;

crates/production/rtx-inference/src/lookahead.rs

@@ -0,0 +1,712 @@
+//! Lookahead Decoding (Fu et al. 2024, arXiv:2402.02057)
+//!
+//! Accelerates autoregressive LLM decoding via Jacobi iteration: instead of
+//! generating one token per forward pass the decoder speculatively generates a
+//! "lookahead window" of W tokens in parallel using an n-gram cache, then
+//! verifies the best continuation against a scoring function that simulates a
+//! single model forward pass.  Accepted tokens advance the sequence without
+//! additional forward passes.
+//!
+//! # Quick start
+//!
+//! ```
+//! use rtx_inference::lookahead::{LookaheadDecoder, LookaheadConfig};
+//!
+//! let config = LookaheadConfig::default();
+//! let mut decoder = LookaheadDecoder::new(config);
+//!
+//! // Warm the cache with some known text (token ids)
+//! decoder.warm_cache(&[1, 2, 3, 4, 5, 6]);
+//!
+//! // score_fn: given the current token sequence, return the single most-likely
+//! // next token (greedy — simulates one model forward pass).
+//! let score_fn = |tokens: &[u32]| -> u32 {
+//!     tokens.last().copied().unwrap_or(0).wrapping_add(1)
+//! };
+//!
+//! let generated = decoder.decode(&[1, 2, 3], 10, score_fn);
+//! assert!(!generated.is_empty());
+//! ```
+
+use std::collections::{HashMap, VecDeque};
+
+// ---------------------------------------------------------------------------
+// NGram
+// ---------------------------------------------------------------------------
+
+/// A single n-gram entry: a prefix of `n-1` tokens paired with the observed
+/// continuation token.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct NGram {
+    /// The `n-1` token prefix (context).
+    pub prefix: Vec<u32>,
+    /// The observed next token.
+    pub next: u32,
+}
+
+// ---------------------------------------------------------------------------
+// NGramCache
+// ---------------------------------------------------------------------------
+
+/// Cache of n-gram statistics accumulated during or before generation.
+///
+/// Internally maps each observed prefix to a frequency table of continuation
+/// tokens.  When the total number of distinct prefixes exceeds `max_size` the
+/// oldest prefix (FIFO) is evicted.
+#[derive(Debug, Clone)]
+pub struct NGramCache {
+    /// The *n* in n-gram: prefix length is `n-1`, context window is `n`.
+    n: usize,
+    /// prefix → { next_token → count }
+    table: HashMap<Vec<u32>, HashMap<u32, usize>>,
+    /// Maximum number of distinct prefixes to store before eviction.
+    max_size: usize,
+    /// FIFO insertion order for eviction.
+    insertion_order: VecDeque<Vec<u32>>,
+    /// Running total of all individual (prefix, next) observations.
+    total_entries: usize,
+}
+
+impl NGramCache {
+    /// Create a new empty cache.
+    ///
+    /// # Arguments
+    /// * `n` — n-gram order (must be ≥ 2; `n == 1` means prefix length 0,
+    ///   which degenerates to a unigram and is not useful here).
+    /// * `max_size` — maximum number of distinct prefixes before FIFO eviction.
+    #[must_use]
+    pub fn new(n: usize, max_size: usize) -> Self {
+        let n = n.max(2); // guard: n must be at least 2 so prefix length ≥ 1
+        Self {
+            n,
+            table: HashMap::new(),
+            max_size: max_size.max(1),
+            insertion_order: VecDeque::new(),
+            total_entries: 0,
+        }
+    }
+
+    /// Slide a window of size `n` over `tokens`, recording every (prefix, next)
+    /// pair.  Does nothing if `tokens.len() < n`.
+    pub fn observe(&mut self, tokens: &[u32]) {
+        let prefix_len = self.n - 1;
+        if tokens.len() < self.n {
+            return;
+        }
+        for i in 0..=(tokens.len() - self.n) {
+            let prefix: Vec<u32> = tokens[i..i + prefix_len].to_vec();
+            let next = tokens[i + prefix_len];
+            self.insert(prefix, next);
+        }
+    }
+
+    /// Return the top-`k` candidate next tokens for `prefix`, sorted by
+    /// observed count descending.  Returns an empty `Vec` if the prefix is
+    /// not in the cache or `prefix.len() != n-1`.
+    #[must_use]
+    pub fn top_candidates(&self, prefix: &[u32], k: usize) -> Vec<u32> {
+        if k == 0 {
+            return Vec::new();
+        }
+        let counts = match self.table.get(prefix) {
+            Some(m) => m,
+            None => return Vec::new(),
+        };
+        let mut pairs: Vec<(u32, usize)> = counts.iter().map(|(&t, &c)| (t, c)).collect();
+        // Stable sort descending by count, ties broken by token id (determinism).
+        pairs.sort_unstable_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
+        pairs.into_iter().take(k).map(|(t, _)| t).collect()
+    }
+
+    /// Number of distinct prefixes currently stored.
+    #[must_use]
+    pub fn num_prefixes(&self) -> usize {
+        self.table.len()
+    }
+
+    /// Total number of individual (prefix, next_token) observations recorded
+    /// (sum of all counts across all entries).
+    #[must_use]
+    pub fn total_observations(&self) -> usize {
+        self.total_entries
+    }
+
+    // ------------------------------------------------------------------
+    // Private helpers
+    // ------------------------------------------------------------------
+
+    fn insert(&mut self, prefix: Vec<u32>, next: u32) {
+        // If prefix is new and we're at capacity, evict the oldest.
+        if !self.table.contains_key(&prefix) {
+            if self.table.len() >= self.max_size {
+                self.evict_oldest();
+            }
+            self.table.insert(prefix.clone(), HashMap::new());
+            self.insertion_order.push_back(prefix.clone());
+        }
+        // Increment count.
+        let count = self
+            .table
+            .get_mut(&prefix)
+            .expect("just inserted above; prefix must exist")
+            .entry(next)
+            .or_insert(0);
+        *count += 1;
+        self.total_entries += 1;
+    }
+
+    fn evict_oldest(&mut self) {
+        if let Some(oldest) = self.insertion_order.pop_front() {
+            if let Some(inner) = self.table.remove(&oldest) {
+                // Deduct evicted counts from the running total.
+                let removed: usize = inner.values().sum();
+                self.total_entries = self.total_entries.saturating_sub(removed);
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// LookaheadConfig
+// ---------------------------------------------------------------------------
+
+/// Configuration for [`LookaheadDecoder`].
+#[derive(Debug, Clone)]
+pub struct LookaheadConfig {
+    /// `W`: number of speculative tokens to draft per step.
+    pub window_size: usize,
+    /// `n` in n-gram; prefix length is `n-1`.
+    pub ngram_order: usize,
+    /// Maximum number of distinct prefixes in the n-gram cache.
+    pub cache_size: usize,
+    /// Top-k n-gram candidates to consider at each draft position.
+    pub max_candidates_per_step: usize,
+    /// Token id that signals end-of-sequence.
+    pub eos_token_id: u32,
+    /// Hard cap on the total output length (initial + generated).
+    pub max_length: usize,
+}
+
+impl Default for LookaheadConfig {
+    fn default() -> Self {
+        Self {
+            window_size: 5,
+            ngram_order: 3,
+            cache_size: 10_000,
+            max_candidates_per_step: 3,
+            eos_token_id: 2,
+            max_length: 256,
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// LookaheadStats
+// ---------------------------------------------------------------------------
+
+/// Cumulative statistics collected during [`LookaheadDecoder::decode`].
+#[derive(Debug, Clone, Default)]
+pub struct LookaheadStats {
+    /// Total tokens appended to the output sequence.
+    pub total_tokens_generated: usize,
+    /// Number of lookahead iterations (decode steps).
+    pub total_steps: usize,
+    /// Tokens accepted directly from n-gram cache predictions.
+    pub total_accepted_from_ngram: usize,
+    /// Tokens accepted from the model (either as correction or no-draft path).
+    pub total_accepted_from_model: usize,
+    /// Average tokens accepted per step (`total_tokens_generated / total_steps`).
+    pub avg_tokens_per_step: f32,
+}
+
+// ---------------------------------------------------------------------------
+// LookaheadDecoder
+// ---------------------------------------------------------------------------
+
+/// CPU-reference lookahead decoder implementing the Jacobi-iteration algorithm
+/// from Fu et al. 2024 (arXiv:2402.02057).
+///
+/// This implementation runs entirely on the CPU and uses a caller-supplied
+/// `score_fn` to simulate a single model forward pass (greedy: returns the
+/// single most-likely next token).  It is suitable for unit testing, benchmarking
+/// the acceptance logic, and as a reference implementation for GPU ports.
+///
+/// # Algorithm overview (per step)
+///
+/// 1. **Lookahead phase** — draft up to `window_size` tokens from the n-gram
+///    cache by repeatedly looking up the most-likely next token for the current
+///    prefix.
+/// 2. **Verification phase** — run `score_fn` greedily over the draft,
+///    accepting tokens while they match and stopping at the first mismatch
+///    (accepting the model's correction).
+/// 3. **Cache update** — observe the newly appended tokens to warm the cache.
+/// 4. **EOS check** — stop if the accepted tail contains the EOS token id.
+pub struct LookaheadDecoder {
+    config: LookaheadConfig,
+    cache: NGramCache,
+    stats: LookaheadStats,
+}
+
+impl LookaheadDecoder {
+    /// Create a new decoder with the given configuration.
+    #[must_use]
+    pub fn new(config: LookaheadConfig) -> Self {
+        let cache = NGramCache::new(config.ngram_order, config.cache_size);
+        Self {
+            config,
+            cache,
+            stats: LookaheadStats::default(),
+        }
+    }
+
+    /// Create a decoder with default configuration.
+    #[must_use]
+    pub fn with_defaults() -> Self {
+        Self::new(LookaheadConfig::default())
+    }
+
+    /// Pre-load n-grams from a reference corpus to warm-start the cache.
+    ///
+    /// Call this before [`decode`](Self::decode) when you have representative
+    /// text available (e.g., a system prompt or previous conversation turns).
+    pub fn warm_cache(&mut self, corpus: &[u32]) {
+        self.cache.observe(corpus);
+    }
+
+    /// Run lookahead decoding.
+    ///
+    /// # Arguments
+    /// * `initial_tokens` — the prompt token ids (not included in the return value).
+    /// * `max_new_tokens` — maximum number of *new* tokens to generate.
+    /// * `score_fn` — closure that accepts the full current token sequence and
+    ///   returns the single greedy-best next token.  This simulates one model
+    ///   forward pass.
+    ///
+    /// # Returns
+    /// The newly generated tokens (does not include `initial_tokens`).
+    pub fn decode<F>(&mut self, initial_tokens: &[u32], max_new_tokens: usize, score_fn: F) -> Vec<u32>
+    where
+        F: Fn(&[u32]) -> u32,
+    {
+        let mut tokens: Vec<u32> = initial_tokens.to_vec();
+        let mut generated: Vec<u32> = Vec::new();
+
+        while generated.len() < max_new_tokens {
+            let remaining = max_new_tokens - generated.len();
+
+            // ----------------------------------------------------------
+            // 1. Lookahead phase: draft up to window_size tokens from cache
+            // ----------------------------------------------------------
+            let mut draft: Vec<u32> = Vec::with_capacity(self.config.window_size);
+            for _ in 0..self.config.window_size.min(remaining) {
+                // Build the prefix from the tail of (tokens + draft so far)
+                let combined_len = tokens.len() + draft.len();
+                let prefix_len = self.config.ngram_order - 1; // n-1
+
+                let prefix: Vec<u32> = if combined_len >= prefix_len {
+                    // Take last prefix_len tokens from the combined sequence.
+                    let start = combined_len - prefix_len;
+                    if draft.len() >= prefix_len {
+                        // The prefix falls entirely within draft.
+                        draft[draft.len() - prefix_len..].to_vec()
+                    } else {
+                        // The prefix spans the end of tokens and the start of draft.
+                        let from_tokens = prefix_len - draft.len();
+                        let token_start = tokens.len().saturating_sub(from_tokens);
+                        tokens[token_start..].iter().chain(draft.iter()).cloned().collect()
+                    }
+                } else {
+                    // Not enough history yet — build what we can.
+                    tokens.iter().chain(draft.iter()).cloned().collect()
+                };
+                let _ = combined_len; // silence unused warning
+
+                let candidates = self
+                    .cache
+                    .top_candidates(&prefix, self.config.max_candidates_per_step);
+                if candidates.is_empty() {
+                    break; // no n-gram available — stop drafting
+                }
+                draft.push(candidates[0]);
+            }
+
+            // ----------------------------------------------------------
+            // 2. Verification phase
+            // ----------------------------------------------------------
+            let mut accepted_this_step: Vec<u32> = Vec::new();
+
+            if !draft.is_empty() {
+                let mut verify_ctx = tokens.clone();
+                for &draft_tok in &draft {
+                    let model_tok = score_fn(&verify_ctx);
+                    if model_tok == draft_tok {
+                        accepted_this_step.push(draft_tok);
+                        verify_ctx.push(draft_tok);
+                        self.stats.total_accepted_from_ngram += 1;
+                    } else {
+                        // Accept model's correction and stop verifying.
+                        accepted_this_step.push(model_tok);
+                        self.stats.total_accepted_from_model += 1;
+                        break;
+                    }
+                    // Don't exceed max_new_tokens in a single step either.
+                    if accepted_this_step.len() >= remaining {
+                        break;
+                    }
+                }
+            } else {
+                // No draft available: just run the model once.
+                let model_tok = score_fn(&tokens);
+                accepted_this_step.push(model_tok);
+                self.stats.total_accepted_from_model += 1;
+            }
+
+            // ----------------------------------------------------------
+            // 3. Update n-gram cache with newly accepted tokens
+            // ----------------------------------------------------------
+            {
+                // Build a small window: last (n-1) tokens from the current
+                // sequence followed by the accepted tokens.
+                let prefix_len = self.config.ngram_order.saturating_sub(1);
+                let context_start = tokens.len().saturating_sub(prefix_len);
+                let new_tokens: Vec<u32> = tokens[context_start..]
+                    .iter()
+                    .chain(accepted_this_step.iter())
+                    .cloned()
+                    .collect();
+                self.cache.observe(&new_tokens);
+            }
+
+            // ----------------------------------------------------------
+            // 4. Extend sequence
+            // ----------------------------------------------------------
+            let to_add = accepted_this_step
+                .len()
+                .min(remaining);
+            let accepted_slice = &accepted_this_step[..to_add];
+
+            tokens.extend_from_slice(accepted_slice);
+            generated.extend_from_slice(accepted_slice);
+            self.stats.total_tokens_generated += accepted_slice.len();
+            self.stats.total_steps += 1;
+
+            // ----------------------------------------------------------
+            // 5. EOS check
+            // ----------------------------------------------------------
+            if accepted_slice.last() == Some(&self.config.eos_token_id) {
+                break;
+            }
+        }
+
+        // Update derived statistic.
+        self.stats.avg_tokens_per_step = self.stats.total_tokens_generated as f32
+            / self.stats.total_steps.max(1) as f32;
+
+        generated
+    }
+
+    /// Access accumulated decoding statistics.
+    #[must_use]
+    pub fn stats(&self) -> &LookaheadStats {
+        &self.stats
+    }
+
+    /// Reset all statistics counters to zero (cache is retained).
+    pub fn reset_stats(&mut self) {
+        self.stats = LookaheadStats::default();
+    }
+
+    /// Access the underlying n-gram cache.
+    #[must_use]
+    pub fn cache(&self) -> &NGramCache {
+        &self.cache
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Helper: score_fn that always predicts last_token + 1 (greedy incrementer).
+    fn incrementer(tokens: &[u32]) -> u32 {
+        tokens.last().copied().unwrap_or(0).wrapping_add(1)
+    }
+
+    // ------------------------------------------------------------------
+    // NGramCache unit tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_ngram_cache_observe_single() {
+        // observe [1,2,3] with n=3: prefix=[1,2], next=3
+        let mut cache = NGramCache::new(3, 100);
+        cache.observe(&[1, 2, 3]);
+        let candidates = cache.top_candidates(&[1, 2], 5);
+        assert_eq!(candidates, vec![3]);
+    }
+
+    #[test]
+    fn test_ngram_cache_top_candidates_most_frequent_first() {
+        // observe token 42 three times and token 7 once for prefix [1,2]
+        let mut cache = NGramCache::new(3, 100);
+        cache.observe(&[1, 2, 42]);
+        cache.observe(&[1, 2, 42]);
+        cache.observe(&[1, 2, 42]);
+        cache.observe(&[1, 2, 7]);
+        let candidates = cache.top_candidates(&[1, 2], 2);
+        assert_eq!(candidates[0], 42, "42 should rank first (count=3)");
+        assert_eq!(candidates[1], 7, "7 should rank second (count=1)");
+    }
+
+    #[test]
+    fn test_ngram_cache_empty_prefix_returns_empty() {
+        let cache = NGramCache::new(3, 100);
+        let candidates = cache.top_candidates(&[99, 100], 5);
+        assert!(candidates.is_empty(), "unknown prefix must return empty vec");
+    }
+
+    #[test]
+    fn test_ngram_cache_observe_window_slides() {
+        // observe [1,2,3,4] with n=3: should store (1,2)→3 and (2,3)→4
+        let mut cache = NGramCache::new(3, 100);
+        cache.observe(&[1, 2, 3, 4]);
+        assert_eq!(cache.top_candidates(&[1, 2], 1), vec![3]);
+        assert_eq!(cache.top_candidates(&[2, 3], 1), vec![4]);
+    }
+
+    #[test]
+    fn test_ngram_cache_top_k_limit() {
+        let mut cache = NGramCache::new(3, 100);
+        // Insert 5 different next tokens for the same prefix.
+        for next in 10u32..15 {
+            cache.observe(&[1, 2, next]);
+        }
+        let candidates = cache.top_candidates(&[1, 2], 3);
+        assert_eq!(candidates.len(), 3, "should return at most k=3 candidates");
+    }
+
+    #[test]
+    fn test_ngram_cache_total_observations() {
+        let mut cache = NGramCache::new(3, 100);
+        // [1,2,3] → 1 observation, [2,3,4] → 1 observation from one call
+        cache.observe(&[1, 2, 3, 4]);
+        assert_eq!(cache.total_observations(), 2);
+        // Add two more.
+        cache.observe(&[5, 6, 7]);
+        assert_eq!(cache.total_observations(), 3);
+    }
+
+    #[test]
+    fn test_ngram_cache_num_prefixes() {
+        let mut cache = NGramCache::new(3, 100);
+        cache.observe(&[1, 2, 3]);
+        assert_eq!(cache.num_prefixes(), 1);
+        cache.observe(&[4, 5, 6]);
+        assert_eq!(cache.num_prefixes(), 2);
+        // Observing the same prefix again must not increase count.
+        cache.observe(&[1, 2, 9]);
+        assert_eq!(cache.num_prefixes(), 2);
+    }
+
+    #[test]
+    fn test_warm_cache_populates_candidates() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        decoder.warm_cache(&[10, 20, 30, 40, 50]);
+        // After warming, the prefix [10,20] should yield at least one candidate.
+        let candidates = decoder.cache().top_candidates(&[10, 20], 3);
+        assert!(!candidates.is_empty(), "warm_cache should populate the cache");
+    }
+
+    // ------------------------------------------------------------------
+    // Decoder integration tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_decode_empty_cache_falls_back_to_model() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig {
+            window_size: 5,
+            ..LookaheadConfig::default()
+        });
+        // No warm_cache call — cache is empty, so every step must call score_fn.
+        let initial = vec![1u32, 2, 3];
+        let generated = decoder.decode(&initial, 5, incrementer);
+        assert_eq!(generated.len(), 5);
+        // All tokens should come from the model (none from ngram).
+        assert_eq!(decoder.stats().total_accepted_from_ngram, 0);
+        assert!(decoder.stats().total_accepted_from_model > 0);
+    }
+
+    #[test]
+    fn test_decode_returns_correct_length() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        let generated = decoder.decode(&[1, 2, 3], 8, incrementer);
+        assert_eq!(generated.len(), 8, "should generate exactly max_new_tokens");
+    }
+
+    #[test]
+    fn test_decode_stops_on_eos() {
+        let eos: u32 = 99;
+        let config = LookaheadConfig {
+            eos_token_id: eos,
+            ..LookaheadConfig::default()
+        };
+        let mut decoder = LookaheadDecoder::new(config);
+        // score_fn always returns the EOS token.
+        let generated = decoder.decode(&[1, 2, 3], 100, |_| eos);
+        // Should stop after the first token (which is EOS).
+        assert!(!generated.is_empty());
+        assert_eq!(*generated.last().unwrap(), eos);
+        assert!(generated.len() < 100, "should stop early on EOS");
+    }
+
+    #[test]
+    fn test_decode_with_matching_cache() {
+        // Warm the cache so it exactly predicts what the incrementer will generate.
+        // incrementer produces: 4,5,6,7,8,... starting from initial=[1,2,3].
+        let corpus: Vec<u32> = (1u32..=20).collect();
+        let config = LookaheadConfig {
+            window_size: 4,
+            ngram_order: 3,
+            eos_token_id: 255,
+            ..LookaheadConfig::default()
+        };
+        let mut decoder = LookaheadDecoder::new(config);
+        decoder.warm_cache(&corpus);
+
+        let generated = decoder.decode(&[1, 2, 3], 10, incrementer);
+        assert_eq!(generated.len(), 10);
+        // Cache matches the model so we should see some ngram acceptances.
+        assert!(
+            decoder.stats().total_accepted_from_ngram > 0,
+            "warm matching cache should yield ngram acceptances"
+        );
+    }
+
+    #[test]
+    fn test_decode_with_mismatched_cache() {
+        // Cache predicts token 99 for every position; model (incrementer) disagrees.
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig {
+            window_size: 4,
+            ngram_order: 2,
+            eos_token_id: 255,
+            ..LookaheadConfig::default()
+        });
+        // Observe [1, 99] → cache predicts 99 after any prefix ending in 1.
+        decoder.warm_cache(&[1, 99, 99, 99, 99, 99]);
+
+        // incrementer starting from [1] will produce 2,3,4,... never 99.
+        let generated = decoder.decode(&[1], 5, incrementer);
+        assert_eq!(generated.len(), 5);
+        // Model should have overridden the cache prediction many times.
+        assert!(decoder.stats().total_accepted_from_model > 0);
+    }
+
+    // ------------------------------------------------------------------
+    // Statistics tests
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_stats_total_tokens_matches_generated_length() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        let generated = decoder.decode(&[1, 2, 3], 12, incrementer);
+        assert_eq!(
+            decoder.stats().total_tokens_generated,
+            generated.len(),
+            "stats counter must equal returned vec length"
+        );
+    }
+
+    #[test]
+    fn test_stats_steps_counted() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        decoder.decode(&[1, 2, 3], 6, incrementer);
+        assert!(
+            decoder.stats().total_steps > 0,
+            "total_steps must be incremented each iteration"
+        );
+    }
+
+    #[test]
+    fn test_stats_avg_tokens_per_step_at_least_one() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        decoder.decode(&[1, 2, 3], 10, incrementer);
+        assert!(
+            decoder.stats().avg_tokens_per_step >= 1.0,
+            "avg_tokens_per_step must be >= 1.0"
+        );
+    }
+
+    #[test]
+    fn test_stats_ngram_accepted_positive_when_cache_matches() {
+        let corpus: Vec<u32> = (1u32..=30).collect();
+        let config = LookaheadConfig {
+            window_size: 4,
+            ngram_order: 3,
+            eos_token_id: 200,
+            ..LookaheadConfig::default()
+        };
+        let mut decoder = LookaheadDecoder::new(config);
+        decoder.warm_cache(&corpus);
+        decoder.decode(&[1, 2, 3], 15, incrementer);
+        assert!(
+            decoder.stats().total_accepted_from_ngram > 0,
+            "matching cache should produce ngram acceptances"
+        );
+    }
+
+    #[test]
+    fn test_reset_stats_clears_all_counters() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        decoder.decode(&[1, 2, 3], 5, incrementer);
+        // Sanity: stats were actually incremented.
+        assert!(decoder.stats().total_steps > 0);
+        decoder.reset_stats();
+        let s = decoder.stats();
+        assert_eq!(s.total_tokens_generated, 0);
+        assert_eq!(s.total_steps, 0);
+        assert_eq!(s.total_accepted_from_ngram, 0);
+        assert_eq!(s.total_accepted_from_model, 0);
+        assert_eq!(s.avg_tokens_per_step, 0.0);
+    }
+
+    #[test]
+    fn test_decode_longer_acceptance_when_cache_warm_and_model_agrees() {
+        // Build a corpus that perfectly matches what incrementer will generate.
+        let corpus: Vec<u32> = (1u32..=50).collect();
+        let config = LookaheadConfig {
+            window_size: 5,
+            ngram_order: 3,
+            cache_size: 10_000,
+            max_candidates_per_step: 3,
+            eos_token_id: 200,
+            max_length: 256,
+        };
+        let mut decoder = LookaheadDecoder::new(config);
+        decoder.warm_cache(&corpus);
+
+        decoder.decode(&[1, 2, 3], 20, incrementer);
+        assert!(
+            decoder.stats().avg_tokens_per_step > 1.0,
+            "avg_tokens_per_step should exceed 1.0 when cache is warm and matches the model"
+        );
+    }
+
+    // Extra coverage for edge conditions and k=0 guard.
+    #[test]
+    fn test_ngram_cache_top_k_zero_returns_empty() {
+        let mut cache = NGramCache::new(3, 100);
+        cache.observe(&[1, 2, 3]);
+        assert!(cache.top_candidates(&[1, 2], 0).is_empty());
+    }
+
+    #[test]
+    fn test_decode_zero_new_tokens_returns_empty() {
+        let mut decoder = LookaheadDecoder::new(LookaheadConfig::default());
+        let generated = decoder.decode(&[1, 2, 3], 0, incrementer);
+        assert!(generated.is_empty(), "max_new_tokens=0 should return empty vec");
+    }
+}

crates/training/rtx-transformers/src/optimizers/mod.rs

@@ -32,6 +32,7 @@ pub use adam::AdamState;
 // pub mod lion;              // TODO: Re-enable in Phase 2
 // pub mod sophia;            // TODO: Re-enable in Phase 2
 pub mod shampoo;
+pub mod soap;
 // pub mod kfac;              // TODO: Re-enable in Phase 2
 // pub mod kfac_advanced;     // TODO: Re-enable in Phase 2
 // pub mod kfac_example;      // TODO: Re-enable in Phase 2
@@ -77,6 +78,7 @@ pub use layer_lr_decay::{LayerLrDecayBuilder, LayerLrDecayConfig, LayerLrSchedul
 // pub use lion::{LionOptimizer, LionConfig};
 // pub use sophia::{SophiaOptimizer, SophiaConfig};
 pub use shampoo::{ShampooOptimizer, ShampooConfig, ShampooParamState};
+pub use soap::{SoapOptimizer, SoapConfig, SoapParamState};
 // pub use kfac::{KFacOptimizer, KFacConfig, LayerType};
 // pub use kfac_advanced::{AdvancedKFacOptimizer, AdvancedKFacConfig, ConvergenceMetrics};
 // pub use ranger::{RangerOptimizer, RangerConfig};

crates/training/rtx-transformers/src/optimizers/soap.rs

@@ -0,0 +1,866 @@
+//! SOAP optimizer — Adam in the eigenbasis of Shampoo's Kronecker-factored preconditioner.
+//!
+//! Implements the SOAP algorithm (Vyas et al. 2024, arXiv:2409.11321 —
+//! "Shampoo As Adam Preconditioner").  SOAP maintains Shampoo's Kronecker-factored
+//! second-moment matrices and re-diagonalises them periodically.  The actual Adam
+//! adaptive update is then computed in the resulting eigenbasis, combining the
+//! curvature information of a second-order method with the per-coordinate
+//! adaptivity of Adam.
+//!
+//! # Algorithm (2-D weight W ∈ ℝ^{m×n})
+//!
+//! ```text
+//! Maintain per-parameter:
+//!   L  ∈ ℝ^{m×m} — accumulated G @ Gᵀ  (left Kronecker factor)
+//!   R  ∈ ℝ^{n×n} — accumulated Gᵀ @ G  (right Kronecker factor)
+//!   Q_L ∈ ℝ^{m×m} — eigenvectors of L (columns)
+//!   Q_R ∈ ℝ^{n×n} — eigenvectors of R
+//!   m1  ∈ ℝ^{m×n} — first  moment in eigenbasis
+//!   m2  ℝ^{m×n} — second moment in eigenbasis (element-wise)
+//!
+//! Each step:
+//!   1. Project:  G_hat = Q_Lᵀ @ G @ Q_R
+//!   2. Adam in eigenbasis (bias-corrected):
+//!        m1 ← β1·m1 + (1−β1)·G_hat
+//!        m2 ← β2·m2 + (1−β2)·G_hat²
+//!        U_hat = (m1/(1−β1ᵗ)) / (sqrt(m2/(1−β2ᵗ)) + ε)
+//!   3. Unproject: U = Q_L @ U_hat @ Q_Rᵀ
+//!   4. Update:    W ← (1 − lr·λ)·W − lr·U
+//!   5. Every `precond_update_freq` steps: re-eigen L and R
+//! ```
+//!
+//! For 1-D parameters (bias vectors) the algorithm reduces to plain Adam — no
+//! projection is performed.
+//!
+//! # Example
+//!
+//! ```rust
+//! use rtx_transformers::optimizers::soap::{SoapConfig, SoapOptimizer};
+//!
+//! let mut opt = SoapOptimizer::new(SoapConfig::default());
+//! opt.register("W", 4, 4);
+//!
+//! let params = vec![1.0_f32; 16];
+//! let grad   = vec![0.1_f32; 16];
+//! let updated = opt.step("W", &params, &grad);
+//! assert_eq!(updated.len(), 16);
+//! ```
+
+use std::collections::HashMap;
+
+// ============================================================================
+// Configuration
+// ============================================================================
+
+/// Configuration for the SOAP optimizer.
+#[derive(Debug, Clone)]
+pub struct SoapConfig {
+    /// Learning rate (default 0.001).
+    pub lr: f32,
+    /// Adam beta1 — first-moment decay (default 0.95).
+    pub beta1: f32,
+    /// Adam beta2 — second-moment decay (default 0.95).
+    pub beta2: f32,
+    /// Epsilon for numerical stability in the Adam denominator (default 1e-8).
+    pub epsilon: f32,
+    /// Decoupled weight-decay coefficient λ (default 0.01).
+    pub weight_decay: f32,
+    /// How many steps between Kronecker-factor eigen-decompositions (default 10).
+    pub precond_update_freq: usize,
+    /// Maximum Jacobi sweeps for eigen-decomposition (default 20).
+    pub max_sweeps: usize,
+    /// Skip eigen-decomposition for dimensions larger than this (default 512).
+    /// Parameters whose row or column dimension exceeds this value fall back to
+    /// plain Adam on the unprojected gradient.
+    pub max_precond_dim: usize,
+}
+
+impl Default for SoapConfig {
+    fn default() -> Self {
+        Self {
+            lr: 0.001,
+            beta1: 0.95,
+            beta2: 0.95,
+            epsilon: 1e-8,
+            weight_decay: 0.01,
+            precond_update_freq: 10,
+            max_sweeps: 20,
+            max_precond_dim: 512,
+        }
+    }
+}
+
+// ============================================================================
+// Per-parameter state
+// ============================================================================
+
+/// Optimizer state stored for a single parameter.
+pub struct SoapParamState {
+    /// Left Kronecker factor L = Σ G @ Gᵀ  [rows × rows], row-major.
+    pub l_factor: Vec<f32>,
+    /// Right Kronecker factor R = Σ Gᵀ @ G  [cols × cols], row-major.
+    pub r_factor: Vec<f32>,
+    /// Eigenvectors of L, columns = eigenvectors  [rows × rows], row-major.
+    pub q_l: Vec<f32>,
+    /// Eigenvectors of R, columns = eigenvectors  [cols × cols], row-major.
+    pub q_r: Vec<f32>,
+    /// First moment in the eigenbasis  [rows × cols], row-major.
+    pub m1: Vec<f32>,
+    /// Second moment (element-wise) in the eigenbasis  [rows × cols], row-major.
+    pub m2: Vec<f32>,
+    /// Number of completed steps for this parameter.
+    pub step: usize,
+    /// Number of rows of the parameter matrix.
+    pub rows: usize,
+    /// Number of columns of the parameter matrix.
+    /// When `cols == 1` the 1-D plain-Adam code path is used.
+    pub cols: usize,
+}
+
+// ============================================================================
+// Optimizer
+// ============================================================================
+
+/// SOAP optimizer — CPU reference implementation (pure Rust, no GPU).
+pub struct SoapOptimizer {
+    config: SoapConfig,
+    states: HashMap<String, SoapParamState>,
+}
+
+impl SoapOptimizer {
+    /// Create a new `SoapOptimizer` with the given configuration.
+    #[must_use]
+    pub fn new(config: SoapConfig) -> Self {
+        Self {
+            config,
+            states: HashMap::new(),
+        }
+    }
+
+    /// Convenience constructor: default config with `lr` overridden.
+    #[must_use]
+    pub fn with_lr(lr: f32) -> Self {
+        Self::new(SoapConfig { lr, ..SoapConfig::default() })
+    }
+
+    /// Register a parameter before the first [`step`](SoapOptimizer::step) call.
+    ///
+    /// `rows` and `cols` describe the logical matrix shape of the (flattened)
+    /// parameter.  Set `cols = 1` to request the 1-D plain-Adam code path.
+    pub fn register(&mut self, name: &str, rows: usize, cols: usize) {
+        let n = rows * cols;
+
+        // Initialise Kronecker factors to zero; eigen-bases to identity.
+        let l_factor = vec![0.0_f32; rows * rows];
+        let r_factor = vec![0.0_f32; cols * cols];
+        let q_l = Self::identity(rows);
+        let q_r = Self::identity(cols);
+
+        self.states.insert(
+            name.to_owned(),
+            SoapParamState {
+                l_factor,
+                r_factor,
+                q_l,
+                q_r,
+                m1: vec![0.0_f32; n],
+                m2: vec![0.0_f32; n],
+                step: 0,
+                rows,
+                cols,
+            },
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Public step
+    // -----------------------------------------------------------------------
+
+    /// Perform one optimizer step for the named parameter.
+    ///
+    /// Returns the updated parameter vector (same length as `params`).
+    ///
+    /// # Panics
+    ///
+    /// Panics if `name` was not previously registered with [`register`](SoapOptimizer::register).
+    pub fn step(&mut self, name: &str, params: &[f32], grad: &[f32]) -> Vec<f32> {
+        // Pull config values out to avoid simultaneous borrow of `self`.
+        let beta1 = self.config.beta1;
+        let beta2 = self.config.beta2;
+        let epsilon = self.config.epsilon;
+        let lr = self.config.lr;
+        let weight_decay = self.config.weight_decay;
+        let precond_freq = self.config.precond_update_freq;
+        let max_sweeps = self.config.max_sweeps;
+        let max_precond_dim = self.config.max_precond_dim;
+
+        let state = self
+            .states
+            .get_mut(name)
+            .unwrap_or_else(|| panic!("SoapOptimizer: parameter '{name}' was not registered"));
+
+        state.step += 1;
+        let step = state.step;
+        let rows = state.rows;
+        let cols = state.cols;
+
+        // Bias-correction denominators.
+        let bc1 = 1.0 - beta1.powi(step as i32);
+        let bc2 = 1.0 - beta2.powi(step as i32);
+
+        // -----------------------------------------------------------------------
+        // 1-D path: plain Adam on the raw gradient, no Kronecker machinery.
+        // -----------------------------------------------------------------------
+        if cols == 1 {
+            for (m, g) in state.m1.iter_mut().zip(grad.iter()) {
+                *m = beta1 * *m + (1.0 - beta1) * g;
+            }
+            for (v, g) in state.m2.iter_mut().zip(grad.iter()) {
+                *v = beta2 * *v + (1.0 - beta2) * g * g;
+            }
+            return params
+                .iter()
+                .zip(state.m1.iter())
+                .zip(state.m2.iter())
+                .map(|((p, m), v)| {
+                    let m_hat = m / bc1;
+                    let v_hat = v / bc2;
+                    let update = m_hat / (v_hat.sqrt() + epsilon);
+                    (1.0 - lr * weight_decay) * p - lr * update
+                })
+                .collect();
+        }
+
+        // -----------------------------------------------------------------------
+        // 2-D path: SOAP (Kronecker-factored preconditioning + Adam in eigenbasis).
+        // -----------------------------------------------------------------------
+
+        // Determine whether we can apply eigenbasis projection.
+        let use_precond = rows <= max_precond_dim && cols <= max_precond_dim;
+
+        // Accumulate Kronecker factors: L += G @ Gᵀ,  R += Gᵀ @ G.
+        if use_precond {
+            let gt = Self::transpose(grad, rows, cols);
+            let g_gt = Self::matmul(grad, &gt, rows, cols, rows);       // [rows × rows]
+            let gt_g = Self::matmul(&gt, grad, cols, rows, cols);       // [cols × cols]
+
+            for (l, v) in state.l_factor.iter_mut().zip(g_gt.iter()) {
+                *l += v;
+            }
+            for (r, v) in state.r_factor.iter_mut().zip(gt_g.iter()) {
+                *r += v;
+            }
+
+            // Re-diagonalise periodically.
+            if step % precond_freq == 0 || step == 1 {
+                let (_, v_l) = Self::jacobi_eigen(&state.l_factor, rows, max_sweeps);
+                let (_, v_r) = Self::jacobi_eigen(&state.r_factor, cols, max_sweeps);
+                state.q_l = v_l;
+                state.q_r = v_r;
+            }
+        }
+
+        // Project gradient into eigenbasis: G_hat = Q_Lᵀ @ G @ Q_R.
+        let g_hat = if use_precond {
+            let q_lt = Self::transpose(&state.q_l, rows, rows);         // [rows × rows]
+            let tmp  = Self::matmul(&q_lt, grad, rows, rows, cols);     // [rows × cols]
+            Self::matmul(&tmp, &state.q_r, rows, cols, cols)            // [rows × cols]
+        } else {
+            grad.to_vec()
+        };
+
+        // Adam in eigenbasis.
+        for (m, g) in state.m1.iter_mut().zip(g_hat.iter()) {
+            *m = beta1 * *m + (1.0 - beta1) * g;
+        }
+        for (v, g) in state.m2.iter_mut().zip(g_hat.iter()) {
+            *v = beta2 * *v + (1.0 - beta2) * g * g;
+        }
+
+        let u_hat: Vec<f32> = state
+            .m1
+            .iter()
+            .zip(state.m2.iter())
+            .map(|(m, v)| {
+                let m_hat = m / bc1;
+                let v_hat = v / bc2;
+                m_hat / (v_hat.sqrt() + epsilon)
+            })
+            .collect();
+
+        // Unproject: U = Q_L @ U_hat @ Q_Rᵀ.
+        let update = if use_precond {
+            let q_rt = Self::transpose(&state.q_r, cols, cols);         // [cols × cols]
+            let tmp  = Self::matmul(&state.q_l, &u_hat, rows, rows, cols); // [rows × cols]
+            Self::matmul(&tmp, &q_rt, rows, cols, cols)                 // [rows × cols]
+        } else {
+            u_hat
+        };
+
+        // Parameter update with decoupled weight decay.
+        //   W ← (1 − lr·λ)·W − lr·U
+        params
+            .iter()
+            .zip(update.iter())
+            .map(|(p, u)| (1.0 - lr * weight_decay) * p - lr * u)
+            .collect()
+    }
+
+    // -----------------------------------------------------------------------
+    // Jacobi eigenvalue algorithm
+    // -----------------------------------------------------------------------
+
+    /// Compute all eigenvalues and eigenvectors of a symmetric matrix `a` of
+    /// size `n × n` (row-major).
+    ///
+    /// Uses the classical Jacobi eigenvalue algorithm with at most `max_sweeps`
+    /// sweeps.  Each sweep iterates over all off-diagonal pairs (p, q) and
+    /// applies a Jacobi plane rotation to zero the (p, q) element.  The
+    /// algorithm converges quadratically for small-to-medium dense symmetric
+    /// matrices.
+    ///
+    /// Returns `(eigenvalues, eigenvectors)` where `eigenvectors` is a
+    /// column-major `n × n` matrix: column `i` is the eigenvector for
+    /// `eigenvalues[i]`.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use rtx_transformers::optimizers::soap::SoapOptimizer;
+    /// // Diagonal matrix: eigenvalues are the diagonal entries.
+    /// let a = vec![3.0_f32, 0.0, 0.0, 7.0];
+    /// let (vals, _vecs) = SoapOptimizer::jacobi_eigen(&a, 2, 20);
+    /// assert!((vals[0] - 3.0).abs() < 1e-5 || (vals[0] - 7.0).abs() < 1e-5);
+    /// ```
+    pub fn jacobi_eigen(a: &[f32], n: usize, max_sweeps: usize) -> (Vec<f32>, Vec<f32>) {
+        debug_assert_eq!(a.len(), n * n);
+
+        let mut a = a.to_vec();                 // working copy — will become diagonal
+        let mut v = Self::identity(n);          // accumulate rotations: v = Q
+
+        for _sweep in 0..max_sweeps {
+            // Find the largest off-diagonal element in magnitude.
+            let mut max_off = 0.0_f32;
+            let mut p = 0usize;
+            let mut q = 1usize;
+            for i in 0..n {
+                for j in (i + 1)..n {
+                    let val = a[i * n + j].abs();
+                    if val > max_off {
+                        max_off = val;
+                        p = i;
+                        q = j;
+                    }
+                }
+            }
+            // Converged: all off-diagonal elements are effectively zero.
+            if max_off < 1e-10 {
+                break;
+            }
+
+            // Compute Jacobi rotation angle so that the (p,q) element becomes zero.
+            // θ = (a[q,q] − a[p,p]) / (2 · a[p,q])
+            // t = sign(θ) / (|θ| + √(1 + θ²))   — smallest root of t²+2θt−1=0
+            let a_pp = a[p * n + p];
+            let a_qq = a[q * n + q];
+            let a_pq = a[p * n + q];
+            let theta = (a_qq - a_pp) / (2.0 * a_pq);
+            let t = if theta >= 0.0 {
+                1.0 / (theta + (1.0 + theta * theta).sqrt())
+            } else {
+                1.0 / (theta - (1.0 + theta * theta).sqrt())
+            };
+            let c = 1.0 / (1.0 + t * t).sqrt();
+            let s = t * c;
+
+            // Apply the symmetric Givens rotation to A.
+            Self::apply_jacobi_rotation(&mut a, n, p, q, c, s);
+            // Accumulate the rotation into the eigenvector matrix V.
+            Self::apply_jacobi_rotation_to_cols(&mut v, n, p, q, c, s);
+        }
+
+        let eigenvalues: Vec<f32> = (0..n).map(|i| a[i * n + i]).collect();
+        (eigenvalues, v)
+    }
+
+    // -----------------------------------------------------------------------
+    // Private helpers
+    // -----------------------------------------------------------------------
+
+    /// Apply a symmetric Jacobi rotation to the symmetric matrix `a` in-place.
+    ///
+    /// After the rotation, `a[p,q]` and `a[q,p]` will be (near) zero, and
+    /// the diagonal elements `a[p,p]` and `a[q,q]` are updated accordingly.
+    fn apply_jacobi_rotation(a: &mut [f32], n: usize, p: usize, q: usize, c: f32, s: f32) {
+        // Snapshot the values that are about to change so intermediate reads
+        // are consistent.
+        let a_pp = a[p * n + p];
+        let a_qq = a[q * n + q];
+        let a_pq = a[p * n + q];
+
+        // Update diagonal elements.
+        a[p * n + p] = c * c * a_pp - 2.0 * s * c * a_pq + s * s * a_qq;
+        a[q * n + q] = s * s * a_pp + 2.0 * s * c * a_pq + c * c * a_qq;
+        a[p * n + q] = 0.0;
+        a[q * n + p] = 0.0;
+
+        // Update all other rows/columns involving p and q.
+        for r in 0..n {
+            if r == p || r == q {
+                continue;
+            }
+            let a_rp = a[r * n + p];
+            let a_rq = a[r * n + q];
+            let new_rp =  c * a_rp - s * a_rq;
+            let new_rq =  s * a_rp + c * a_rq;
+            a[r * n + p] = new_rp;
+            a[p * n + r] = new_rp;
+            a[r * n + q] = new_rq;
+            a[q * n + r] = new_rq;
+        }
+    }
+
+    /// Accumulate a Jacobi column rotation into the eigenvector matrix `v`.
+    ///
+    /// This updates columns `p` and `q` of `v` (each of length `n`).
+    fn apply_jacobi_rotation_to_cols(v: &mut [f32], n: usize, p: usize, q: usize, c: f32, s: f32) {
+        for r in 0..n {
+            let v_rp = v[r * n + p];
+            let v_rq = v[r * n + q];
+            v[r * n + p] =  c * v_rp - s * v_rq;
+            v[r * n + q] =  s * v_rp + c * v_rq;
+        }
+    }
+
+    /// Dense matrix multiply: `A[m × k] @ B[k × n] → C[m × n]` (row-major).
+    fn matmul(a: &[f32], b: &[f32], m: usize, k: usize, n: usize) -> Vec<f32> {
+        debug_assert_eq!(a.len(), m * k);
+        debug_assert_eq!(b.len(), k * n);
+
+        let mut c = vec![0.0_f32; m * n];
+        for i in 0..m {
+            for p in 0..k {
+                let a_ip = a[i * k + p];
+                for j in 0..n {
+                    c[i * n + j] += a_ip * b[p * n + j];
+                }
+            }
+        }
+        c
+    }
+
+    /// Transpose `A[m × n]` → `Aᵀ[n × m]` (row-major).
+    fn transpose(a: &[f32], m: usize, n: usize) -> Vec<f32> {
+        debug_assert_eq!(a.len(), m * n);
+
+        let mut at = vec![0.0_f32; n * m];
+        for i in 0..m {
+            for j in 0..n {
+                at[j * m + i] = a[i * n + j];
+            }
+        }
+        at
+    }
+
+    /// Construct the `n × n` identity matrix (row-major).
+    fn identity(n: usize) -> Vec<f32> {
+        let mut eye = vec![0.0_f32; n * n];
+        for i in 0..n {
+            eye[i * n + i] = 1.0;
+        }
+        eye
+    }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // Tolerance for floating-point comparisons.
+    const TOL: f32 = 1e-4;
+
+    fn approx_eq(a: f32, b: f32, tol: f32) -> bool {
+        (a - b).abs() < tol
+    }
+
+    // -----------------------------------------------------------------------
+    // Configuration
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_default_config() {
+        let cfg = SoapConfig::default();
+        assert!(approx_eq(cfg.lr,           0.001, 1e-7), "lr default");
+        assert!(approx_eq(cfg.beta1,        0.95,  1e-7), "beta1 default");
+        assert!(approx_eq(cfg.beta2,        0.95,  1e-7), "beta2 default");
+        assert!(approx_eq(cfg.epsilon,      1e-8,  1e-14), "epsilon default");
+        assert!(approx_eq(cfg.weight_decay, 0.01,  1e-7), "weight_decay default");
+        assert_eq!(cfg.precond_update_freq, 10,  "precond_update_freq default");
+        assert_eq!(cfg.max_sweeps,          20,  "max_sweeps default");
+        assert_eq!(cfg.max_precond_dim,     512, "max_precond_dim default");
+    }
+
+    #[test]
+    fn test_with_lr() {
+        let opt = SoapOptimizer::with_lr(0.005);
+        assert!(approx_eq(opt.config.lr, 0.005, 1e-7), "lr override");
+        // Other fields stay at defaults.
+        assert!(approx_eq(opt.config.beta1, 0.95, 1e-7));
+        assert_eq!(opt.config.max_sweeps, 20);
+    }
+
+    // -----------------------------------------------------------------------
+    // Linear-algebra helpers (tested via public wrappers / direct calls)
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_identity_matrix() {
+        // 3×3 identity should have ones on the diagonal and zeros elsewhere.
+        let eye = SoapOptimizer::identity(3);
+        assert_eq!(eye.len(), 9);
+        for i in 0..3 {
+            for j in 0..3 {
+                let expected = if i == j { 1.0 } else { 0.0 };
+                assert!(
+                    approx_eq(eye[i * 3 + j], expected, 1e-9),
+                    "identity[{i},{j}] wrong"
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_matmul_2x2() {
+        // [[1,2],[3,4]] @ [[5,6],[7,8]] = [[19,22],[43,50]]
+        let a = vec![1.0_f32, 2.0, 3.0, 4.0];
+        let b = vec![5.0_f32, 6.0, 7.0, 8.0];
+        let c = SoapOptimizer::matmul(&a, &b, 2, 2, 2);
+        let expected = [19.0_f32, 22.0, 43.0, 50.0];
+        for (r, e) in c.iter().zip(expected.iter()) {
+            assert!(approx_eq(*r, *e, TOL), "matmul: got {r}, expected {e}");
+        }
+    }
+
+    #[test]
+    fn test_transpose_2x3() {
+        // A = [[1,2,3],[4,5,6]]  →  Aᵀ = [[1,4],[2,5],[3,6]]
+        let a = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
+        let at = SoapOptimizer::transpose(&a, 2, 3);
+        let expected = vec![1.0_f32, 4.0, 2.0, 5.0, 3.0, 6.0];
+        assert_eq!(at, expected, "transpose mismatch");
+    }
+
+    // -----------------------------------------------------------------------
+    // Jacobi eigen-decomposition
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_jacobi_2x2_diagonal() {
+        // A diagonal matrix is already diagonalised; eigenvalues are the diagonal entries.
+        let a = vec![3.0_f32, 0.0, 0.0, 7.0];
+        let (vals, _vecs) = SoapOptimizer::jacobi_eigen(&a, 2, 20);
+        // Eigenvalues may be returned in any order.
+        let mut sorted = vals.clone();
+        sorted.sort_by(|x, y| x.partial_cmp(y).unwrap());
+        assert!(approx_eq(sorted[0], 3.0, TOL), "smallest eigenvalue");
+        assert!(approx_eq(sorted[1], 7.0, TOL), "largest  eigenvalue");
+    }
+
+    #[test]
+    fn test_jacobi_2x2_symmetric() {
+        // A = [[2,1],[1,2]] has eigenvalues 3 and 1.
+        let a = vec![2.0_f32, 1.0, 1.0, 2.0];
+        let (vals, _vecs) = SoapOptimizer::jacobi_eigen(&a, 2, 20);
+        let mut sorted = vals.clone();
+        sorted.sort_by(|x, y| x.partial_cmp(y).unwrap());
+        assert!(approx_eq(sorted[0], 1.0, TOL), "smallest eigenvalue");
+        assert!(approx_eq(sorted[1], 3.0, TOL), "largest  eigenvalue");
+    }
+
+    #[test]
+    fn test_jacobi_eigenvectors_orthogonal() {
+        // For a 3×3 symmetric positive-definite matrix, Vᵀ @ V should be ≈ I.
+        let a = vec![
+            4.0_f32, 1.0, 0.5,
+            1.0,     3.0, 0.8,
+            0.5,     0.8, 2.0,
+        ];
+        let n = 3;
+        let (_vals, vecs) = SoapOptimizer::jacobi_eigen(&a, n, 30);
+
+        // Compute Vᵀ @ V.
+        let vt = SoapOptimizer::transpose(&vecs, n, n);
+        let vtv = SoapOptimizer::matmul(&vt, &vecs, n, n, n);
+
+        // Check against identity.
+        for i in 0..n {
+            for j in 0..n {
+                let expected = if i == j { 1.0 } else { 0.0 };
+                assert!(
+                    approx_eq(vtv[i * n + j], expected, TOL),
+                    "Vᵀ·V[{i},{j}] = {}, expected {}",
+                    vtv[i * n + j],
+                    expected
+                );
+            }
+        }
+    }
+
+    #[test]
+    fn test_jacobi_reconstruction() {
+        // V @ diag(λ) @ Vᵀ ≈ A  (within tolerance).
+        let a = vec![
+            5.0_f32, 2.0,
+            2.0,     3.0,
+        ];
+        let n = 2;
+        let (vals, vecs) = SoapOptimizer::jacobi_eigen(&a, n, 20);
+
+        // Build diag(λ) as a dense matrix.
+        let mut diag = vec![0.0_f32; n * n];
+        for i in 0..n {
+            diag[i * n + i] = vals[i];
+        }
+
+        let vt = SoapOptimizer::transpose(&vecs, n, n);
+        // V @ diag(λ)
+        let v_diag = SoapOptimizer::matmul(&vecs, &diag, n, n, n);
+        // V @ diag(λ) @ Vᵀ
+        let reconstructed = SoapOptimizer::matmul(&v_diag, &vt, n, n, n);
+
+        for idx in 0..n * n {
+            assert!(
+                approx_eq(reconstructed[idx], a[idx], TOL),
+                "reconstruction[{idx}]: got {}, expected {}",
+                reconstructed[idx],
+                a[idx]
+            );
+        }
+    }
+
+    // -----------------------------------------------------------------------
+    // Optimizer step — structural
+    // -----------------------------------------------------------------------
+
+    #[test]
+    #[should_panic(expected = "not registered")]
+    fn test_register_then_step_panics_if_unregistered() {
+        let mut opt = SoapOptimizer::with_lr(0.01);
+        let params = vec![1.0_f32; 4];
+        let grad   = vec![0.1_f32; 4];
+        let _ = opt.step("unregistered_param", &params, &grad);
+    }
+
+    #[test]
+    fn test_step_1d_plain_adam() {
+        // 1-D parameter (cols=1) should use plain Adam; params must change.
+        let mut opt = SoapOptimizer::with_lr(0.01);
+        opt.register("bias", 4, 1);
+
+        let params  = vec![1.0_f32; 4];
+        let grad    = vec![1.0_f32; 4];
+        let updated = opt.step("bias", &params, &grad);
+
+        assert_eq!(updated.len(), 4, "output length");
+        let changed = params.iter().zip(updated.iter()).any(|(p, u)| (p - u).abs() > 1e-9);
+        assert!(changed, "1-D step should change params");
+    }
+
+    #[test]
+    fn test_step_2d_updates_params() {
+        // 2-D parameter: params must change after one step.
+        let mut opt = SoapOptimizer::with_lr(0.001);
+        opt.register("W", 4, 4);
+
+        let params  = vec![1.0_f32; 16];
+        let grad    = vec![0.1_f32; 16];
+        let updated = opt.step("W", &params, &grad);
+
+        assert_eq!(updated.len(), 16, "output length");
+        let changed = params.iter().zip(updated.iter()).any(|(p, u)| (p - u).abs() > 1e-9);
+        assert!(changed, "2-D step should change params");
+    }
+
+    #[test]
+    fn test_step_reduces_loss_direction() {
+        // For a quadratic loss L = 0.5·||W||² the gradient equals W.
+        // One optimizer step with a positive learning rate should move W toward zero.
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr: 0.01,
+            weight_decay: 0.0, // isolate Adam update
+            ..SoapConfig::default()
+        });
+        opt.register("W", 3, 3);
+
+        let params: Vec<f32> = vec![2.0_f32; 9];
+        let grad = params.clone(); // ∇L = W
+
+        let updated = opt.step("W", &params, &grad);
+
+        let norm_before: f32 = params.iter().map(|x| x * x).sum::<f32>().sqrt();
+        let norm_after:  f32 = updated.iter().map(|x| x * x).sum::<f32>().sqrt();
+        assert!(
+            norm_after < norm_before,
+            "update should reduce ||W||: before={norm_before}, after={norm_after}"
+        );
+    }
+
+    #[test]
+    fn test_kronecker_factors_accumulate() {
+        // After two steps the L and R Kronecker factors should be non-zero.
+        let mut opt = SoapOptimizer::with_lr(0.001);
+        opt.register("W", 3, 3);
+
+        let params = vec![1.0_f32; 9];
+        let grad   = vec![0.5_f32; 9];
+
+        opt.step("W", &params, &grad);
+        let p2 = opt.step("W", &params, &grad);
+
+        let state = opt.states.get("W").unwrap();
+
+        let l_norm: f32 = state.l_factor.iter().map(|x| x * x).sum::<f32>().sqrt();
+        let r_norm: f32 = state.r_factor.iter().map(|x| x * x).sum::<f32>().sqrt();
+
+        assert!(l_norm > 0.0, "L factor should be non-zero after 2 steps");
+        assert!(r_norm > 0.0, "R factor should be non-zero after 2 steps");
+
+        // Sanity: second step produced something.
+        assert_eq!(p2.len(), 9);
+    }
+
+    #[test]
+    fn test_weight_decay_shrinks_params() {
+        // With non-zero weight decay and near-zero gradient, params should shrink.
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr: 0.1,
+            weight_decay: 0.5, // strong decay to make the effect visible
+            ..SoapConfig::default()
+        });
+        opt.register("W", 2, 2);
+
+        let params = vec![1.0_f32; 4];
+        // Use a tiny grad so the Adam term is tiny compared to weight decay.
+        let grad = vec![1e-6_f32; 4];
+
+        let updated = opt.step("W", &params, &grad);
+
+        for &u in &updated {
+            assert!(u < 1.0_f32, "weight decay should shrink params, got {u}");
+        }
+    }
+
+    #[test]
+    fn test_bias_correction_step1() {
+        // At step=1 with beta1=0.95:
+        //   m1_raw = (1−0.95) * g = 0.05 * g
+        //   bc1    = 1 − 0.95^1  = 0.05
+        //   m1_hat = m1_raw / bc1 = g
+        //
+        // Similarly for beta2=0.95:
+        //   m2_raw = 0.05 * g²
+        //   bc2    = 0.05
+        //   m2_hat = g²
+        //   denom  = sqrt(g²) + ε = |g| + ε
+        //
+        // So update ≈ g / (|g| + ε)  for element-wise scalar g.
+        // For g=1.0: update ≈ 1.0 / (1.0 + 1e-8) ≈ 1.0.
+        // new_param = (1 − lr·λ)·p − lr·update
+        //
+        // We verify that:
+        //   1. the update direction is negative (params decrease for positive grad)
+        //   2. the magnitude is consistent with bias-correction at step 1.
+        let lr = 0.01_f32;
+        let wd = 0.0_f32;
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr,
+            weight_decay: wd,
+            ..SoapConfig::default()
+        });
+        opt.register("b", 1, 1); // 1-D — direct Adam
+
+        let p0 = vec![0.0_f32];
+        let g  = vec![1.0_f32];
+        let p1 = opt.step("b", &p0, &g);
+
+        // bias-corrected update ≈ 1.0 / (1.0 + 1e-8) ≈ 1.0
+        // new_param ≈ 0.0 - 0.01 * 1.0 = -0.01
+        let expected = -lr; // ≈ -0.01
+        assert!(
+            approx_eq(p1[0], expected, 1e-3),
+            "step-1 bias correction: got {}, expected ≈ {}",
+            p1[0],
+            expected
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Additional correctness checks
+    // -----------------------------------------------------------------------
+
+    #[test]
+    fn test_step_1d_multiple_steps_converge() {
+        // With a constant gradient and weight_decay=0, repeated Adam steps on a
+        // 1-D param should monotonically move the parameter in the direction of
+        // steepest descent (toward -∞ for positive gradient), and the step size
+        // should stabilise as the moments converge.
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr: 0.01,
+            weight_decay: 0.0,
+            ..SoapConfig::default()
+        });
+        opt.register("b", 1, 1);
+
+        let mut p = vec![1.0_f32];
+        let g = vec![1.0_f32];
+
+        for _ in 0..60 {
+            p = opt.step("b", &p, &g);
+        }
+        // With lr=0.01 and a bias-corrected Adam update ≈ 1.0 per step,
+        // after 60 steps from p=1.0 the parameter should be well below 0.5.
+        assert!(p[0] < 0.5, "param should decrease monotonically: {}", p[0]);
+    }
+
+    #[test]
+    fn test_step_2d_precond_freq() {
+        // Requesting precond_update_freq=1 means re-eigen every step; should not panic.
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr: 0.001,
+            precond_update_freq: 1,
+            ..SoapConfig::default()
+        });
+        opt.register("W", 4, 4);
+
+        let params = vec![0.5_f32; 16];
+        let grad   = vec![0.1_f32; 16];
+
+        for _ in 0..5 {
+            let _ = opt.step("W", &params, &grad);
+        }
+    }
+
+    #[test]
+    fn test_step_2d_large_dim_falls_back() {
+        // When rows or cols exceed max_precond_dim the optimizer should fall
+        // back to plain Adam and still produce output of the right length.
+        let mut opt = SoapOptimizer::new(SoapConfig {
+            lr: 0.001,
+            max_precond_dim: 2, // deliberately tiny
+            ..SoapConfig::default()
+        });
+        opt.register("W", 4, 4); // 4 > 2 → fallback path
+
+        let params: Vec<f32> = (0..16).map(|i| i as f32 * 0.1).collect();
+        let grad:   Vec<f32> = vec![0.01_f32; 16];
+        let updated = opt.step("W", &params, &grad);
+
+        assert_eq!(updated.len(), 16);
+        let changed = params.iter().zip(updated.iter()).any(|(p, u)| (p - u).abs() > 1e-9);
+        assert!(changed, "fallback path should still update params");
+    }
+}

crates/training/rtx-transformers/src/training/mod.rs

@@ -55,6 +55,9 @@ pub use draft_distill::{
 pub use gradient_noise_scale::{GnsEstimate, GnsTracker, GradientNoiseScale};
 pub use model_ema::{ModelEma, ModelEmaConfig};
 
+pub mod swa;
+pub use swa::{SwaBuffer, SwaSchedule, SwaTrainer, SwagBuffer};
+
 /// Training state structure
 #[derive(Debug, Clone)]
 pub struct TrainingState {

crates/training/rtx-transformers/src/training/swa.rs

@@ -0,0 +1,916 @@
+//! Stochastic Weight Averaging (SWA) and SWA-Gaussian (SWAG).
+//!
+//! **SWA** (Izmailov et al. 2018, <https://arxiv.org/abs/1803.05407>) finds flatter loss
+//! minima and achieves better generalisation than standard SGD by averaging model weights
+//! across the final portion of a cosine-annealing training run.  Unlike EMA, SWA
+//! accumulates a simple arithmetic mean over discrete snapshots taken at the end of each
+//! cosine cycle.
+//!
+//! **SWAG** (Maddox et al. 2019, <https://arxiv.org/abs/1902.02476>) extends SWA to
+//! Bayesian uncertainty estimation.  It fits a Gaussian posterior over the weight space
+//! using the running mean and a low-rank + diagonal covariance estimated from the SWA
+//! trajectory, enabling cheap Monte-Carlo uncertainty sampling at inference time.
+//!
+//! # Quick-start
+//!
+//! ```rust
+//! use rtx_transformers::training::{SwaSchedule, SwaTrainer};
+//!
+//! let schedule = SwaSchedule::new(1e-2, 1000, 5000);
+//! let mut trainer = SwaTrainer::new(schedule);
+//!
+//! // Simulate a training loop:
+//! for step in 0..7001 {
+//!     let params = [("weight", &[0.1_f32, 0.2, 0.3][..])];
+//!     let snapshotted = trainer.update(&params);
+//!     let _ = snapshotted; // true when SWA average was updated
+//! }
+//! let swa_w = trainer.get_swa_params("weight");
+//! assert!(swa_w.is_some());
+//! ```
+
+use std::collections::HashMap;
+
+// ---------------------------------------------------------------------------
+// LCG-based normal sampler (Box-Muller) — used by SWAG for deterministic tests
+// ---------------------------------------------------------------------------
+
+/// Draw one standard-normal sample using a 64-bit LCG and the Box-Muller transform.
+///
+/// The LCG parameters are the Knuth/MMIX constants; the same values are used by
+/// several standard-library RNGs.  The caller advances the seed in place so that
+/// successive calls produce an independent stream.
+fn lcg_normal(seed: &mut u64) -> f32 {
+    *seed = seed
+        .wrapping_mul(6_364_136_223_846_793_005)
+        .wrapping_add(1_442_695_040_888_963_407);
+    let u1 = (*seed >> 11) as f32 / (1u64 << 53) as f32;
+
+    *seed = seed
+        .wrapping_mul(6_364_136_223_846_793_005)
+        .wrapping_add(1_442_695_040_888_963_407);
+    let u2 = (*seed >> 11) as f32 / (1u64 << 53) as f32;
+
+    // Box-Muller: guard log(0) by clamping u1 away from 0.
+    let u1 = u1.max(1e-10_f32);
+    (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos()
+}
+
+// ---------------------------------------------------------------------------
+// SwaSchedule
+// ---------------------------------------------------------------------------
+
+/// Cosine-annealing learning-rate schedule used by SWA.
+///
+/// Before `swa_start_step` the schedule holds a constant `lr_init`.  Afterwards
+/// it applies cosine annealing within each cycle of `cycle_length` steps, with
+/// the minimum value `lr_min` (defaults to `lr_init * 0.01`).
+///
+/// A *snapshot* (averaging event) is triggered at the end of every full cycle
+/// after `swa_start_step`.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_transformers::training::SwaSchedule;
+///
+/// let sched = SwaSchedule::new(1e-2, 1000, 5000);
+/// assert_eq!(sched.lr_at_step(0), 1e-2);
+/// assert!(sched.should_snapshot(5999));   // end of first cycle
+/// assert!(!sched.should_snapshot(5500));  // mid-cycle
+/// ```
+#[derive(Debug, Clone)]
+pub struct SwaSchedule {
+    /// Starting / peak learning rate for the cosine annealing cycle.
+    pub lr_init: f32,
+    /// Minimum (trough) learning rate.  Default: `lr_init * 0.01`.
+    pub lr_min: f32,
+    /// Number of steps per cosine annealing cycle.
+    pub cycle_length: usize,
+    /// Global step at which SWA averaging begins.
+    pub swa_start_step: usize,
+}
+
+impl SwaSchedule {
+    /// Construct a schedule with `lr_min = lr_init * 0.01`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `cycle_length == 0`.
+    pub fn new(lr_init: f32, cycle_length: usize, swa_start_step: usize) -> Self {
+        assert!(cycle_length > 0, "cycle_length must be > 0");
+        Self {
+            lr_init,
+            lr_min: lr_init * 0.01,
+            cycle_length,
+            swa_start_step,
+        }
+    }
+
+    /// Learning rate at a given global `step`.
+    ///
+    /// - `step < swa_start_step` → constant `lr_init`.
+    /// - `step >= swa_start_step` → cosine annealing:
+    ///
+    /// ```text
+    /// phase = (step - swa_start_step) % cycle_length
+    /// lr = lr_min + 0.5 * (lr_init - lr_min) * (1 + cos(π * phase / cycle_length))
+    /// ```
+    pub fn lr_at_step(&self, step: usize) -> f32 {
+        if step < self.swa_start_step {
+            return self.lr_init;
+        }
+        let phase = (step - self.swa_start_step) % self.cycle_length;
+        let t = phase as f32 / self.cycle_length as f32;
+        self.lr_min + 0.5 * (self.lr_init - self.lr_min) * (1.0 + (std::f32::consts::PI * t).cos())
+    }
+
+    /// Returns `true` when a snapshot should be taken (end of a cycle after SWA start).
+    ///
+    /// A snapshot is triggered on step `s` iff:
+    /// - `s >= swa_start_step`, and
+    /// - `(s - swa_start_step + 1) % cycle_length == 0`
+    pub fn should_snapshot(&self, step: usize) -> bool {
+        if step < self.swa_start_step {
+            return false;
+        }
+        (step - self.swa_start_step + 1) % self.cycle_length == 0
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SwaBuffer
+// ---------------------------------------------------------------------------
+
+/// Accumulates the arithmetic running average of parameter snapshots for SWA.
+///
+/// Each call to [`update`](SwaBuffer::update) incorporates one new parameter snapshot
+/// using the online mean formula:
+///
+/// ```text
+/// avg = avg + (new - avg) / (n_models + 1)
+/// ```
+///
+/// This is numerically stable and requires only one pass over the data.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_transformers::training::SwaBuffer;
+///
+/// let mut buf = SwaBuffer::new();
+/// buf.update(&[("w", &[1.0_f32, 3.0])]);
+/// buf.update(&[("w", &[3.0_f32, 5.0])]);
+/// let avg = buf.get("w").unwrap();
+/// assert!((avg[0] - 2.0).abs() < 1e-6);
+/// assert!((avg[1] - 4.0).abs() < 1e-6);
+/// ```
+#[derive(Debug, Clone)]
+pub struct SwaBuffer {
+    /// Running mean of each parameter tensor (name → flat `Vec<f32>`).
+    pub avg_params: HashMap<String, Vec<f32>>,
+    /// Number of snapshots averaged so far.
+    pub n_models: usize,
+}
+
+impl Default for SwaBuffer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SwaBuffer {
+    /// Create an empty buffer.
+    pub fn new() -> Self {
+        Self {
+            avg_params: HashMap::new(),
+            n_models: 0,
+        }
+    }
+
+    /// Incorporate a new parameter snapshot into the running mean.
+    ///
+    /// `params` is a slice of `(name, flat_data)` pairs.  All parameter names must
+    /// be consistent across calls; new names are initialised on first sight.
+    pub fn update(&mut self, params: &[(&str, &[f32])]) {
+        let n = self.n_models as f32;
+        let denom = n + 1.0;
+
+        for &(name, data) in params {
+            match self.avg_params.get_mut(name) {
+                Some(avg) => {
+                    // Online update: avg += (new - avg) / (n + 1)
+                    for (a, &x) in avg.iter_mut().zip(data.iter()) {
+                        *a += (x - *a) / denom;
+                    }
+                }
+                None => {
+                    // First snapshot for this parameter: initialise directly.
+                    self.avg_params.insert(name.to_string(), data.to_vec());
+                }
+            }
+        }
+
+        self.n_models += 1;
+    }
+
+    /// Return the current SWA average for `name`, or `None` if not yet seen.
+    pub fn get(&self, name: &str) -> Option<&[f32]> {
+        self.avg_params.get(name).map(Vec::as_slice)
+    }
+
+    /// All averaged parameter names in arbitrary order.
+    pub fn param_names(&self) -> Vec<&str> {
+        self.avg_params.keys().map(String::as_str).collect()
+    }
+
+    /// Reset: clear all averages and snapshot count.
+    pub fn reset(&mut self) {
+        self.avg_params.clear();
+        self.n_models = 0;
+    }
+
+    /// Number of distinct parameter tensors tracked.
+    pub fn num_params(&self) -> usize {
+        self.avg_params.len()
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SwagBuffer
+// ---------------------------------------------------------------------------
+
+/// Fits a diagonal + low-rank Gaussian posterior over the SWA trajectory (SWAG).
+///
+/// The posterior approximation is:
+///
+/// ```text
+/// θ ~ N(μ, Σ)
+/// Σ ≈ (1/2) Σ_diag + (1/(2(K-1))) D Dᵀ
+/// ```
+///
+/// where `μ` is the running mean, `Σ_diag = diag(E[θ²] - E[θ]²)` is the
+/// element-wise variance, `D` is the matrix of deviation columns
+/// `(θ_k - μ_k)` for the last `max_rank` snapshots, and `K` is the number
+/// of snapshots collected.
+///
+/// # Sampling
+///
+/// [`sample`](SwagBuffer::sample) draws a parameter vector as:
+///
+/// ```text
+/// θ_s = μ + (1/√2) Σ_diag^{1/2} z₁ + (1/√(2(K-1))) D z₂
+/// ```
+///
+/// where `z₁ ~ N(0, I_p)` and `z₂ ~ N(0, I_K)` are standard normal vectors.
+#[derive(Debug, Clone)]
+pub struct SwagBuffer {
+    /// Running mean E[θ] for each parameter.
+    pub mean: HashMap<String, Vec<f32>>,
+    /// Running second moment E[θ²] (element-wise) for each parameter.
+    pub sq_mean: HashMap<String, Vec<f32>>,
+    /// Number of snapshots collected.
+    pub n_models: usize,
+    /// Maximum number of deviation columns stored (low-rank rank cap).
+    pub max_rank: usize,
+    /// Deviation columns `θ_k - μ_k` at snapshot time, capped at `max_rank`.
+    ///
+    /// Stored as `name → Vec<Vec<f32>>` where the outer Vec is across snapshots.
+    pub deviations: HashMap<String, Vec<Vec<f32>>>,
+}
+
+impl SwagBuffer {
+    /// Create an empty SWAG buffer with the given low-rank cap.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `max_rank == 0`.
+    pub fn new(max_rank: usize) -> Self {
+        assert!(max_rank > 0, "max_rank must be > 0");
+        Self {
+            mean: HashMap::new(),
+            sq_mean: HashMap::new(),
+            n_models: 0,
+            max_rank,
+            deviations: HashMap::new(),
+        }
+    }
+
+    /// Incorporate a new snapshot: updates the running mean, second moment, and
+    /// deviation columns.
+    ///
+    /// The deviation column is computed as `snapshot - mean_before_update`.  We then
+    /// append it to the deviation list, dropping the oldest entry if we have reached
+    /// `max_rank`.
+    pub fn update(&mut self, params: &[(&str, &[f32])]) {
+        let n = self.n_models as f32;
+        let denom = n + 1.0;
+
+        for &(name, data) in params {
+            // ---- mean & sq_mean ----
+            match self.mean.get_mut(name) {
+                Some(mu) => {
+                    let sq = self.sq_mean.get_mut(name).expect("sq_mean desync");
+
+                    // Capture deviation BEFORE updating mean (θ_k - μ_k)
+                    let deviation: Vec<f32> = data.iter().zip(mu.iter()).map(|(&x, &m)| x - m).collect();
+
+                    // Online mean update
+                    for (m, &x) in mu.iter_mut().zip(data.iter()) {
+                        *m += (x - *m) / denom;
+                    }
+                    // Online second-moment update
+                    for (s, &x) in sq.iter_mut().zip(data.iter()) {
+                        *s += (x * x - *s) / denom;
+                    }
+
+                    // Append deviation column; evict oldest if at capacity
+                    let cols = self.deviations.entry(name.to_string()).or_default();
+                    if cols.len() >= self.max_rank {
+                        cols.remove(0);
+                    }
+                    cols.push(deviation);
+                }
+                None => {
+                    // First snapshot: mean = data, sq_mean = data², no deviation yet.
+                    self.mean.insert(name.to_string(), data.to_vec());
+                    self.sq_mean
+                        .insert(name.to_string(), data.iter().map(|&x| x * x).collect());
+                    self.deviations.insert(name.to_string(), Vec::new());
+                }
+            }
+        }
+
+        self.n_models += 1;
+    }
+
+    /// Element-wise variance estimate: `Var[θ] ≈ E[θ²] - E[θ]²` (diagonal component).
+    ///
+    /// Returns `None` if `name` has not been seen.  Values are clamped to `[0, ∞)` to
+    /// avoid floating-point artifacts producing tiny negatives.
+    pub fn diagonal_variance(&self, name: &str) -> Option<Vec<f32>> {
+        let mu = self.mean.get(name)?;
+        let sq = self.sq_mean.get(name)?;
+        Some(
+            mu.iter()
+                .zip(sq.iter())
+                .map(|(&m, &s)| (s - m * m).max(0.0))
+                .collect(),
+        )
+    }
+
+    /// Sample a parameter vector from the SWAG posterior for `name`.
+    ///
+    /// Returns `None` if `name` has not been seen or there are fewer than 2 snapshots.
+    ///
+    /// The sample is:
+    /// ```text
+    /// θ_s = μ + (1/√2) Σ_diag^{1/2} z₁ + (1/√(2(K-1))) D z₂
+    /// ```
+    ///
+    /// Uses an internal LCG seeded by `rng_seed` for fully deterministic output —
+    /// identical seeds produce identical samples.
+    pub fn sample(&self, name: &str, rng_seed: u64) -> Option<Vec<f32>> {
+        let mu = self.mean.get(name)?;
+        let var = self.diagonal_variance(name)?;
+        let cols = self.deviations.get(name)?;
+        let k = cols.len();
+
+        // Need at least one deviation column (requires >= 2 snapshots).
+        if k == 0 {
+            return None;
+        }
+
+        let p = mu.len();
+        let mut seed = rng_seed;
+
+        // z₁: diagonal contribution
+        let z1: Vec<f32> = (0..p).map(|_| lcg_normal(&mut seed)).collect();
+
+        // z₂: low-rank contribution (k-dimensional)
+        let z2: Vec<f32> = (0..k).map(|_| lcg_normal(&mut seed)).collect();
+
+        let diag_scale = (0.5_f32).sqrt();
+        let lr_scale = if k > 1 {
+            (1.0 / (2.0 * (k - 1) as f32)).sqrt()
+        } else {
+            0.0 // only one deviation — low-rank term contributes nothing meaningful
+        };
+
+        // θ_s = μ + diag_scale * √Var ⊙ z₁ + lr_scale * D z₂
+        let mut sample: Vec<f32> = mu
+            .iter()
+            .zip(var.iter())
+            .zip(z1.iter())
+            .map(|((&m, &v), &z)| m + diag_scale * v.sqrt() * z)
+            .collect();
+
+        // Low-rank term: add lr_scale * D z₂ column by column
+        for (j, col) in cols.iter().enumerate() {
+            let scale = lr_scale * z2[j];
+            for (s, &d) in sample.iter_mut().zip(col.iter()) {
+                *s += scale * d;
+            }
+        }
+
+        Some(sample)
+    }
+
+    /// Number of snapshots collected so far.
+    pub fn num_snapshots(&self) -> usize {
+        self.n_models
+    }
+}
+
+// ---------------------------------------------------------------------------
+// SwaTrainer
+// ---------------------------------------------------------------------------
+
+/// High-level SWA trainer: wraps a [`SwaSchedule`], [`SwaBuffer`], and an optional
+/// [`SwagBuffer`] and drives the snapshot logic automatically.
+///
+/// Call [`update`](SwaTrainer::update) after every optimizer step.  It advances the
+/// internal step counter, checks [`SwaSchedule::should_snapshot`], and, if a snapshot
+/// is due, forwards the parameters to the averaging buffers.
+///
+/// # Example
+///
+/// ```rust
+/// use rtx_transformers::training::{SwaSchedule, SwaTrainer};
+///
+/// let schedule = SwaSchedule::new(1e-2, 100, 500);
+/// let mut trainer = SwaTrainer::new(schedule);
+///
+/// for step in 0..700 {
+///     let params = [("w", &[1.0_f32, 2.0][..])];
+///     trainer.update(&params);
+/// }
+/// assert!(trainer.n_snapshots() > 0);
+/// ```
+pub struct SwaTrainer {
+    /// Learning-rate schedule and snapshot timing.
+    pub schedule: SwaSchedule,
+    /// Arithmetic average accumulator.
+    pub swa_buffer: SwaBuffer,
+    /// Optional SWAG covariance accumulator.
+    pub swag_buffer: Option<SwagBuffer>,
+    /// Internal step counter (0-indexed, incremented inside `update`).
+    step: usize,
+}
+
+impl SwaTrainer {
+    /// Create a trainer with SWA averaging only (no SWAG).
+    pub fn new(schedule: SwaSchedule) -> Self {
+        Self {
+            schedule,
+            swa_buffer: SwaBuffer::new(),
+            swag_buffer: None,
+            step: 0,
+        }
+    }
+
+    /// Create a trainer with both SWA averaging and SWAG posterior estimation.
+    pub fn with_swag(schedule: SwaSchedule, max_rank: usize) -> Self {
+        Self {
+            schedule,
+            swa_buffer: SwaBuffer::new(),
+            swag_buffer: Some(SwagBuffer::new(max_rank)),
+            step: 0,
+        }
+    }
+
+    /// Process one training step.
+    ///
+    /// Increments the step counter and, if the current step is a snapshot boundary,
+    /// passes `params` to the SWA and (if configured) SWAG buffers.
+    ///
+    /// Returns `true` if a snapshot was taken on this step.
+    pub fn update(&mut self, params: &[(&str, &[f32])]) -> bool {
+        let snapshot = self.schedule.should_snapshot(self.step);
+        if snapshot {
+            self.swa_buffer.update(params);
+            if let Some(ref mut swag) = self.swag_buffer {
+                swag.update(params);
+            }
+        }
+        self.step += 1;
+        snapshot
+    }
+
+    /// Current global step (number of `update` calls completed so far, pre-increment).
+    ///
+    /// After the first `update` call returns, `current_step()` returns 1.
+    pub fn current_step(&self) -> usize {
+        self.step
+    }
+
+    /// Learning rate that should be used at the *next* step (i.e., for `self.step`).
+    ///
+    /// Call this before `update` to get the LR to pass to your optimizer.
+    pub fn current_lr(&self) -> f32 {
+        self.schedule.lr_at_step(self.step)
+    }
+
+    /// Total number of snapshots (averaging events) performed so far.
+    pub fn n_snapshots(&self) -> usize {
+        self.swa_buffer.n_models
+    }
+
+    /// SWA-averaged parameters for `name`, or `None` if no snapshots yet.
+    pub fn get_swa_params(&self, name: &str) -> Option<&[f32]> {
+        self.swa_buffer.get(name)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ------------------------------------------------------------------
+    // SwaSchedule
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_schedule_lr_before_swa_start() {
+        let s = SwaSchedule::new(0.1, 1000, 5000);
+        // Any step before swa_start_step should return lr_init exactly.
+        assert_eq!(s.lr_at_step(0), 0.1);
+        assert_eq!(s.lr_at_step(4999), 0.1);
+    }
+
+    #[test]
+    fn test_schedule_lr_at_swa_start() {
+        // At exactly swa_start_step, phase == 0 → cos(0) == 1 → lr == lr_init.
+        let s = SwaSchedule::new(0.1, 1000, 5000);
+        let lr = s.lr_at_step(5000);
+        assert!(
+            (lr - 0.1).abs() < 1e-6,
+            "lr at swa_start_step should equal lr_init, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_schedule_lr_cosine_min() {
+        // At the very end of a cycle (phase == cycle_length - 1), lr ≈ lr_min.
+        // More precisely: phase = cycle_length - 1, t = (cycle_length-1)/cycle_length → ~1
+        // cos(π * t) → cos(π * (1 - 1/L)) which for large L → cos(π) = -1 → lr_min.
+        // For a cycle_length of 4, at phase=3: t=3/4, cos(3π/4) = -√2/2 ≈ -0.707
+        // We test the final step (phase = cycle_length - 1) for a large cycle to see lr near lr_min.
+        let lr_init = 0.1_f32;
+        let lr_min = lr_init * 0.01;
+        let cycle = 10_000_usize;
+        let start = 0_usize;
+        let s = SwaSchedule { lr_init, lr_min, cycle_length: cycle, swa_start_step: start };
+        // step = cycle - 1 → phase = cycle - 1, t ≈ 1 − 1/cycle ≈ 1
+        let lr = s.lr_at_step(cycle - 1);
+        // Should be close to lr_min but not exactly (phase = cycle-1, not cycle)
+        // At phase=cycle-1: t=(cycle-1)/cycle, cos(π*t)=cos(π-π/cycle)≈-1+small
+        // lr ≈ lr_min + 0.5*(lr_init-lr_min)*(1+(-1+ε)) ≈ lr_min + tiny
+        assert!(
+            lr < lr_init * 0.02,
+            "lr near end of cycle should be close to lr_min ({lr_min}), got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_schedule_lr_cosine_max() {
+        // At the start of a new cycle (phase == 0), lr == lr_init.
+        let s = SwaSchedule::new(0.05, 200, 1000);
+        // Second cycle starts at step = swa_start_step + cycle_length = 1200
+        let lr = s.lr_at_step(1200);
+        assert!(
+            (lr - 0.05).abs() < 1e-6,
+            "lr at cycle start should equal lr_init, got {lr}"
+        );
+    }
+
+    #[test]
+    fn test_schedule_should_snapshot_false_before_start() {
+        let s = SwaSchedule::new(0.1, 100, 500);
+        for step in 0..500 {
+            assert!(
+                !s.should_snapshot(step),
+                "should not snapshot before swa_start_step at step {step}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_schedule_should_snapshot_at_cycle_end() {
+        // First cycle ends at swa_start_step + cycle_length - 1.
+        let s = SwaSchedule::new(0.1, 100, 500);
+        let first_snapshot_step = 500 + 100 - 1; // = 599
+        assert!(
+            s.should_snapshot(first_snapshot_step),
+            "should snapshot at end of first cycle (step {first_snapshot_step})"
+        );
+        // Second cycle end
+        let second_snapshot_step = 500 + 200 - 1; // = 699
+        assert!(
+            s.should_snapshot(second_snapshot_step),
+            "should snapshot at end of second cycle (step {second_snapshot_step})"
+        );
+    }
+
+    #[test]
+    fn test_schedule_should_snapshot_mid_cycle() {
+        let s = SwaSchedule::new(0.1, 100, 500);
+        // Mid-cycle: step 550 → (550 - 500 + 1) = 51 % 100 == 51 ≠ 0
+        assert!(!s.should_snapshot(550), "should not snapshot mid-cycle");
+        assert!(!s.should_snapshot(501));
+        assert!(!s.should_snapshot(598));
+    }
+
+    // ------------------------------------------------------------------
+    // SwaBuffer
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_swa_buffer_new_empty() {
+        let buf = SwaBuffer::new();
+        assert_eq!(buf.n_models, 0);
+        assert_eq!(buf.num_params(), 0);
+        assert!(buf.get("w").is_none());
+    }
+
+    #[test]
+    fn test_swa_buffer_first_update() {
+        // After one update the average must equal the snapshot exactly.
+        let mut buf = SwaBuffer::new();
+        buf.update(&[("w", &[1.0_f32, 2.0, 3.0])]);
+        let avg = buf.get("w").unwrap();
+        assert_eq!(avg, &[1.0_f32, 2.0, 3.0]);
+        assert_eq!(buf.n_models, 1);
+    }
+
+    #[test]
+    fn test_swa_buffer_second_update() {
+        // After two updates, avg == (snap1 + snap2) / 2.
+        let mut buf = SwaBuffer::new();
+        buf.update(&[("w", &[1.0_f32, 3.0])]);
+        buf.update(&[("w", &[3.0_f32, 5.0])]);
+        let avg = buf.get("w").unwrap();
+        assert!(
+            (avg[0] - 2.0).abs() < 1e-6,
+            "expected 2.0 got {}",
+            avg[0]
+        );
+        assert!(
+            (avg[1] - 4.0).abs() < 1e-6,
+            "expected 4.0 got {}",
+            avg[1]
+        );
+    }
+
+    #[test]
+    fn test_swa_buffer_incremental_correctness() {
+        // Four updates: verify the mean matches a manual reference.
+        let snaps: &[&[f32]] = &[
+            &[0.0, 4.0],
+            &[2.0, 2.0],
+            &[4.0, 0.0],
+            &[6.0, 6.0],
+        ];
+        let mut buf = SwaBuffer::new();
+        for snap in snaps {
+            buf.update(&[("w", snap)]);
+        }
+        let expected = [3.0_f32, 3.0_f32]; // (0+2+4+6)/4, (4+2+0+6)/4
+        let avg = buf.get("w").unwrap();
+        for (a, e) in avg.iter().zip(expected.iter()) {
+            assert!(
+                (a - e).abs() < 1e-5,
+                "expected {e} got {a}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_swa_buffer_multiple_params() {
+        let mut buf = SwaBuffer::new();
+        buf.update(&[("a", &[1.0_f32, 2.0]), ("b", &[10.0_f32])]);
+        buf.update(&[("a", &[3.0_f32, 4.0]), ("b", &[20.0_f32])]);
+        let a = buf.get("a").unwrap();
+        let b = buf.get("b").unwrap();
+        assert!((a[0] - 2.0).abs() < 1e-6, "a[0]={}", a[0]);
+        assert!((a[1] - 3.0).abs() < 1e-6, "a[1]={}", a[1]);
+        assert!((b[0] - 15.0).abs() < 1e-6, "b[0]={}", b[0]);
+    }
+
+    #[test]
+    fn test_swa_buffer_reset() {
+        let mut buf = SwaBuffer::new();
+        buf.update(&[("w", &[1.0_f32])]);
+        buf.update(&[("w", &[2.0_f32])]);
+        assert_eq!(buf.n_models, 2);
+
+        buf.reset();
+
+        assert_eq!(buf.n_models, 0);
+        assert!(buf.avg_params.is_empty());
+        assert!(buf.get("w").is_none());
+    }
+
+    // ------------------------------------------------------------------
+    // SwagBuffer
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_swag_buffer_diagonal_variance_zero() {
+        // After two identical snapshots, variance should be 0 (or near 0).
+        let mut swag = SwagBuffer::new(10);
+        swag.update(&[("w", &[1.0_f32, 2.0, 3.0])]);
+        swag.update(&[("w", &[1.0_f32, 2.0, 3.0])]);
+        let var = swag.diagonal_variance("w").unwrap();
+        for (i, &v) in var.iter().enumerate() {
+            assert!(
+                v.abs() < 1e-5,
+                "variance[{i}] should be ~0 for identical snapshots, got {v}"
+            );
+        }
+    }
+
+    #[test]
+    fn test_swag_buffer_diagonal_variance_nonzero() {
+        // Different snapshots should produce positive variance.
+        let mut swag = SwagBuffer::new(10);
+        swag.update(&[("w", &[0.0_f32, 0.0])]);
+        swag.update(&[("w", &[2.0_f32, 4.0])]);
+        let var = swag.diagonal_variance("w").unwrap();
+        assert!(var[0] > 0.0, "var[0] should be > 0, got {}", var[0]);
+        assert!(var[1] > 0.0, "var[1] should be > 0, got {}", var[1]);
+    }
+
+    #[test]
+    fn test_swag_buffer_sample_shape() {
+        // The sample should have the same length as the parameter vector.
+        let mut swag = SwagBuffer::new(5);
+        swag.update(&[("w", &[1.0_f32, 2.0, 3.0, 4.0])]);
+        swag.update(&[("w", &[2.0_f32, 3.0, 4.0, 5.0])]);
+        let sample = swag.sample("w", 42).unwrap();
+        assert_eq!(sample.len(), 4, "sample length should match parameter size");
+    }
+
+    #[test]
+    fn test_swag_buffer_sample_differs_per_seed() {
+        let mut swag = SwagBuffer::new(5);
+        swag.update(&[("w", &[0.0_f32, 0.0, 0.0])]);
+        swag.update(&[("w", &[1.0_f32, 1.0, 1.0])]);
+        let s1 = swag.sample("w", 1).unwrap();
+        let s2 = swag.sample("w", 99999).unwrap();
+        // Different seeds should (almost certainly) produce different samples.
+        let same = s1.iter().zip(s2.iter()).all(|(a, b)| (a - b).abs() < 1e-9);
+        assert!(!same, "different seeds should produce different samples");
+    }
+
+    #[test]
+    fn test_swag_buffer_sample_none_before_snapshots() {
+        // With only one snapshot there are no deviation columns, so sample returns None.
+        let mut swag = SwagBuffer::new(5);
+        swag.update(&[("w", &[1.0_f32, 2.0])]);
+        assert!(
+            swag.sample("w", 0).is_none(),
+            "sample should return None with only one snapshot (no deviation columns)"
+        );
+    }
+
+    #[test]
+    fn test_swag_buffer_max_rank_eviction() {
+        // Deviation columns should be capped at max_rank.
+        let mut swag = SwagBuffer::new(3);
+        for i in 0..10 {
+            swag.update(&[("w", &[i as f32])]);
+        }
+        let cols = swag.deviations.get("w").unwrap();
+        assert!(
+            cols.len() <= 3,
+            "deviation columns should be capped at max_rank=3, got {}",
+            cols.len()
+        );
+    }
+
+    #[test]
+    fn test_swag_buffer_num_snapshots() {
+        let mut swag = SwagBuffer::new(5);
+        assert_eq!(swag.num_snapshots(), 0);
+        swag.update(&[("w", &[1.0_f32])]);
+        assert_eq!(swag.num_snapshots(), 1);
+        swag.update(&[("w", &[2.0_f32])]);
+        assert_eq!(swag.num_snapshots(), 2);
+    }
+
+    // ------------------------------------------------------------------
+    // SwaTrainer
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_swa_trainer_snapshots_at_cycle_end() {
+        // schedule: lr_init=0.01, cycle=100, start=500
+        // First snapshot at step 599 (0-indexed).
+        let s = SwaSchedule::new(0.01, 100, 500);
+        let mut trainer = SwaTrainer::new(s);
+        let params = [("w", &[1.0_f32][..])];
+
+        let mut snapshot_steps = Vec::new();
+        for step in 0..700 {
+            let took = trainer.update(&params);
+            if took {
+                snapshot_steps.push(step);
+            }
+        }
+
+        // Snapshots should occur at steps 599 and 699.
+        assert_eq!(snapshot_steps, vec![599, 699]);
+    }
+
+    #[test]
+    fn test_swa_trainer_lr_schedule() {
+        let s = SwaSchedule::new(0.1, 1000, 5000);
+        let trainer = SwaTrainer::new(s);
+        // Before SWA starts: current_lr at step 0 should be lr_init.
+        assert_eq!(trainer.current_lr(), 0.1);
+    }
+
+    #[test]
+    fn test_swa_trainer_n_snapshots_counts() {
+        let s = SwaSchedule::new(0.01, 10, 50);
+        let mut trainer = SwaTrainer::new(s);
+        let params = [("w", &[1.0_f32][..])];
+
+        for _ in 0..80 {
+            trainer.update(&params);
+        }
+        // Snapshots at steps 59 and 69 and 79 → 3 snapshots
+        assert_eq!(trainer.n_snapshots(), 3);
+    }
+
+    #[test]
+    fn test_swa_trainer_get_params_after_snapshot() {
+        let s = SwaSchedule::new(0.01, 10, 0);
+        let mut trainer = SwaTrainer::new(s);
+        // No snapshots yet
+        assert!(trainer.get_swa_params("w").is_none());
+        let params = [("w", &[5.0_f32, 6.0][..])];
+        // First snapshot at step 9.
+        for _ in 0..10 {
+            trainer.update(&params);
+        }
+        let avg = trainer.get_swa_params("w").unwrap();
+        assert_eq!(avg, &[5.0_f32, 6.0]);
+    }
+
+    #[test]
+    fn test_swa_trainer_with_swag() {
+        let s = SwaSchedule::new(0.01, 10, 0);
+        let trainer = SwaTrainer::with_swag(s, 5);
+        assert!(
+            trainer.swag_buffer.is_some(),
+            "swag_buffer should be Some after with_swag"
+        );
+    }
+
+    #[test]
+    fn test_swa_trainer_step_increments() {
+        let s = SwaSchedule::new(0.01, 100, 500);
+        let mut trainer = SwaTrainer::new(s);
+        assert_eq!(trainer.current_step(), 0);
+        trainer.update(&[("w", &[1.0_f32])]);
+        assert_eq!(trainer.current_step(), 1);
+        trainer.update(&[("w", &[1.0_f32])]);
+        assert_eq!(trainer.current_step(), 2);
+    }
+
+    #[test]
+    fn test_swa_trainer_zero_snapshots_initially() {
+        let s = SwaSchedule::new(0.01, 100, 500);
+        let trainer = SwaTrainer::new(s);
+        assert_eq!(trainer.n_snapshots(), 0);
+    }
+
+    // ------------------------------------------------------------------
+    // LCG sampler sanity check
+    // ------------------------------------------------------------------
+
+    #[test]
+    fn test_lcg_normal_deterministic() {
+        // Same seed must produce the same value.
+        let mut s1 = 12345_u64;
+        let mut s2 = 12345_u64;
+        let v1 = lcg_normal(&mut s1);
+        let v2 = lcg_normal(&mut s2);
+        assert_eq!(v1, v2, "lcg_normal must be deterministic for the same seed");
+    }
+
+    #[test]
+    fn test_lcg_normal_differs_per_seed() {
+        let mut s1 = 1_u64;
+        let mut s2 = 9999_u64;
+        let v1 = lcg_normal(&mut s1);
+        let v2 = lcg_normal(&mut s2);
+        assert_ne!(v1, v2, "different seeds should produce different values");
+    }
+}