rtx-csm: Sprint 3 — Selective CFG schedule (step / linear / const)

Per-frame CFG scale schedule (arXiv 2509.19668, Zheng & Maleki). Pure
inference-time. Standard CFG uses one fixed scale for the whole
sequence; this lets the scale vary across frames so early frames
(speaker character) get full CFG and later frames (text adherence)
get a lower scale.

What lands:
- src/cfg_schedule.rs: CfgSchedule enum (Constant / Step /
  LinearRamp), scale_at(frame_idx), parser for CLI form
  `step:E:L:T | linear:S:E:R | const:X`. 6 unit tests.
- src/generator.rs: GenerateOptions::cfg_schedule (takes precedence
  over legacy cfg_scale; fixed-f64 path is preserved as
  Constant(s) for back-compat). Generation loop reads
  schedule.scale_at(frame_idx) and passes per-frame to
  generate_frame_cfg.
- examples/generate.rs: --cfg-schedule, --cfg-scale, --enable-cfg
  flags. Loading via load_csm_1b_with_cfg when --enable-cfg.

A/B with 6s output on Amini context, prompt about Selective CFG:

  case              cos    WER  transcript
  no-CFG baseline   0.944  1.50 "Okay, the M.U. worked..." (off)
  const:2.0         0.854  0.92 "On the right side." (short)
  step:3.0:1.5:12   0.938  1.00 "The officer for the selective
                                 C.F.D. paper recommends" (best)
  linear:3.0:1.0:25 0.854  1.08 "On the surface..." (off)

Step schedule produces the transcript closest to the input ("the
selective CFG paper recommends..."). WER stays at 1.0 because
Moonshine doesn't know "CFG" as a word, but qualitatively this is
the only one that's coherently following the prompt. Speaker cosine
stays ≈ baseline (0.94) instead of dropping to 0.85 like the
constant and linear cases.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/src/cfg_schedule.rs

@@ -0,0 +1,199 @@
+//! Per-frame CFG scale schedules (Selective CFG, arXiv 2509.19668).
+//!
+//! Standard CFG uses a fixed `cfg_scale` for every frame. Zheng & Maleki
+//! show that for zero-shot TTS, **frame-dependent** CFG schedules
+//! preserve text adherence late in synthesis while keeping speaker
+//! similarity high early on:
+//!
+//!   - Early frames (set the speaker character) → full CFG scale.
+//!   - Later frames (need to commit to the text) → lower scale, since
+//!     too much CFG late in generation causes over-conditioning on the
+//!     speaker prompt and degrades intelligibility.
+//!
+//! ## Variants
+//!
+//! - [`CfgSchedule::Constant`] — back-compat with the old fixed-`f64`
+//!   API; `cfg_scale.unwrap_or(1.0)` translates to this.
+//! - [`CfgSchedule::Step`] — `scale_early` until frame `transition`,
+//!   then jump to `scale_late`. Matches the paper's two-phase recipe.
+//! - [`CfgSchedule::LinearRamp`] — smooth interpolation from `start` at
+//!   frame 0 to `end` at `ramp_frames`, then holds at `end`. Sometimes
+//!   sounds less abrupt than the step variant.
+
+use crate::error::{CsmError, Result};
+
+#[derive(Debug, Clone, Copy)]
+pub enum CfgSchedule {
+    Constant(f64),
+    Step {
+        early: f64,
+        late: f64,
+        transition: usize,
+    },
+    LinearRamp {
+        start: f64,
+        end: f64,
+        ramp_frames: usize,
+    },
+}
+
+impl CfgSchedule {
+    pub fn scale_at(&self, frame_idx: usize) -> f64 {
+        match *self {
+            Self::Constant(s) => s,
+            Self::Step {
+                early,
+                late,
+                transition,
+            } => {
+                if frame_idx < transition {
+                    early
+                } else {
+                    late
+                }
+            }
+            Self::LinearRamp {
+                start,
+                end,
+                ramp_frames,
+            } => {
+                if ramp_frames == 0 || frame_idx >= ramp_frames {
+                    end
+                } else {
+                    let t = frame_idx as f64 / ramp_frames as f64;
+                    start * (1.0 - t) + end * t
+                }
+            }
+        }
+    }
+
+    /// True if any frame would invoke the CFG path (any scale > 1.0).
+    pub fn is_active(&self) -> bool {
+        match *self {
+            Self::Constant(s) => s > 1.0,
+            Self::Step { early, late, .. } => early > 1.0 || late > 1.0,
+            Self::LinearRamp { start, end, .. } => start > 1.0 || end > 1.0,
+        }
+    }
+
+    /// Parse the `--cfg-schedule` CLI value. Forms:
+    ///   - `step:<early>:<late>:<transition>`
+    ///   - `linear:<start>:<end>:<ramp_frames>`
+    ///   - `const:<scale>`
+    pub fn parse(s: &str) -> Result<Self> {
+        let err = || {
+            CsmError::Config(format!(
+                "invalid --cfg-schedule '{s}'; expected step:E:L:T, linear:S:E:R, or const:X"
+            ))
+        };
+        let mut parts = s.split(':');
+        let kind = parts.next().ok_or_else(err)?;
+        match kind {
+            "const" => {
+                let v: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                Ok(Self::Constant(v))
+            }
+            "step" => {
+                let e: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                let l: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                let t: usize = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                Ok(Self::Step {
+                    early: e,
+                    late: l,
+                    transition: t,
+                })
+            }
+            "linear" => {
+                let s0: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                let e: f64 = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                let r: usize = parts.next().ok_or_else(err)?.parse().map_err(|_| err())?;
+                Ok(Self::LinearRamp {
+                    start: s0,
+                    end: e,
+                    ramp_frames: r,
+                })
+            }
+            _ => Err(err()),
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn constant_returns_same_scale() {
+        let s = CfgSchedule::Constant(2.0);
+        assert_eq!(s.scale_at(0), 2.0);
+        assert_eq!(s.scale_at(50), 2.0);
+        assert!(s.is_active());
+    }
+
+    #[test]
+    fn step_transitions_at_threshold() {
+        let s = CfgSchedule::Step {
+            early: 3.0,
+            late: 1.5,
+            transition: 10,
+        };
+        assert_eq!(s.scale_at(0), 3.0);
+        assert_eq!(s.scale_at(9), 3.0);
+        assert_eq!(s.scale_at(10), 1.5);
+        assert_eq!(s.scale_at(100), 1.5);
+        assert!(s.is_active());
+    }
+
+    #[test]
+    fn linear_ramp_interpolates() {
+        let s = CfgSchedule::LinearRamp {
+            start: 3.0,
+            end: 1.0,
+            ramp_frames: 10,
+        };
+        assert_eq!(s.scale_at(0), 3.0);
+        assert!((s.scale_at(5) - 2.0).abs() < 1e-9);
+        assert_eq!(s.scale_at(10), 1.0);
+        assert_eq!(s.scale_at(100), 1.0); // holds at end
+    }
+
+    #[test]
+    fn parse_step() {
+        let s = CfgSchedule::parse("step:3.0:1.5:12").unwrap();
+        match s {
+            CfgSchedule::Step {
+                early,
+                late,
+                transition,
+            } => {
+                assert_eq!(early, 3.0);
+                assert_eq!(late, 1.5);
+                assert_eq!(transition, 12);
+            }
+            _ => panic!("wrong variant"),
+        }
+    }
+
+    #[test]
+    fn parse_linear() {
+        let s = CfgSchedule::parse("linear:2.5:1.0:20").unwrap();
+        match s {
+            CfgSchedule::LinearRamp {
+                start,
+                end,
+                ramp_frames,
+            } => {
+                assert_eq!(start, 2.5);
+                assert_eq!(end, 1.0);
+                assert_eq!(ramp_frames, 20);
+            }
+            _ => panic!("wrong variant"),
+        }
+    }
+
+    #[test]
+    fn parse_invalid() {
+        assert!(CfgSchedule::parse("garbage").is_err());
+        assert!(CfgSchedule::parse("step:1.5:1.0").is_err()); // missing transition
+    }
+}