rtx-interpret: seeded SAE init + coupled per-unit optimizer reset

SparseAutoencoder::new_seeded draws encoder/decoder weights via
randn_seeded with per-tensor SplitMix64-derived seeds (same
derivation as MambaBlock::new_seeded), so identical (config, seed)
gives bit-exact SAEs — without it, cross-instance loss comparisons
are noise (measured downstream: 0.09 vs 0.65 starts on identical
data).

SAETrainer::reinitialize_neuron couples the encoder-unit weight
reinit with zeroing that unit's optimizer moment rows (encoder row,
bias slot, decoder column), so external generate-and-test callers
can't reset weights while leaving optimizer state stale — previously
only the trainer's internal dead-neuron resampling did both. Note
train_step's update rule is plain SGD today, so the moment reset is
inert until the Adam path is switched on; the coupling is the
contract either way, and a doctored-checkpoint test pins the
row/column semantics.

Also drops a vacuous assert!(true) smoke test that failed clippy's
assertions_on_constants.

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01B1feFAQxjbCRHePUdxuNra

crates/core/rtx-interpret/src/lib.rs

@@ -82,12 +82,6 @@ pub use types::{AttributionMetadata, AttributionOutput, Baseline, PerturbationOu
 mod tests {
     use super::*;
 
-    #[test]
-    fn test_crate_compiles() {
-        // Smoke test to ensure crate structure is sound
-        assert!(true);
-    }
-
     #[test]
     fn test_error_types_accessible() {
         let err = InterpretError::tensor("test");

crates/core/rtx-interpret/src/sae/mod.rs

@@ -213,13 +213,48 @@ pub struct SparseAutoencoder {
 impl SparseAutoencoder {
     /// Create a new Sparse Autoencoder
     pub fn new(config: SAEConfig) -> Result<Self> {
+        Self::build(config, None)
+    }
+
+    /// Create a new Sparse Autoencoder with a deterministic seed.
+    ///
+    /// Unlike [`Self::new`], the encoder and decoder weights are drawn
+    /// via `Tensor::randn_seeded` with per-tensor seeds derived from
+    /// `seed` by SplitMix64 (the same derivation as
+    /// `MambaBlock::new_seeded`), so two calls with identical
+    /// `(config, seed)` produce bit-exact identical SAEs. Required for
+    /// reproducible training runs and any cross-instance loss
+    /// comparison — with unseeded init, two SAEs on identical data can
+    /// start at wildly different losses.
+    pub fn new_seeded(config: SAEConfig, seed: u64) -> Result<Self> {
+        Self::build(config, Some(seed))
+    }
+
+    /// Shared constructor for [`Self::new`] (random) and
+    /// [`Self::new_seeded`] (`Some(seed)` ⇒ deterministic).
+    fn build(config: SAEConfig, seed: Option<u64>) -> Result<Self> {
         let d_model = config.d_model;
         let d_sae = config.d_sae();
         let device = &config.device;
 
+        // Per-tensor seed derivation by SplitMix64 so one operator seed
+        // maps stably onto the internal tensors (mirrors MambaBlock).
+        let mut state = seed.unwrap_or(0);
+        let mut rand_t = |shape: &[usize]| -> std::result::Result<Tensor, rtx_tensor::TensorError> {
+            state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
+            let mut z = state;
+            z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
+            z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
+            let s = z ^ (z >> 31);
+            match seed {
+                Some(_) => Tensor::randn_seeded(shape, device, s),
+                None => Tensor::randn(shape, device),
+            }
+        };
+
         // Initialize encoder weights (He initialization scaled)
         let encoder_std = config.encoder_init_scale * (2.0 / d_model as f32).sqrt();
-        let encoder = Tensor::randn(&[d_sae, d_model], device)?.mul_scalar(encoder_std)?;
+        let encoder = rand_t(&[d_sae, d_model])?.mul_scalar(encoder_std)?;
 
         let encoder_bias = Tensor::zeros([d_sae], device)?;
 
@@ -228,7 +263,7 @@ impl SparseAutoencoder {
         let decoder = if config.tied_weights {
             encoder.transpose(0, 1)?
         } else {
-            Tensor::randn(&[d_model, d_sae], device)?.mul_scalar(decoder_std)?
+            rand_t(&[d_model, d_sae])?.mul_scalar(decoder_std)?
         };
 
         let decoder_bias = Tensor::zeros([d_model], device)?;
@@ -776,6 +811,33 @@ mod tests {
         assert_eq!(sae.decoder().shape(), &[64, 256]);
     }
 
+    #[test]
+    fn test_sae_new_seeded_deterministic() {
+        let mk = |seed| {
+            let config = SAEConfig::new(16).with_expansion(2);
+            SparseAutoencoder::new_seeded(config, seed).unwrap()
+        };
+        let a = mk(42);
+        let b = mk(42);
+        assert_eq!(
+            a.encoder().to_cpu().unwrap(),
+            b.encoder().to_cpu().unwrap(),
+            "same seed must give bit-identical encoders"
+        );
+        assert_eq!(
+            a.decoder().to_cpu().unwrap(),
+            b.decoder().to_cpu().unwrap(),
+            "same seed must give bit-identical decoders"
+        );
+
+        let c = mk(43);
+        assert_ne!(
+            a.encoder().to_cpu().unwrap(),
+            c.encoder().to_cpu().unwrap(),
+            "different seeds must give different encoders"
+        );
+    }
+
     #[test]
     fn test_sae_forward() {
         let config = SAEConfig::new(32).with_expansion(2);

crates/core/rtx-interpret/src/sae/training.rs

@@ -629,6 +629,26 @@ impl SAETrainer {
         Ok(())
     }
 
+    /// Reinitialize one encoder unit AND reset its optimizer state, as
+    /// one operation.
+    ///
+    /// The coupled form of
+    /// [`SparseAutoencoder::reinitialize_encoder_neuron`]: callers
+    /// driving their own generate-and-test unit replacement (rather
+    /// than the trainer's internal dead-neuron resampling, which
+    /// already couples the two) must use this instead of reaching
+    /// through [`Self::sae_mut`], so the reset unit's optimizer
+    /// first/second-moment rows (encoder row, encoder-bias slot,
+    /// decoder column) are zeroed together with the weights. Note the
+    /// current `train_step` update rule is plain SGD — the moment
+    /// reset becomes load-bearing when the Adam path is switched on,
+    /// and the coupling is the contract either way.
+    pub fn reinitialize_neuron(&mut self, neuron_idx: usize, new_weights: &[f32]) -> Result<()> {
+        self.sae
+            .reinitialize_encoder_neuron(neuron_idx, new_weights)?;
+        self.reset_optimizer_states_for_neurons(&[neuron_idx])
+    }
+
     /// Get the SAE
     pub fn sae(&self) -> &SparseAutoencoder {
         &self.sae
@@ -944,6 +964,73 @@ mod tests {
         assert!((post_warmup_lr - 1e-3).abs() < 1e-5);
     }
 
+    #[test]
+    fn test_reinitialize_neuron_zeros_optimizer_rows() {
+        let d_model = 8usize;
+        let sae_config = SAEConfig::new(d_model).with_expansion(2);
+        let sae = SparseAutoencoder::new(sae_config).unwrap();
+        let mut trainer = SAETrainer::new(sae, SAETrainerConfig::default());
+        let d_sae = trainer.sae().d_sae();
+
+        let device = Device::cpu();
+        let x = Tensor::randn(&[4, d_model], &device).unwrap();
+        trainer.train_step(&x).unwrap(); // allocates optimizer state
+
+        // Doctor the checkpoint's optimizer moments to all-ones so the
+        // per-unit reset is observable (train_step's SGD path leaves
+        // them zero).
+        let path =
+            std::env::temp_dir().join(format!("rtx_sae_reinit_test_{}.json", std::process::id()));
+        trainer.save_checkpoint(&path).unwrap();
+        let raw = std::fs::read_to_string(&path).unwrap();
+        let mut ckpt: SAECheckpoint = serde_json::from_str(&raw).unwrap();
+        for state in [
+            ckpt.encoder_adam.as_mut(),
+            ckpt.encoder_bias_adam.as_mut(),
+            ckpt.decoder_adam.as_mut(),
+        ]
+        .into_iter()
+        .flatten()
+        {
+            state.m.iter_mut().for_each(|v| *v = 1.0);
+            state.v.iter_mut().for_each(|v| *v = 1.0);
+        }
+        std::fs::write(&path, serde_json::to_string(&ckpt).unwrap()).unwrap();
+        trainer.load_checkpoint(&path).unwrap();
+
+        let neuron = 3usize;
+        let new_w = vec![0.5f32; d_model];
+        trainer.reinitialize_neuron(neuron, &new_w).unwrap();
+
+        trainer.save_checkpoint(&path).unwrap();
+        let raw = std::fs::read_to_string(&path).unwrap();
+        let ckpt: SAECheckpoint = serde_json::from_str(&raw).unwrap();
+        std::fs::remove_file(&path).ok();
+
+        let enc = ckpt.encoder_adam.unwrap();
+        for i in 0..d_model {
+            assert_eq!(enc.m[neuron * d_model + i], 0.0);
+            assert_eq!(enc.v[neuron * d_model + i], 0.0);
+        }
+        assert_eq!(enc.m[0], 1.0, "other units' encoder rows untouched");
+
+        let encb = ckpt.encoder_bias_adam.unwrap();
+        assert_eq!(encb.m[neuron], 0.0);
+        assert_eq!(encb.v[neuron], 0.0);
+        assert_eq!(encb.m[0], 1.0, "other units' bias slots untouched");
+
+        let dec = ckpt.decoder_adam.unwrap();
+        for row in 0..d_model {
+            assert_eq!(dec.m[row * d_sae + neuron], 0.0);
+            assert_eq!(dec.v[row * d_sae + neuron], 0.0);
+        }
+        assert_eq!(dec.m[0], 1.0, "other units' decoder columns untouched");
+
+        // The weight reinit itself went through.
+        let enc_w = trainer.sae().encoder().to_cpu().unwrap();
+        assert_eq!(&enc_w[neuron * d_model..(neuron + 1) * d_model], &new_w[..]);
+    }
+
     #[test]
     fn test_activation_dataset() {
         let device = Device::cpu();