rtx-csm: decoder activation capture — architectural hypothesis validated

Adds capture_decoder_activations on Model and ModelBackend, plus
--target-module decoder|backbone on examples/steering_extract.

Decoder mode runs a text-only prompt through the backbone (no
capture), Mimi-encodes the audio separately to grab the middle
frame's c0 token, teacher-forces that c0, and captures one mean-
pooled-over-seq vector per decoder layer. Result: 4 layers ×
1024 embed dim per call, much faster than backbone capture
(text-only prompts are short).

A/B with the canonical "Today I want to share..." prompt at seed
7 (the previously-identified low-WER seed):

  case          cos    WER    transcript
  baseline      0.76   0.36   "And today I want to share something some
                              funnel distraits"  (high baseline at this
                                                  seed)
  [email protected]  0.75   5.00   "Today, I want to share some needs of my
                              prey"  (backbone destroys content)
  [email protected]   0.86   0.93   "I'm not that tall. I'm not that tall."
                              (fluent but repetitive — biggest cos)
  [email protected]   0.61   2.57   over-steered
  [email protected]   0.61   2.14   broken

[email protected] is the largest speaker_cosine boost we've measured AND
produces clean English. Backbone steering at the same seed destroyed
content fidelity. This validates the architectural hypothesis: the
backbone carries semantic content (what the model says), the depth
decoder carries acoustic detail (how it sounds). Steering the
decoder shifts voice character without disturbing word content the
way backbone steering does.

Open issues: [email protected] produces repetitive output ("I'm not that
tall" three times). Likely lower scale (~0.5) plus the existing
repetition guard would fix it; left for follow-up.

Decoder vector magnitudes are ~10× smaller than backbone (norm 0.85
at deepest layer vs 14.9), so the appropriate scale is ~10× higher
than the backbone recipe (1.0 vs 0.1-0.3).

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
2026-04-29 12:59:28 -07:00
co-authored by Claude Opus 4.7
parent a1fa72d151
commit 99a7e53aa8
3 changed files with 123 additions and 26 deletions
+40
View File
@@ -840,6 +840,46 @@ impl Model {
Ok(captured)
}
/// Teacher-forced single-frame capture at the depth decoder. Builds
/// the same prompt as `capture_backbone_activations`, runs the
/// backbone (no capture), then teacher-forces the decoder with the
/// ground-truth c0 and captures one mean-pooled-over-seq activation
/// per decoder layer (4 vectors @ 1024 dim for CSM-1B).
///
/// `target_c0` should be the Mimi-encoded codebook-0 token for the
/// audio frame the model would predict next given the prompt — i.e.
/// the FIRST audio token AFTER the prompt's audio prefix. The caller
/// is responsible for slicing the manifest's frame_codes accordingly.
pub fn capture_decoder_activations(
&mut self,
tokens: &Tensor,
tokens_mask: &Tensor,
target_c0: u32,
) -> Result<Vec<Tensor>> {
let saved_b_steering = self.backbone.steering.take();
let saved_d_steering = self.decoder.steering.take();
self.backbone.clear_kv_cache();
self.decoder.clear_kv_cache();
let embeds = self.build_embeds(tokens, tokens_mask)?;
let h = self.backbone.forward(&embeds, 0)?;
let c0_t = Tensor::from_slice(&[target_c0], (1, 1), &self.decoder.device)?;
let c0_embed = self.audio_embeddings.forward(&c0_t)?;
let curr_h = Tensor::cat(&[h, c0_embed], 1)?;
let proj_h = curr_h.apply(&self.projection)?;
self.decoder.start_capture();
let _decoder_h = self.decoder.forward(&proj_h, 0)?;
let captured = self
.decoder
.take_capture()
.ok_or_else(|| crate::error::CsmError::Config("decoder capture empty".into()))?;
self.backbone.steering = saved_b_steering;
self.decoder.steering = saved_d_steering;
Ok(captured)
}
/// Build the per-frame embedding tensor `(B, S, D)` from packed token slots.
/// Shared by `generate_frame` and `generate_frame_cfg`.
fn build_embeds(&self, tokens: &Tensor, tokens_mask: &Tensor) -> Result<Tensor> {