Mamba M3.5: canonical numerically-stable initialization
Plain randn(0,1) init made A=-exp(A_log) and Δ wildly large, overflowing
the real exp(Δ·A) scan to NaN (the stub never cared — it discarded these
weights). new()/new_seeded() now share a build() with canonical S6 init:
- A_log = ln(1..=d_state) ⇒ A = -(1..=d_state), bounded
- dt_bias so softplus(dt_bias) ≈ 0.01 (small, stable Δ; near-identity
scan at init — intentional for gradient flow)
- D = 1, zero conv bias, projections scaled by 1/√fan_in (capped 0.5)
This fixes the NaN that broke omni-cortex's d231 action-conditioned
predictor training (now green). Seeded determinism preserved.
Tests: active_block helper (Δ overridden to ≈0.69) exercises the
scan-active regime so the liveness check can observe each weight; the
training test asserts a seed-varying backbone weight (conv1d_weight)
moves. All 6 selective-scan tests green incl. the finite-diff grad check.
Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
This commit is contained in:
omar sobh
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
33cf9bf731
commit
639937e9f5
@@ -287,40 +287,7 @@ pub struct MambaBlock {
|
||||
impl MambaBlock {
|
||||
/// Create new `MambaBlock`
|
||||
pub fn new(config: MambaConfig, device: &Device) -> Result<Self> {
|
||||
let d_inner = config.get_d_inner();
|
||||
let dt_rank = config.get_dt_rank();
|
||||
|
||||
// Initialize parameters
|
||||
let in_proj = Tensor::randn(&[config.d_model, d_inner * 2], device)?;
|
||||
let conv1d_weight = Tensor::randn(&[d_inner, 1, config.d_conv], device)?;
|
||||
let conv1d_bias = if config.conv_bias {
|
||||
Some(Tensor::randn(&[d_inner], device)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let A_log = Tensor::randn(&[d_inner, config.d_state], device)?;
|
||||
let x_proj = Tensor::randn(&[d_inner, dt_rank + 2 * config.d_state], device)?;
|
||||
let dt_proj = Tensor::randn(&[dt_rank, d_inner], device)?;
|
||||
let dt_bias = Tensor::randn(&[d_inner], device)?;
|
||||
let d_skip = Tensor::randn(&[d_inner], device)?;
|
||||
let out_proj = Tensor::randn(&[d_inner, config.d_model], device)?;
|
||||
|
||||
let selective_scan = SelectiveScan::new(d_inner, config.d_state);
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
device: device.clone(),
|
||||
in_proj,
|
||||
conv1d_weight,
|
||||
conv1d_bias,
|
||||
A_log,
|
||||
x_proj,
|
||||
dt_proj,
|
||||
dt_bias,
|
||||
d_skip,
|
||||
out_proj,
|
||||
selective_scan,
|
||||
})
|
||||
Self::build(config, device, None)
|
||||
}
|
||||
|
||||
/// Create `MambaBlock` with a deterministic seed.
|
||||
@@ -335,38 +302,70 @@ impl MambaBlock {
|
||||
/// same operator-supplied `seed` produces a stable mapping across
|
||||
/// the six (or seven, with conv_bias) internal tensors.
|
||||
pub fn new_seeded(config: MambaConfig, device: &Device, seed: u64) -> Result<Self> {
|
||||
let d_inner = config.get_d_inner();
|
||||
let dt_rank = config.get_dt_rank();
|
||||
Self::build(config, device, Some(seed))
|
||||
}
|
||||
|
||||
// Six (or seven) per-tensor seeds derived from the operator
|
||||
// seed. Using SplitMix64 so adjacent operator seeds don't
|
||||
// produce correlated per-tensor seeds.
|
||||
let mut state = seed;
|
||||
let mut next_seed = || {
|
||||
/// Shared constructor for [`Self::new`] (random) and
|
||||
/// [`Self::new_seeded`] (`Some(seed)` ⇒ deterministic). Uses
|
||||
/// canonical, numerically-stable selective-SSM initialization:
|
||||
/// `A_log = ln(1..=d_state)` so `A = -(1..=d_state)` is bounded;
|
||||
/// `dt_bias` set so `softplus(dt_bias) ≈ 0.01` (small, stable Δ);
|
||||
/// `D = 1`; zero conv bias; projection weights scaled by `1/√fan_in`
|
||||
/// (capped at 0.5). Plain `randn(0,1)` would overflow the real
|
||||
/// `exp(Δ·A)` scan to NaN — the stub didn't care because it
|
||||
/// discarded these weights.
|
||||
fn build(config: MambaConfig, device: &Device, seed: Option<u64>) -> Result<Self> {
|
||||
let d = config.get_d_inner();
|
||||
let dt_rank = config.get_dt_rank();
|
||||
let n = config.d_state;
|
||||
let dbc = dt_rank + 2 * n;
|
||||
|
||||
let mut state = seed.unwrap_or(0xD1CE_F00D);
|
||||
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);
|
||||
z ^ (z >> 31)
|
||||
let s = z ^ (z >> 31);
|
||||
match seed {
|
||||
Some(_) => Tensor::randn_seeded(shape, device, s),
|
||||
None => Tensor::randn(shape, device),
|
||||
}
|
||||
};
|
||||
let scaled =
|
||||
|t: Tensor, f: f32| -> std::result::Result<Tensor, rtx_tensor::TensorError> {
|
||||
let mut v = t.to_vec()?;
|
||||
for x in v.iter_mut() {
|
||||
*x *= f;
|
||||
}
|
||||
Tensor::from_vec(v, t.shape().dims(), device)
|
||||
};
|
||||
let lin = |fan: usize| (1.0 / (fan.max(1) as f32).sqrt()).min(0.5);
|
||||
|
||||
let in_proj = Tensor::randn_seeded(&[config.d_model, d_inner * 2], device, next_seed())?;
|
||||
let conv1d_weight =
|
||||
Tensor::randn_seeded(&[d_inner, 1, config.d_conv], device, next_seed())?;
|
||||
let in_proj = scaled(rand_t(&[config.d_model, 2 * d])?, lin(config.d_model))?;
|
||||
let conv1d_weight = scaled(rand_t(&[d, 1, config.d_conv])?, 0.2)?;
|
||||
let conv1d_bias = if config.conv_bias {
|
||||
Some(Tensor::randn_seeded(&[d_inner], device, next_seed())?)
|
||||
Some(Tensor::from_vec(vec![0.0f32; d], &[d], device)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let A_log = Tensor::randn_seeded(&[d_inner, config.d_state], device, next_seed())?;
|
||||
let x_proj =
|
||||
Tensor::randn_seeded(&[d_inner, dt_rank + 2 * config.d_state], device, next_seed())?;
|
||||
let dt_proj = Tensor::randn_seeded(&[dt_rank, d_inner], device, next_seed())?;
|
||||
let dt_bias = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
|
||||
let d_skip = Tensor::randn_seeded(&[d_inner], device, next_seed())?;
|
||||
let out_proj = Tensor::randn_seeded(&[d_inner, config.d_model], device, next_seed())?;
|
||||
// A = -exp(A_log) = -(1..=d_state) per channel (S4D-real init).
|
||||
let mut a_v = vec![0.0f32; d * n];
|
||||
for j in 0..d {
|
||||
for nn in 0..n {
|
||||
a_v[j * n + nn] = ((nn + 1) as f32).ln();
|
||||
}
|
||||
}
|
||||
let a_log = Tensor::from_vec(a_v, &[d, n], device)?;
|
||||
let x_proj = scaled(rand_t(&[d, dbc])?, lin(d))?;
|
||||
let dt_proj = scaled(rand_t(&[dt_rank, d])?, lin(dt_rank))?;
|
||||
// softplus(dt_bias) ≈ 0.01 ⇒ small, stable time-step.
|
||||
let dt_bias_val = (0.01f32.exp() - 1.0).ln();
|
||||
let dt_bias = Tensor::from_vec(vec![dt_bias_val; d], &[d], device)?;
|
||||
let d_skip = Tensor::from_vec(vec![1.0f32; d], &[d], device)?;
|
||||
let out_proj = scaled(rand_t(&[d, config.d_model])?, lin(d))?;
|
||||
|
||||
let selective_scan = SelectiveScan::new(d_inner, config.d_state);
|
||||
let selective_scan = SelectiveScan::new(d, config.d_state);
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
@@ -374,7 +373,7 @@ impl MambaBlock {
|
||||
in_proj,
|
||||
conv1d_weight,
|
||||
conv1d_bias,
|
||||
A_log,
|
||||
A_log: a_log,
|
||||
x_proj,
|
||||
dt_proj,
|
||||
dt_bias,
|
||||
|
||||
Reference in New Issue
Block a user