MaRK: Markov-adapted Recurrent Kernels for Dynamic Operator Conditioning in State Space Models

AI summary

Overview

Research area: Machine learning, specifically sequence modeling with State Space Models (SSMs) and conditional generation via masked discrete diffusion.

Technical level: Advanced. The paper combines architectural design with Linear Parameter-Varying (LPV) control theory, Lyapunov-based stability certificates, and Rademacher complexity generalization bounds, though the core intuition is accessible.

Scope: The paper proposes and evaluates a method for conditioning a frozen SSM by modulating its internal recurrent parameters directly, rather than injecting context into the input stream or rescaling activations.

What This Paper Is About

Pre-trained State Space Models are usually adapted for conditional generation tasks such as diffusion by two external mechanisms: injecting the conditioning vector into the input stream, or applying adaptive layer normalization (AdaLN) to activations. Both approaches leave the model's underlying temporal dynamics, the recurrence (A), read-in (B), read-out (C), skip (D), and discretization (Δ) matrices, unchanged. MaRK instead maps a context vector (here, the diffusion timestep) directly into bounded modulations of those five parameters of a frozen backbone, so the model's input-output memory kernel itself becomes context-dependent.

Key Contributions

1. LPV-SSM operator conditioning. The authors cast conditional generation as Linear Parameter-Varying control over SSM operators, directly modulating A, B, C, D, and Δ as a view beyond input injection and AdaLN.
2. The MaRK kernel family. Three structural adapter geometries are designed that map conditioning vectors to low-rank modulations: Hypernet, Chebyshev polynomial, and Discrete Cosine Transform (DCT).
3. Stability certificates and operator recovery. Bounded parameterization via tanh and softplus is used to obtain analytic Affine Quadratic Stability (AQS) certificates, and Markov-operator recovery is studied in synthetic LPV identification experiments.
4. Markov operator adaptation. The paper provides evidence that MaRK conditions the Markov parameter sequence, with timestep-specific lag-norm profiles revealing distinct temporal-memory signatures across adapters.

Main Findings

• Chebyshev performs best overall. Averaged CART validation loss across the five reported datasets was 2.55 for Chebyshev (11.0M adapter parameters), 2.59 for DCT (7.1M), and 3.77 for Hypernet (6.3M). Chebyshev led DCT on PTB (2.42 vs. 2.48), Lambada (2.40 vs. 2.47), and Arxiv (2.55 vs. 2.61), and tied on WikiText (2.55 vs. 2.55); AG News was 2.85 for Chebyshev and 2.86 for DCT, described as inside seed noise.

• Geometry matters more than parameter count. Hypernet (6.3M) and DCT (7.1M) differ by 1.18 nats despite similar budgets, while the additional 3.9M parameters separating DCT from Chebyshev buy only 0.04 nats. The paper attributes this to basis smoothness acting as a structural regularizer: Chebyshev uses five global polynomials; DCT uses eight frequencies with a learned spectral decay.

• Modulating A is necessary. Freezing A increased validation loss by +0.18 to +0.25 points. The B,C-only variant degraded by +0.34 to +0.51 points, which the authors state shows MaRK cannot be reduced to standard Mamba-style selection. A-only outperformed Δ-only by up to 0.75 points, and All Except Δ achieved the second-best score for every kernel. Single-parameter configurations trailed the Full setting by up to +0.50 points.

• All 23 layers certify as stable. Using the identity matrix P = I as a common Lyapunov function with Δ_min = 0.001 and worst-case modulation at δ = −β_A, every layer of the pre-trained Hydra SSM returned a positive stability margin ε_ℓ = 1 − max_i Ā²_{i,ℓ} at extreme modulation bounds.

• Synthetic operator recovery follows near-parametric rates. Markov Operator Error decayed systematically as the number of synthetic trajectories grew from N = 50 to 1600, with fitted empirical exponents of N^−0.54 (Hypernet), N^−0.66 (Chebyshev), and N^−0.64 (DCT), against a theoretical O(1/√N) prediction (Proposition 4.1).

• Timestep-conditioned memory profiles are distinct and stable. Evaluating head-averaged Markov norms over lags k ≤ 4096 at t ∈ {0, 0.25, 0.5, 0.75, 1.0}, the authors found Hypernet produced the broadest upward shift in memory persistence, Chebyshev produced a tightly ordered family of smooth decay curves, and DCT produced the most selective redistribution, amplifying early-lag response at high-noise timesteps and crossing the frozen backbone at intermediate lags.

• Runtime overhead amortizes with sequence length. At sequence length 4096, MaRK variants added 14.0% (Hypernet), 15.8% (Chebyshev), and 14.9% (DCT) latency over frozen Hydra, with memory overhead below 1%. At sequence length 512 the overhead was +65.6%, +104.0%, and +88.7% respectively, driven by a constant per-layer conditioning cost of roughly 37 ms to 42 ms. Measured on an NVIDIA A100-SXM4-80GB at batch size 8, with 15 seeds and 100 forward passes per seed.

Methodology in Plain English

The authors start from a released 111M-parameter, 23-layer Hydra-BERT SSM checkpoint and keep it completely frozen. Around it they attach small adapter networks whose only job is to read the diffusion timestep context (encoded with standard DDPM/DiT sinusoidal encoding and a two-layer MLP) and output small, bounded adjustments to the five continuous-time matrices that define the SSM's dynamics.

To keep the recurrence stable, the adjustments are constrained. Diagonal parameters such as the log-space A and the Δ bias receive bounded shifts through tanh. The dense B and C matrices receive bounded scales and shifts applied through a low-rank factorization, avoiding a full-rank explosion. Because Hydra already builds its discrete eigenvalue as exp(A·Δ) with A < 0 and Δ > 0 via softplus, every modulated eigenvalue stays in the interval (0,1), which lets the authors prove stability analytically with the identity matrix as a Lyapunov function rather than solving a general semidefinite program.

The three adapter variants differ only in the basis used to map context to modulation: a direct MLP projection (Hypernet), a sum of Chebyshev polynomials evaluated at a learned latent coordinate (Chebyshev), and a spectrally damped combination of fixed Type-II DCT cosine basis functions (DCT). Training uses a masked discrete diffusion objective in which each token is masked with probability t and the loss is computed only on masked positions, reweighted per token by the Context-Adaptive Token-Level Noise Rescheduling (CART) scheme, which gives stronger weight to masked tokens surrounded by more visible context.

Validation proceeds on three fronts: synthetic LPV identification (recovering a known Markov operator from noisy trajectories), Markov-operator diagnostics (inspecting whether the induced memory kernel actually changes with timestep), and head-to-head adapter ablations under the same frozen backbone and objective.

Why This Matters

The paper reframes a fairly common question, how do you steer a pre-trained sequence model, as a question about the model's dynamics rather than its activations. If conditioning is applied to the recurrent operator itself, then the model's transfer function, impulse response, and memory decay become functions of the context. That is a stronger claim than feature re-projection, because the Markov parameter sequence is invariant under similarity transformations of the state-space coordinates, meaning two systems with the same sequence are behaviorally indistinguishable.

The authors also note a practical by-product: because modulation enters through low-rank auxiliary maps on a frozen backbone, parameter-efficient fine-tuning falls out of the adaptation mechanism structurally. Only 6.3M to 11M trainable auxiliary parameters were needed to move a bidirectional model into an iterative diffusion regime.

Potential applications, extrapolated from the paper's own framing of non-stationary sequence modeling:

• Diffusion-based language generation, where the reverse process runs over many timesteps and the denoising model may benefit from timestep-specific temporal dynamics.
• Adapting large pre-trained sequence models on constrained hardware, since the base weights stay frozen and only small adapters are trained.
• Long-context sequence tasks where the required memory horizon changes across the input, since A governs persistent memory while B and C govern instantaneous throughput.
• Any setting with a continuously varying conditioning signal (time, control input, noise level) that could serve as the LPV scheduling variable.

Industry relevance. The work was supported by the NVIDIA AI Technology Center, with authors affiliated with City University of Hong Kong, NVIDIA AI Technology Center, and NVIDIA Santa Clara. The latency and memory measurements are reported on an A100-SXM4-80GB, indicating the design was evaluated with deployment-side constraints in mind. The paper explicitly notes that DCT is cheaper than Chebyshev at inference, framing Chebyshev as the accuracy ceiling and DCT as the better parameter-latency trade.

Future Directions

• Systematic comparison against established conditioning frameworks. The authors state that comparing against input injection and AdaLN Zero on downstream benchmarks remains an important future direction, and that the reported results are an internal design ablation rather than a full scaling comparison against separately trained diffusion language models.
• Scaling behavior. How MaRK holds under scaling laws is named as an open direction, and the limitations are discussed further in Appendix A.5.
• Broadening the statistical guarantees. Proposition 4.1 is explicitly conditional on matched assumptions (bounded representations, persistent excitation, sub-Gaussian inputs and noise, exponential β-mixing). Relaxing those premises would be a natural extension.
• Extending beyond masked discrete diffusion. The conclusion points toward principled adaptation in non-stationary sequence modeling more generally, which would include conditioning signals other than diffusion timesteps.

Target Audience

This paper is best suited to researchers and engineers working on state space models, diffusion language models, and parameter-efficient fine-tuning who already have some grounding in dynamical systems or linear control theory. Readers interested in the intersection of SSM architecture design and Lyapunov stability analysis will get the most from Sections 2 and 2.1, while practitioners focused on adapter design trade-offs will find Section 4.4 and the ablation tables most directly useful. Readers with no background in LPV systems or stability theory should expect a substantial learning curve on the theoretical sections.

Authors’ abstract

State Space Models (SSMs) offer an efficient alternative to Transformers for sequence modeling, yet conditioning pre-trained SSMs for iterative generation typically operates outside the recurrent operator, through input injection or activation modulation. While such mechanisms expose the model to conditioning information, they leave the underlying temporal dynamics fixed. We introduce MaRK (Markov-adapted Recurrent Kernels), a dynamic operator-conditioning framework that maps context vectors directly into bounded modulations of a frozen SSM's recurrence ($A$), read-in ($B$), read-out ($C$), skip ($D$), and discretization ($Δ$) parameters. Viewed through the lens of LPV-SSM systems, MaRK induces a context-indexed family of Markov parameter sequences, allowing each diffusion timestep to reshape the model's input-output memory kernel. We instantiate MaRK on a frozen 111M-parameter Hydra SSM backbone and study three adapter geometries: Hypernet, Chebyshev polynomial, and Discrete Cosine Transform kernels. Since these adapters modify the Markov parameter sequence through low-rank auxiliary maps on the frozen backbone, parameter-efficient fine-tuning arises as a structural consequence of the adaptation mechanism itself, requiring only 6.3--11M trainable auxiliary parameters to transition from a bidirectional objective to an iterative diffusion regime. The bounded recurrence parameterization further yields an analytic Affine Quadratic Stability certificate for the modulated recurrence. Through synthetic LPV recovery experiments and Markov-operator diagnostics, we show that MaRK recovers coordinate-invariant temporal operators under matched assumptions and produces distinct, stable timestep-conditioned memory profiles. Empirically, the Chebyshev variant yields the strongest performance, achieving an average validation loss of 2.55, followed by the DCT (2.59) and Hypernet (3.77) geometries.