Struct NestedPrefixPenalty 

pub struct NestedPrefixPenalty {
    pub target: PsiSlice,
    pub target_tier: PenaltyTier,
    pub prefix_sizes: Vec<usize>,
    pub shell_weights: Vec<f64>,
    pub eps: f64,
    pub rho_indices: Vec<usize>,
    pub weight_schedule: Option<ScalarWeightSchedule>,
}

Nested-prefix sparsity penalty used by the Matryoshka SAE (Bussmann/Nabeshima/Karvonen/Nanda, ICML 2025, arXiv:2503.17547).

Given K nested prefix sizes m_1 < m_2 < ... < m_K ≤ F over the latent dimension F, and per-shell weights λ_k = w_k · exp(ρ_k), the penalty is

  P(t; ρ) = Σ_k λ_k · Σ_{i=0}^{m_k - 1} sqrt(t_i² + ε²)

summed over all rows of the latent target. Equivalently, coordinate i contributes with effective weight W_i = Σ_{k: m_k > i} λ_k, so the earliest atoms (small i) are penalized by every shell (= strongest L¹) and the latest atoms only by the outermost shell. This is exactly the mask-weighted sum-of-L¹ over K prefixes used to enforce shell-wise reconstruction during Matryoshka training.

Closed forms (per row, summed across all rows):

  ∂P/∂t_i      = W_i · t_i / sqrt(t_i² + ε²)
  Hess_diag(i) = W_i · ε² / (t_i² + ε²)^{3/2}           (PSD)
  ∂P/∂ρ_k      = λ_k · Σ_{i < m_k} sqrt(t_i² + ε²)

target lays out n_rows × latent_dim in row-major order (row * F + col). latent_dim is taken from PsiSlice::latent_dim; if absent we fall back to the maximum prefix size, which is the standard Matryoshka convention.

Fields

target: PsiSlice

target_tier: PenaltyTier

prefix_sizes: Vec<usize>

Sorted strictly-increasing prefix sizes m_1 < m_2 < ... < m_K.

shell_weights: Vec<f64>

Per-shell base weights w_k. The effective strength is λ_k = w_k · exp(ρ_k).

eps: f64

Smoothing parameter ε > 0 for the smoothed-L¹ surrogate sqrt(x² + ε²); the Hessian needs ε > 0 for differentiability at 0.

rho_indices: Vec<usize>

Local ρ indices for the K per-shell log-strengths.

weight_schedule: Option<ScalarWeightSchedule>

Implementations

impl NestedPrefixPenalty

pub fn new( target: PsiSlice, target_tier: PenaltyTier, prefix_sizes: Vec<usize>, shell_weights: Vec<f64>, eps: f64, ) -> Result<Self, String>

Build a new nested-prefix penalty.

Errors when:

• prefix_sizes is empty.
• prefix_sizes is not strictly increasing.
• any prefix exceeds the latent dimension (when known).
• shell_weights.len() != prefix_sizes.len().
• eps <= 0 (the smoothed-L¹ gradient 1/sqrt(x²+ε²) and Hessian ε²/(x²+ε²)^{3/2} both need ε > 0).

pub fn with_weight_schedule(self, schedule: ScalarWeightSchedule) -> Self

Attach a global annealing schedule shared by all shell weights. The REML loop still picks per-shell ρ_k on top of this baseline.

Trait Implementations

impl AnalyticPenalty for NestedPrefixPenalty

fn tier(&self) -> PenaltyTier

Tier the target lives in (β or ext-coord).

fn value(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> f64

Scalar penalty contribution P(target; ρ). The strength factor exp(ρ) (or whatever parameterization the penalty uses) is folded in.

fn grad_target( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, ) -> Array1<f64>

Gradient ∂P/∂target, same length as target.

fn hessian_diag( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, ) -> Option<Array1<f64>>

Diagonal of the Hessian diag(∂²P/∂target²) when the Hessian is block-diagonal. Returns None for penalties whose Hessian is dense (Isometry); those implement Self::hvp instead. The default signals “no closed-form diagonal” by returning None for any non-empty target — concrete penalties either override with their own analytic diagonal or rely on the matrix-free hvp path.

fn grad_rho( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, ) -> Array1<f64>

Gradient of the penalty value w.r.t. each owned ρ-axis. Length equals Self::rho_count.

fn rho_count(&self) -> usize

Number of REML-selectable hyperparameter axes this penalty contributes to the outer ρ vector.

fn name(&self) -> &str

Human-readable identifier for diagnostics / logging.

fn apply_schedule(&mut self, iter: usize)

Update any attached scalar weight schedule at the given REML outer iteration. Penalties without schedules keep their stored weight.

fn hvp( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, v: ArrayView1<'_, f64>, ) -> Array1<f64>

Hessian-vector product H v = (∂²P/∂target²) v, in closed form.

fn psd_majorizer_diag( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, ) -> Option<Array1<f64>>

Diagonal of a PSD majorizer of the Hessian — the positive re-weighted-ℓ₂ / MM surrogate diag(B(target; ρ)) with B ⪰ ∂²P/∂target² everywhere and B ⪰ 0. This is a different operator from Self::hessian_diag: for nonconvex penalties (log sparsity, JumpReLU) the exact Hessian is indefinite, but the inner Newton / PIRLS solve and the log-det / preconditioner pipeline require a PSD curvature block. For convex penalties the majorizer coincides with the exact Hessian, so the default simply delegates to Self::hessian_diag; nonconvex penalties override.

fn psd_majorizer_hvp( &self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>, v: ArrayView1<'_, f64>, ) -> Array1<f64>

Matrix-vector product against the PSD majorizer B(target; ρ) v (see Self::psd_majorizer_diag). For convex penalties this is the exact Hessian-vector product, so the default delegates to Self::hvp; nonconvex penalties override to return their PSD surrogate instead of the indefinite true Hessian.

impl Clone for NestedPrefixPenalty

fn clone(&self) -> NestedPrefixPenalty

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for NestedPrefixPenalty

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl PenaltyManifest for NestedPrefixPenalty

const KIND_TAG: &'static str = "nested_prefix"

const PYTHON_WRAPPER: &'static str = "NestedPrefixPenalty"

const ROW_BLOCK_DIAGONAL: bool = true

fn dispatch_tier(&self) -> PenaltyTier