osobhandClaude Opus 4.7
aa274f2210
Adds the apply hook for ActAdd-style activation steering on the Llama
backbone. Inspired by EmoSteer-TTS (arXiv 2508.03543), but adapted: the
paper is flow-matching-specific (DiT layers, 32 CFM steps, per-token
attribution search via mel synthesis), none of which apply to CSM's
autoregressive Llama-over-Mimi-tokens. What's portable is the
underlying difference-in-means construction with residual-stream
addition — the standard ActAdd / contrastive-steering pattern.
What lands:
- src/steering.rs: LayerSteering type, per-layer (1, embed_dim) tensors,
global scale, safetensors load with keys `layer_<i>_steering`. Three
unit tests covering empty/no-op, dimension validation, and apply math.
- src/csm_fork.rs LlamaModel: optional `steering: Option<LayerSteering>`
field, applied after every layer's forward inside the for-loop. Adds
~3 LOC to the hot path; gated by the Option so unsteered generation
has zero cost beyond a None check.
- src/csm_fork.rs Model::set_backbone_steering: installs steering only
on the conditional backbone (cfg_backbone is intentionally left
un-steered so CFG correctly subtracts an unsteered baseline).
- src/generator.rs Generator::set_steering: errors on quantized
backend (only FP supported for now).
- examples/generate.rs: --steering-vec / --steering-scale flags.
- examples/steering_random.rs: smoke helper that writes random Gaussian
vectors so the apply path can be exercised end-to-end before the
real corpus extractor lands. Box-Muller via seeded rand to avoid an
extra rand_distr dep.
Smoke test (16-layer random Gaussian, stddev=0.05, scale=0.5):
- baseline (no steering, same seed/text): 3.04 s @ RMS -19.5 dB
- steered (random vectors): 1.84 s @ RMS -16.2 dB,
EOT triggered earlier
Output clearly differs — pathway is wired correctly. Random vectors
aren't musically meaningful; that's Phase B.
Phase B (next session): corpus extractor that runs forward passes over
emotion-labeled audio (we already have audio_to_manifest emitting
emotion_tag rows), captures per-layer post-residual activations, and
computes the difference-in-means between emotion_X and neutral pools.
Then A/B with quality_eval.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
79 lines
2.7 KiB
Rust
79 lines
2.7 KiB
Rust
//! Smoke-test helper: write a random `LayerSteering` safetensors so the
|
|
//! apply hook in `LlamaModel` can be exercised end-to-end before the real
|
|
//! emotion-extraction binary lands (Sprint 2 Phase B).
|
|
//!
|
|
//! The vectors are random Gaussian noise — they do NOT encode any
|
|
//! meaningful direction. Output should sound DIFFERENT from the unsteered
|
|
//! generation but no quality claim. This binary exists purely so the
|
|
//! plumbing is proven before we invest in corpus extraction.
|
|
//!
|
|
//! Usage:
|
|
//! ```bash
|
|
//! target/release/examples/steering_random \
|
|
//! --num-layers 16 --embed-dim 2048 \
|
|
//! --out /tmp/steering_random.safetensors
|
|
//!
|
|
//! target/release/examples/generate \
|
|
//! --text "Activation steering smoke test." \
|
|
//! --speaker 0 --max-audio-ms 3000 --seed 42 \
|
|
//! --steering-vec /tmp/steering_random.safetensors \
|
|
//! --steering-scale 0.5 \
|
|
//! --out /tmp/steered.wav
|
|
//! ```
|
|
|
|
use anyhow::Result;
|
|
use candle_core::{DType, Device, Tensor};
|
|
use clap::Parser;
|
|
|
|
#[derive(Debug, Parser)]
|
|
struct Cli {
|
|
#[arg(long)]
|
|
out: std::path::PathBuf,
|
|
/// CSM-1B has 16 backbone layers (BackboneFlavor::Llama1B).
|
|
#[arg(long, default_value_t = 16)]
|
|
num_layers: usize,
|
|
/// CSM-1B's Llama backbone embed_dim = 2048.
|
|
#[arg(long, default_value_t = 2048)]
|
|
embed_dim: usize,
|
|
/// Stddev of the per-layer Gaussian. Smaller = subtler steering.
|
|
#[arg(long, default_value_t = 0.05)]
|
|
stddev: f32,
|
|
#[arg(long, default_value_t = 42)]
|
|
seed: u64,
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
let cli = Cli::parse();
|
|
let dev = Device::Cpu;
|
|
|
|
// Box-Muller from a seeded uniform RNG so re-runs are reproducible
|
|
// without an extra `rand_distr` dep.
|
|
use rand::{Rng, SeedableRng};
|
|
let mut rng = rand::rngs::StdRng::seed_from_u64(cli.seed);
|
|
let mut gauss = || -> f32 {
|
|
let u1: f32 = rng.gen_range(1e-9..1.0);
|
|
let u2: f32 = rng.r#gen();
|
|
let r = (-2.0 * u1.ln()).sqrt();
|
|
let theta = 2.0 * std::f32::consts::PI * u2;
|
|
r * theta.cos() * cli.stddev
|
|
};
|
|
|
|
let mut tensors: Vec<(String, Tensor)> = Vec::with_capacity(cli.num_layers);
|
|
for i in 0..cli.num_layers {
|
|
let v: Vec<f32> = (0..cli.embed_dim).map(|_| gauss()).collect();
|
|
let t = Tensor::from_vec(v, (cli.embed_dim,), &dev)?.to_dtype(DType::F32)?;
|
|
tensors.push((format!("layer_{i}_steering"), t));
|
|
}
|
|
|
|
let map: std::collections::HashMap<String, Tensor> = tensors.into_iter().collect();
|
|
candle_core::safetensors::save(&map, &cli.out)?;
|
|
println!(
|
|
"wrote {} layer vectors (embed_dim={}, stddev={}) to {}",
|
|
cli.num_layers,
|
|
cli.embed_dim,
|
|
cli.stddev,
|
|
cli.out.display(),
|
|
);
|
|
Ok(())
|
|
}
|