Sharp Stability Threshold and Certification for Designing Stable Residual Architectures

Hyemin GuMichael Tyrrell
University of Massachusetts AmherstSRI International
hgu@umass.edu``michael.tyrrell@sri.com
Tuhin SahaiMarkos A. Katsoulakis
SRI InternationalUniversity of Massachusetts Amherst
tuhin.sahai@sri.com``markos@umass.edu


Abstract

We propose the sublinear-growth principle for deep residual architectures — a sharp stability threshold on the input-magnitude exponent of every residual block’s velocity field: \[\left\|v(x, t)\right\| \le c\,\left\|x\right\|^q + b, \qquad q \in [0, 1].\] The threshold \(q = 1\) is established via two independent arguments. Classical ODE theory gives a global forward flow on \([0, T]\) at \(q \le 1\) and exhibits divergent velocity fields at any \(q > 1\). The optimal-control analysis, via the Hamilton–Jacobi–Bellman equation, sharpens this to a selection statement: the training optimum is bang-bang on the boundary of the admissible class, so the optimum at \(q > 1\) blows up while the optimum at \(q \le 1\) is safe by construction. The exponent criterion \(q \le 1\) is thereby a necessary and sufficient condition for stable training. It clarifies architectural placements that ensure the stability of training and inference, explaining, for instance, the stabilizing role of layer normalization. The sublinear-growth velocity fields form the right function space on which forward dynamics, adjoint sensitivity, and architectural composition are all well-controlled. An arithmetic of input-magnitude exponents under the five operations that build residual blocks enables efficient certification of \(q_k \le 1\) at the level of architectural primitives, in place of ad hoc trial and error in the search for stable neural architectural designs. A parameter-free modification reduces the supercritical Mamba block from \(q = 5\) to \(q = 1\) without layer normalization, demonstrating this point. Experiments on Mamba and PatchTST confirm that the \(q \le 1\) variants train stably: the criterion is the input-magnitude exponent, not the presence of a normalization layer.

0.0.0.1 Acknowledgements.

This material is based upon work of the authors supported by the Defense Advanced Research Projects Agency (DARPA) under Agreement No. HR00112590112. Approved for public release; distribution is unlimited.

1 Introduction↩︎

Normalization layers — LayerNorm [@ba2016layernorm], RMSNorm [@zhang2019rmsnorm], BatchNorm [@ioffe2015batchnorm], GroupNorm [@wu2018groupnorm], InstanceNorm [@ulyanov2016instancenorm] — are standard components of modern deep residual architectures and appear in essentially every model across vision, language, speech, and generative modeling [@vaswani2017attention; @dosovitskiy2020vit; @gu2024mamba]. Their primary empirical role is to stabilize training convergence: they suppress runaway activation magnitudes, regularize gradient flow, and reduce sensitivity to initialization and learning rate [@santurkar2018howbn; @bjorck2018understandingbn; @huang2023normreview].

The mechanism behind this stabilization has been investigated empirically [@santurkar2018howbn; @bjorck2018understandingbn; @brock2021signalprop; @huang2023normreview], and a number of architectural alternatives that dispense with normalization have been proposed. Fixup [@zhangFixup2019] and ReZero [@bachlechner2021rezero] use specialized initialization to keep residual branches stable; self-normalizing networks [@klambauer2017snn] use a designed activation function; NFNets [@brock2021nfnets] combine weight normalization with adaptive gradient clipping; nGPT [@loshchilov2024ngpt] performs learning on the hypersphere; AERO [@jha2024aero] restricts the nonlinearity to softmax only; Dynamic Tanh [@zhu2025dyt] replaces layer normalization with an elementwise scaled tanh function. Reached by trial and empirical observation, these works show that normalization is not unique in its stabilizing properties: different computational actions encoded in architectural primitives can play the same role.

A theoretical understanding of stable training spans the architectural and learning-objective aspects within continuous-time control, mean-field-game, and Hamilton–Jacobi–Bellman (HJB) perspectives on deep learning. Haber and Ruthotto [@haber2017stable] address the architectural aspect, formulating residual networks as ODE flows and proposing architecture designs for which the continuous-time dynamics are well-posed. Zhang and Katsoulakis [@zhang2023mfg] address the learning-objective aspect, casting generative-model training as MFG optimality conditions and using the resulting HJB structure to design well-posed training regularizers for continuous normalizing flows, score-based generative models, and Wasserstein gradient flows.

Recent works in alignment with these aspects are by Kan et al. [@kan2025stability; @kan2026optimal]. [@kan2026optimal] works on the learning-objective axis, regularizing transformer training through OC-derived terms for generalization, robustness, and efficiency. [@kan2025stability] works on the architectural axis, analyzing Pre-LN and Peri-LN [@kim2025peri] through the HJB perspective. Across the architectural works, the theoretical perspective is applied to specific architectures or their components; to the best of our knowledge, no single criterion has been identified that spans a wide class of architectural primitives and normalization choices and is sharp enough to guide architectural search or discovery.

We propose such a principle, one that overarches architectural components encountered in deep residual architectures under a single criterion — the sublinear-growth principle: the velocity field of every residual block should grow at most sublinearly in the input magnitude, \[\left\|v(x, t)\right\| \le c\,\left\|x\right\|^q + b, \qquad q \in [0, 1], \label{eq:intro-bound}\tag{1}\] with constants \(c, b \ge 0\) uniform in the parameters \(\theta\). The motivation comes directly from the global-existence-versus-blow-up dichotomy in classical ODE theory (Section 2): at \(q \le 1\) the continuous-time ODE \(\dot{X} = v(X, t)\) admits a globally defined flow on \([0, T]\), while at any \(q > 1\) (the supercritical regime) there exist velocity fields whose forward trajectory blows up in finite time, and the discrete forward pass inherits this divergence by overflowing within a modest depth. The exponent \(q = 1\) is therefore the threshold of forward non-explosion, and a residual architecture built on a \(q > 1\) admissible class carries an unavoidable risk of overflow, independently of initialization or optimization.

The optimal-control analysis of Section 3 sharpens this picture. Casting training as an optimal-control problem on the velocity field \(v\), pointwise maximization of the Hamiltonian on the admissible class yields a closed-form optimal velocity field that is bang-bang (Theorem 1): at every \((X, t)\), \(v^\ast\) saturates the magnitude bound \(\left\|v\right\| \le c\,\left\|X\right\|^q + b\). Consequently, optimization cannot dodge the boundary of the admissible class. At \(q > 1\), this boundary contains velocity fields whose forward trajectory blows up in finite time, and the optimum selects one of them; at \(q \le 1\), the boundary is safe and the discrete residual stack inherits explicit forward state bounds and backward Jacobian bounds (Theorems 5 and 7), with constants uniform in the admissible parameters. The exponent \(q = 1\) is therefore the stability threshold under the OC framework: both necessary (the optimum at \(q > 1\) blows up) and sufficient (forward and backward stability at \(q \le 1\)) for stable training.

Together, the necessary and sufficient conditions from Sections 23 identify the family of velocity fields satisfying 1 as the right function space; we call it the sublinear-growth class. In Section 4 we make the principle 1 operational at the level of architecture design by developing an arithmetic of input-magnitude exponents under operations on architectural primitives. The input-magnitude exponent is the vocabulary in which standard primitives already sit — layer normalizations such as Peri-LN, bounded activations such as Dynamic Tanh [@zhu2025dyt], linear layers, Lipschitz activations, and self-attention. Each carries an input-magnitude exponent catalogued in Table 1, and the five operations that combine primitives into residual blocks transform these exponents by simple arithmetic rules (Theorem 9). Modern residual architectures have landed, by empirical practice, at block-level compositions of these primitives with \(q \le 1\); the most familiar instance is wrapping a sub-block with layer normalization. We articulate this implicit, algorithmic decision as an explicit, mathematically grounded, and sharp criterion for architectural certification: a single sharp threshold at \(q = 1\), certified per-block from primitive-level exponents and between-block across the residual stack. This unifies existing case-by-case analyses and widens the space of stable design candidates for architectural discovery.

We validate the principle 1 on two time-series forecasting backbones: Mamba [@gu2024mamba; @cmamba2024], whose native selective state-space operator drives the block-level exponent to \(q = 5\), and PatchTST [@nie2023patchtst], whose primitives are all \(q = 1\) by construction. For Mamba we evaluate five variants: the conventional RMSNorm baseline, Pre-LN, and Peri-LN (three normalization wrappers that each collapse the block to \(q = 0\)); a structural linear-growth modification that reduces the block to \(q = 1\) without any normalization layer; and a no-normalization control at the native \(q = 5\). For PatchTST we evaluate the same four normalization choices, the linear-growth modification being redundant since the block is already \(q = 1\). Across ETTm1 and Weather, the \(q \le 1\) variants train stably on every configuration, while the Mamba free-velocity control diverges at a rate that grows sharply with depth. The pattern dissociates the input-magnitude exponent from the presence of normalization: Mamba without normalization is supercritical and blows up, PatchTST without normalization is at \(q = 1\) and does not.

1.0.0.1 Contributions.

  1. As a design principle for the velocity field of every residual block, we propose the sublinear-growth bound \(\left\|v(x, t)\right\| \le c\left\|x\right\|^q + b\) with \(q \in [0, 1]\), and identify \(q = 1\) as a sharp stability threshold, via two independent arguments: classical non-explosion of the ODE (Section 2) and an optimal-control argument that identifies the training optimum as bang-bang on the boundary of the admissible class (Section 3).

  2. We derive forward state bounds and backward Jacobian bounds for the discrete residual architecture below the threshold (\(q \le 1\)), sharpened at the optimal-control optimum (Section 3.2).

  3. We make the principle operational at the level of architecture design by developing an arithmetic of input-magnitude exponents under five operations on primitives. Closure of the set of primitives under these operations (Theorem 9) reduces certification of \(q_k \le 1\) for a candidate block to inspection of its primitives, catalogued in Table 1 (Section 4).

  4. We rectify the supercritical Mamba block as a demonstrative example of this arithmetic: an elementary normalization-free modification reduces its input-magnitude exponent from \(q = 5\) to \(q = 1\) without any learned parameters (Section 5).

2 ODE View of Residual Architectures↩︎

A residual architecture is viewed throughout as a time discretization of the continuous-time ODE generated by a velocity field \(v\). Under this view, divergence of the forward trajectory is recast as finite-time blow-up of the ODE solution, and classical theory says exactly what prevents it: a growth condition on \(v\), hence on the residual block. That condition is the object of this section. Classical theory is available in continuous time; the discrete-time counterpart of each statement is recorded alongside it below.

2.0.0.1 Discrete and continuous-time dynamics.

Let \(X_k\) denote the hidden state at depth \(k\), taking values in a finite-dimensional Euclidean space \(\mathbb{R}^N\) (\(N = dn\) for sequence models with \(d\) channels and \(n\) tokens, via vectorization) with Euclidean inner product \(\langle \cdot, \cdot\rangle\) and induced norm \(\left\|\cdot\right\|\). A residual architecture evolves \(X_k\) through depth \(D\) layers via \[X_{k+1} = X_k + \Delta t \cdot v(X_k; \theta_k), \qquad k = 0, 1, \dots, D-1, \label{eq:discrete-dynamics}\tag{2}\] where \(\theta_k\) are the layer-\(k\) parameters, \(v\) is the velocity field, and \(\Delta t > 0\) is the step size (equivalently, the residual scaling; \(\Delta t = 1\) is the standard choice). We refer to \(v(X_k; \theta_k)\) interchangeably as the velocity at depth \(k\) and as the depth-\(k\) residual block, depending on whether the ODE or the architectural perspective is foregrounded. The continuous-time limit \(\Delta t \to 0\), \(D \to \infty\) with \(T = D\Delta t\) fixed yields the ODE \[\frac{dX(t)}{dt} = v(X(t), t; \theta(t)), \qquad X(0) = X_0, \qquad t \in [0, T] \label{eq:continuous-dynamics}\tag{3}\] with continuous-time velocity field \(v\). We refer to 2 as the forward trajectory of the network and to 3 as its continuous-time limit, as in the neural ODE [@chen2018neuralode] view of residual networks used here as an analytical limit. The mapping from inputs to outputs at depth \(D\) is the flow \(\Phi_D : X_0 \mapsto X_D\).

2.1 Non-explosion via classical ODE theory↩︎

The ODE perspective establishes non-explosion of the forward trajectory of the network: existence of a global solution \(X(t)\) on \([0, T]\) at each fixed velocity field. We recall the classical conditions under which 3 admits such a solution, formulated on the Euclidean space \(\mathbb{R}^N\) with norm \(\left\|\cdot\right\|\).

Definition 1 (Sublinear and linear growth conditions). A velocity field \(v : \mathbb{R}^N \times [0, T] \to \mathbb{R}^N\) satisfies a sublinear-growth condition with exponent \(q \in [0, 1]\) if there exist constants \(c, b \ge 0\) such that \[\left\|v(x, t)\right\| \le c\left\|x\right\|^q + b \qquad \text{for all } (x, t) \in \mathbb{R}^N \times [0, T]. \label{eq:lg-classical}\tag{4}\] The case \(q = 1\) is the linear-growth condition of classical ODE theory; the cases \(q \in [0, 1)\) are strictly sublinear.

2.1.0.1 Sublinear-growth threshold for non-explosion.

Condition 4 identifies \(q = 1\) as the threshold below which ODE solutions do not explode, by two well-known classical facts.

  • Global existence of solutions at \(q \le 1\). If \(v\) satisfies 4 with \(q = 1\) and \(c > 0\), and is continuous in \((x, t)\), then for every \(X_0 \in \mathbb{R}^N\) the ODE 3 admits a global solution on \([0, T]\), and any such solution satisfies the Grönwall bound \[\left\|X(t)\right\| \le e^{ct}\left\|X_0\right\| + (b/c)(e^{ct} - 1). \label{eq:gronwall-bound}\tag{5}\]

    Strictly sublinear \(q \in [0, 1)\) gives the same global existence with a sub-exponential (polynomial) growth bound, via Peano’s existence theorem and a scalar comparison argument [@hartman1964ode; @coddington1955ode]. Taking \(c = 0\) in 4 at any \(q\) leaves \(\left\|v\right\| \le b\), an instance of the bounded-velocity case \(q = 0\), for which the bound is \(\left\|X(t)\right\| \le \left\|X_0\right\| + bt\).

  • Finite-time blow-up at \(q > 1\). For every \(q > 1\) and every \(c > 0\), the velocity field \(v(x) := c\left\|x\right\|^{q-1} x\) obeys \(\left\|v(x)\right\| \le c\left\|x\right\|^q\) and reduces 3 to the scalar ODE \(\dot{r} = c r^q\) for \(r(t) := \left\|X(t)\right\|\), whose explicit solution blows up at \[t^\ast = \frac{\left\|X_0\right\|^{1-q}}{c(q-1)}. \label{eq:blowup-time}\tag{6}\]

Together, the two facts identify \(q = 1\) as the threshold input-magnitude exponent at which the continuous-time flow \(\Phi_T\) is globally defined on \([0, T]\): every \(q \le 1\) admits a global flow, while \(q > 1\) admits velocity fields whose flow is not defined on the whole interval. The discrete recursion 2 inherits a corresponding behavior, examined next.

2.1.0.2 Consequence of finite-time blow-up in discrete-time.

Within the regime \(c > 0\) of Definition 1, the forward-Euler recursion 2 stays below its continuous-time limit 3 , so continuous-time analysis is conservative for the discrete recursion. This conservativeness is not, however, a practical benefit at \(q > 1\). Concretely, the velocity \(v(x) = c \left\|x\right\|^{q-1} x\) reduces 2 to the scalar recursion \(r_{k+1} = r_k + \Delta t\, c\, r_k^q\) for \(r_k := \left\|X_k\right\|\), the forward-Euler discretization of \(\dot{r} = c r^q\); its per-step ratio \(r_{k+1}/r_k = 1 + \Delta t\, c\, r_k^{q-1}\) is itself increasing in \(r_k\), so the iterates grow super-exponentially and any fixed numerical range is exhausted within a modest depth, as Figure 1 illustrates for exponents \(q \in \left\{2, 5\right\}\).

2.1.0.3 From non-explosion to architectural design.

The consequence at the architectural level is direct: any residual architecture whose residual block has \(q > 1\) uniformly in \(\theta\) inherits the risk of forward blow-up. We therefore design or modify residual blocks so that the block-level input-magnitude exponent satisfies \(q \le 1\). Standard residual blocks are built from a set of architectural primitives (Table 1) combined by serial composition or parallel addition. The GPT-2 block [@radford2019gpt2code] composes layer normalization, self-attention, and an MLP in sequence; the Hymba block [@dong2024hymba] sums attention and state-space heads at the same depth. The arithmetic of Section 4 then computes the block-level \(q\) from those of the constituent primitives, enabling principled exploration of the architecture design space at \(q \le 1\) in place of ad hoc trial and error.

3 Optimal Control of Residual Architectures↩︎

Section 2 flagged finite-time blow-up of velocity fields at input-magnitude exponent \(q > 1\) through the ODE perspective, but did not close the story: whether this danger actually constrains training depends on where the OC optimum sits. This section closes the loop through the HJB framework. The OC optimum is bang-bang: it saturates the magnitude bound on the velocity field. Optimization therefore cannot dodge the dangerous velocities identified by the ODE analysis; the architecture must not admit \(q_k > 1\) at all depth \(k\) for stable training.

3.0.0.1 Discrete and continuous-time learning problem.

Alongside the forward trajectory 2 , training the residual architecture solves the optimization \[\min_{\theta}\, \mathbb{E}_{(X_0, y) \sim \rho_0}\, G(X_D, y) \quad \text{subject to } \eqref{eq:discrete-dynamics}, \label{eq:learning-discrete}\tag{7}\] with terminal cost \(G : \mathbb{R}^N \times \mathcal{Y} \to \mathbb{R}\) evaluated at the depth-\(D\) output \(X_D\) and the label \(y\). This is the training loss on the network’s task-specific output, typically a small prediction or classification head applied to \(X_D\). The parameter set is \(\theta = (\theta_0, \ldots, \theta_{D-1})\): each \(\theta_k\) parameterizes the depth-\(k\) block \(v(\cdot;\theta_k)\) as a composition of architectural primitives, and the resulting per-block velocity fields form the admissible class \(\mathcal{U}_{\mathrm{ad}}\). Forward evaluation of 2 is the forward pass of training; the optimization is typically conducted by backpropagation, which transports the discrete adjoint \(\lambda_k := \nabla_{X_k}\, G(X_D, y)\) backward in depth via the chain rule \[\lambda_k = \bigl(I + \Delta t\, \nabla_X v(X_k; \theta_k)\bigr)^{\!\top} \lambda_{k+1}, \qquad \lambda_D = \nabla_{X_D}\, G(X_D, y), \label{eq:discrete-adjoint}\tag{8}\] with the parameter gradient assembled along the way as \(\lambda_{k+1}^{\!\top}\, \nabla_{\theta_k} v(X_k; \theta_k)\). In the continuous-time limit this forward–backward procedure is the adjoint sensitivity method [@pontryagin1962], formulated for neural networks as the neural ODE [@chen2018neuralode]: the optimization becomes \[\min_{v \in \mathcal{U}_{\mathrm{ad}}}\, \mathbb{E}_{(X_0, y) \sim \rho_0}\, G(X(T), y) \quad \text{subject to } \frac{dX(t)}{dt} = v(X(t), t),\;X(0) = X_0, \label{eq:learning-continuous}\tag{9}\] with \(v\) ranging over the admissible class \(\mathcal{U}_{\mathrm{ad}}\), and the backward dynamics is governed by the adjoint equation \[\frac{d\lambda(t)}{dt} = -\bigl(\nabla_X v(X(t), t)\bigr)^{\!\top} \lambda(t), \qquad \lambda(T) = \nabla_X G(X(T), y), \label{eq:neural-ode-adjoint}\tag{10}\] solved backward in \(t\) from \(T\) to \(0\).

Ideally, both the forward and backward dynamics are well-posed via the choice of admissible class \(\mathcal{U}_{\mathrm{ad}}\), agreeing in both discrete and continuous-time formulations. Section 3.1 analyzes the OC optimum on a parametric admissible class with input-magnitude exponent \(q \ge 0\), deriving from the bang-bang nature of the optimum the stability threshold \(q \le 1\) for architectural design.

3.1 Optimal-control analysis of the admissible class \(\mathcal{U}_{\mathrm{ad}}^q\)↩︎

We now specialize the admissible set \(\mathcal{U}_{\mathrm{ad}}\) to the parametric class \[\mathcal{U}_{\mathrm{ad}}^q(X) := \left\{v : \left\|v\right\| \le c\left\|X\right\|^q + b\right\} \label{eq:Uad-generic}\tag{11}\] at input-magnitude exponent \(q \ge 0\). A member \(v\) of this class is a depth-\(k\) residual block (\(t = k\Delta t\)) composed from architectural primitives (attention, normalization, MLP, gating, activations); choosing the primitives is the choice of control. The optimum \(v^\ast\) admits a pointwise characterization. This leaves two structural questions open about \(v^\ast\) as a globally defined feedback law. First, does its forward trajectory stay finite on \([0, T]\)? Second, does backpropagation through this dynamics produce a well-defined gradient signal — the gradient of a function of state and time, rather than only a sequence of numerical Jacobian–vector products 8 ? Both are settled by the Hamilton–Jacobi–Bellman framework: the value function of the OC problem satisfies the HJB equation 13 , and \(q \le 1\) is necessary and sufficient for this equation to be well-posed on \([0, T]\): at \(q > 1\) the bang-bang trajectory blows up in finite time, and the linear-growth hypothesis on the Hamiltonian \(H_q\) that underlies the HJB argument no longer holds (necessary); at \(q \le 1\) linear growth holds, and the bang-bang optimum is a globally defined feedback law with bounded forward trajectory (sufficient).

3.1.0.1 Hamiltonian formulation and optimal-velocity.

The Hamiltonian associated with \(\mathcal{U}_{\mathrm{ad}}(X)\) is \[H(X, P) := \sup_{v \in \mathcal{U}_{\mathrm{ad}}(X)}\left\{-\langle P, v\rangle\right\}, \label{eq:Hamiltonian-def}\tag{12}\] where \(P\) is the costate. In continuous time, the value function of the OC problem 9 is \(u_y(X, t) := \inf_{v \in \mathcal{U}_{\mathrm{ad}}}\left\{G(X(T), y) \mid X(t) = X\right\}\) and satisfies the Hamilton–Jacobi–Bellman (HJB) equation \[\begin{align} -\partial_t u_y(X, t) + H(X, \nabla u_y(X, t)) &= 0, \quad (X, t) \in \mathbb{R}^{d \times n} \times [0, T), \label{eq:HJB-PDE} \\ u_y(X, T) &= G(X, y), \nonumber \end{align}\tag{13}\] in the viscosity sense [@crandall1992user; @tran2021hje]. In discrete time, the value function of the OC problem 7 is \(u_k(X)\) at depth \(k\), the optimal cost-to-go, and satisfies the Bellman recursion of dynamic programming and Markov chain approximation [@kushner2001stochastic; @bertsekas2017dp]: \[u_D(X) = G(X, y), \qquad u_k(X) = \inf_{v \in \mathcal{U}_{\mathrm{ad}}}\left\{\ell_k(X, v)\,\Delta t + u_{k+1}(X + \Delta t\, v)\right\}, \label{eq:bellman}\tag{14}\] solved backward in \(k\) from terminal data \(u_D = G\), with indicator running cost \(\ell_k(X, v) = \delta_{\mathcal{U}_{\mathrm{ad}}}(v)\). Under step size \(\Delta t = T/D\), \(u_k^{\Delta t}\) at \(t = k\Delta t\) converges to the viscosity solution of 13 as \(\Delta t \to 0\) by Barles–Souganidis convergence [@barles1991convergence].

The optimal velocity is obtained in closed form by maximizing the Hamiltonian at the value-function gradient, \(\nabla u_y(X, t)\) in continuous time and \(\nabla u_{k+1}(X)\) in discrete time: \[v^\ast(X, t) := -\nabla_P H(X, \nabla u_y(X, t)), \qquad v_k^\ast(X) := -\nabla_P H(X, \nabla u_{k+1}(X)), \label{eq:v-opt-hamiltonian}\tag{15}\] with the same Hamiltonian \(H\) in both, evaluated at different gradients.

3.1.0.2 Bang-bang optimum on \(\mathcal{U}_{\mathrm{ad}}^q\) and the stability threshold \(q\leq 1\).

Maximizing \(\langle v, P\rangle\) over the Euclidean ball \(\mathcal{U}_{\mathrm{ad}}^q(X)\) is a one-variable problem, yielding closed forms for \(H_q\) and \(v^\ast\).

Theorem 1 (Bang-bang optimum on \(\mathcal{U}_{\mathrm{ad}}^q\)). Let \(\mathcal{U}_{\mathrm{ad}}^q\) be defined by 11 at \(q \ge 0\). The Hamiltonian 12 and the optimal velocity fields 15 have closed forms \[H_q(X, P) = (c\left\|X\right\|^q + b)\,\left\|P\right\|, \label{eq:H-generic}\tag{16}\] \[v^\ast(X, t) = -\bigl(c\left\|X\right\|^q + b\bigr)\, \frac{\nabla u_y(X, t)}{\left\|\nabla u_y(X, t)\right\|}, \qquad v_k^\ast(X) = -\bigl(c\left\|X\right\|^q + b\bigr)\, \frac{\nabla u_{k+1}(X)}{\left\|\nabla u_{k+1}(X)\right\|}, \label{eq:v-opt-generic}\tag{17}\] on \(\left\{\nabla u_y(X, t) \ne 0\right\}\) and \(\left\{\nabla u_{k+1}(X) \ne 0\right\}\) respectively. The optimum is bang-bang: \[\left\|v^\ast(X, t)\right\| = \left\|v_k^\ast(X)\right\| = c\left\|X\right\|^q + b, \label{eq:bang-bang}\tag{18}\] attained on the boundary of \(\mathcal{U}_{\mathrm{ad}}^q(X)\) antiparallel to the corresponding costate.

Proof. By Cauchy–Schwarz, \(\sup_{v \in \mathcal{U}_{\mathrm{ad}}^q(X)}\left\{-\langle P, v\rangle\right\} = (c\left\|X\right\|^q + b)\,\left\|P\right\|\), attained at \(v = -(c\left\|X\right\|^q + b)\,P/\left\|P\right\|\) on \(\left\{P \ne 0\right\}\). Substituting \(P = \nabla u_y(X, t)\) in 15 gives the continuous closed form, with norm \(c\left\|X\right\|^q + b\) on the stated set; the discrete case follows from the same Cauchy–Schwarz with \(P = \nabla u_{k+1}(X)\). ◻

The bang-bang optimum saturates the boundary of the admissible class, not its interior. At \(q > 1\) this boundary contains velocity fields whose forward trajectories blow up in finite time on \([0, T]\) (Section 2); at \(q \le 1\) every boundary velocity field admits a global trajectory. The stability threshold \(q \le 1\) is therefore necessary for the OC optimum to remain globally defined on \([0, T]\). Figure 1 illustrates the mechanism: at initialization the effective coefficient \(c\) of \(v_\theta\) is small and all trajectories appear linear on \([0, T]\) regardless of \(q\), so a supercritical architecture is indistinguishable from a subcritical one in the observation window; as training grows \(c\) toward the admissible-class boundary, the supercritical trajectories cross into finite-time blow-up while the subcritical ones remain bounded.

Figure 1: State trajectories \|X_k\| under the residual-stack recursion r_{k+1} = r_k + c\, r_k^q with r_0 = 1 and \Delta t = 1, where c is a schematic proxy for a training parameter — small at initialization (left, c = 0.05) and growing toward the admissible-class ceiling as training approaches the optimal-control optimum (right, c = 1). The axis k is the discrete layer index. At small c, all q \le 2 trajectories appear near-linear across k \in[0, D] and a q > 1 architecture is indistinguishable from a q = 0 one at moderate depth. At larger c, the supercritical trajectories (q > 1) overflow within a few layers while q \le 1 remain bounded — the architecture with q > 1 therefore harbors a latent instability, invisible at initialization, that surfaces once c grows enough to bring the overflow depth k_\text{overflow}(c) into [0, D].

Remark 2 (Geometric realizations for sequence-model hidden states). For sequence-model hidden states \(X \in \mathbb{R}^{d \times n}\), the magnitude bound 11 admits two geometric realizations of the admissible set: an entrywise box (per-\((i, j)\) bound on \(\left|v_{ij}\right|\)) and a per-token \(\ell^2\) ball (per-column bound on \(\left\|v_{:,j}\right\|_2\)).

3.1.0.3 Well-posedness of the value function.

Theorem 1 and the stability threshold ensure that the optimal trajectory of the OC problem exists on \([0, T]\) at \(q \le 1\). A separate question is whether the value function is well-posed — existence, uniqueness, and continuous dependence on the terminal datum \(G\). Uniqueness and sufficient regularity of \(u_y\) are needed for the costate \(P(t) := \nabla_X u_y(X^\ast(t), t)\) to be well-defined a.e. along the optimal trajectory — Theorem 7 in Section 3.2 avoids this route altogether, working with the discrete adjoint \(\lambda_k\) directly via Danskin’s envelope identity under the sole assumption that an OC-optimal velocity field \(v^\ast\) exists.

For the continuous value function \(u_y\) on the unbounded spatial domain \(\mathbb{R}^{d \times n}\), the Crandall–Lions comparison principle delivers well-posedness when \(H\) satisfies a modulus condition \[\left|H(X_1, P) - H(X_2, P)\right| \le \omega\bigl(\left\|X_1 - X_2\right\|(1 + \left\|P\right\|)\bigr) \label{eq:CL-modulus}\tag{19}\] for some modulus \(\omega\), together with a linear-growth control \[\left|H(X, P)\right| \le C\,(1 + \left\|X\right\|)(1 + \left\|P\right\|) \label{eq:CL-linear-growth}\tag{20}\] uniform in \((X, P)\) [@crandall1992user; @tran2021hje]. Without the modulus condition, viscosity solutions can exist but be non-unique, and continuous dependence then loses meaning.

Proposition 3 (Continuous-time HJB well-posedness on \(\mathcal{U}_{\mathrm{ad}}^q\) at \(q \le 1\)). Let \(H_q\) be the Hamiltonian of 16 , and let terminal datum \(G\) be bounded uniformly continuous. When \(q \in (0, 1)\), additionally assume \(b > 0\) and \(G\) globally Lipschitz. Then at \(q \in [0, 1]\), the HJB equation 13 admits a unique bounded viscosity solution \(u_y \in C(\mathbb{R}^{d \times n} \times [0, T])\) coinciding with the value function of the OC problem 9 , and \[\left\|u_y^{G_1} - u_y^{G_2}\right\|_{L^\infty(\mathbb{R}^{d \times n} \times [0, T])} \le \left\|G_1 - G_2\right\|_{L^\infty(\mathbb{R}^{d \times n})} \label{eq:hjb-stability}\qquad{(1)}\] for any two admissible terminal data \(G_1, G_2\). At \(q > 1\), the Hamiltonian \(H_q\) violates the linear-growth control 20 .

Proof. Linear-growth control. For \(q \in [0, 1]\), \(\left\|X\right\|^q \le 1 + \left\|X\right\|\) gives \(\left|H_q(X, P)\right| = (c\left\|X\right\|^q + b)\left\|P\right\| \le (c + b)(1 + \left\|X\right\|)(1 + \left\|P\right\|)\), so 20 holds with \(C = c + b\). At \(q > 1\), \(\left|H_q(X, P)\right| \ge c\left\|X\right\|^q\left\|P\right\|\) is superlinear in \(\left\|X\right\|\), and 20 fails.

Modulus analysis. Using \(\left|a^q - b^q\right| \le \left|a - b\right|^q\) for \(a, b \ge 0\), \(q \in [0, 1]\), together with \(\bigl|\left\|X_1\right\| - \left\|X_2\right\|\bigr| \le \left\|X_1 - X_2\right\|\), \[\left|H_q(X_1, P) - H_q(X_2, P)\right| \;\le\; c\,\left\|X_1 - X_2\right\|^q\,\left\|P\right\|. \label{eq:holder-q}\tag{21}\] At \(q = 0\), \(H_0(X, P) = b\left\|P\right\|\) is \(X\)-independent and 19 holds trivially. At \(q = 1\), 21 yields \(c\left\|X_1 - X_2\right\|\left\|P\right\| \le c\left\|X_1 - X_2\right\|(1 + \left\|P\right\|)\), so 19 holds with \(\omega(s) = cs\). At \(q \in (0, 1)\), 19 does not hold globally: parameterizing \(\left\|X_1 - X_2\right\|(1 + \left\|P\right\|) = s\) fixed with \(\left\|X_1 - X_2\right\| = s/(1 + \left\|P\right\|)\), the right-hand side of 21 becomes \(c s^q \cdot \left\|P\right\|/(1 + \left\|P\right\|)^q\), whose supremum over \(\left\|P\right\| \ge 0\) diverges because \(q < 1\). On any bounded gradient set \(\left\{\left\|P\right\| \le R\right\}\), however, 21 yields the \(R\)-dependent modulus \(\omega_R(s) = cR\,s^q\).

Well-posedness. At \(q = 0\), \(H_0\) is \(X\)-independent, so comparison for \(X\)-independent Hamiltonians is standard [@crandall1992user]; only bounded uniform continuity of \(G\) is required. Existence and ?? follow from Perron’s method. At \(q = 1\), 19 with \(\omega(s) = cs\) (from Step 2) together with 20 gives well-posedness by the Crandall–Lions comparison principle and Perron’s method under bounded uniformly continuous \(G\) [@crandall1992user; @tran2021hje]. At \(q \in (0, 1)\), the failure of 19 globally (Step 2) precludes the standard CL route. Instead, the additional hypotheses \(b > 0\) and Lipschitz \(G\) enable the following alternative. The coercivity \[H_q(X, P) \ge b\,\left\|P\right\| \quad \text{uniformly in } X, \label{eq:coercivity}\tag{22}\] together with the Lipschitz terminal datum, implies that any bounded viscosity solution is Lipschitz in \(X\) on \(\mathbb{R}^{d \times n} \times [0, T]\) with a constant \(L\) depending only on \(b\), \(T\), and \(\mathrm{Lip}(G)\) [@bardi1997optimal; @fleming2006controlled]. Comparison between two Lipschitz solutions requires 19 only on the bounded gradient set \(\left\{\left\|P\right\| \le L\right\}\), where the \(R\)-dependent modulus \(\omega_L(s) = cL\,s^q\) established in Step 2 closes the standard doubling-variable argument [@crandall1992user]. Existence follows from Perron’s method with the value function of 9 as a viscosity solution, and ?? follows from the standard \(L^\infty\)-estimate on the difference of two value functions with common Hamiltonian. ◻

The discrete value function \(u_k\) of the Bellman recursion 14 admits a simpler route to well-posedness that bypasses viscosity-solution machinery. The per-block Bellman operator \[\mathcal{T}_k[\phi](X) := \inf_{v \in \mathcal{U}_{\mathrm{ad}}}\left\{\ell_k(X, v)\,\Delta t + \phi(X + \Delta t\, v)\right\} \label{eq:bellman-operator}\tag{23}\] is monotone and constant-shift invariant by its infimum structure; these two properties deliver well-posedness directly.

Proposition 4 (Discrete-time HJB well-posedness on \(\mathcal{U}_{\mathrm{ad}}^q\)). Let \(\mathcal{T}_k\) be the per-block Bellman operator 23 . Then \(\mathcal{T}_k\) is monotone: \(\phi \le \psi \Rightarrow \mathcal{T}_k[\phi] \le \mathcal{T}_k[\psi]\); and constant-shift invariant: \(\mathcal{T}_k[\phi + c] = \mathcal{T}_k[\phi] + c\) for \(c \in \mathbb{R}\). For any \(q \ge 0\) and bounded terminal data \(u_D = G\), the Bellman recursion 14 admits a unique discrete value function \(u_k : \mathbb{R}^{d \times n} \to \mathbb{R}\) at each depth \(k \in \left\{0, 1, \ldots, D-1\right\}\) via the backward iteration \(u_k = \mathcal{T}_k[u_{k+1}]\), and \[\left\|u_k^{G_1} - u_k^{G_2}\right\|_{L^\infty(\mathbb{R}^{d \times n})} \le \left\|G_1 - G_2\right\|_{L^\infty(\mathbb{R}^{d \times n})} \label{eq:bellman-stability}\qquad{(2)}\] for any two bounded terminal data \(G_1, G_2\) [@kushner2001stochastic; @bertsekas2017dp].

Proof. Monotonicity follows from the infimum structure of 23 : if \(\phi \le \psi\), then for every \(v \in \mathcal{U}_{\mathrm{ad}}\), \(\ell_k(X, v)\Delta t + \phi(X + \Delta t\, v) \le \ell_k(X, v)\Delta t + \psi(X + \Delta t\, v)\), and the infimum over \(v\) preserves the inequality. Constant-shift invariance follows from \(\inf_v\{f(v) + c\} = \inf_v f(v) + c\) for any constant \(c\) and function \(f\). Existence and uniqueness of \(u_k\) via the backward iteration \(u_k = \mathcal{T}_k[u_{k+1}]\) from \(u_D = G\) are then immediate from the recursive structure with bounded \(G\). The \(L^\infty\) contraction ?? follows from monotonicity plus constant-shift invariance: \(G_1 \le G_2 + \left\|G_1 - G_2\right\|_\infty \Rightarrow u_D^{G_1} \le u_D^{G_2} + \left\|G_1 - G_2\right\|_\infty \Rightarrow u_k^{G_1} \le u_k^{G_2} + \left\|G_1 - G_2\right\|_\infty\) by backward induction (monotonicity preserves the inequality through \(\mathcal{T}_k\); constant-shift invariance carries the constant through unchanged). Symmetric argument gives the reverse inequality. ◻

Propositions 3 and 4 make both continuous-time and discrete-time value functions well-defined. Still, for the value-function gradient along the optimal trajectory to be well-defined, so that the optimal velocity fields \(v^\ast\) are correspondingly well-defined, we need additional regularity on the value function, addressed below.

3.1.0.4 Regularity of the value function for the optimal velocity.

Under the Crandall–Lions hypotheses of Proposition 3 together with Lipschitz continuity of the terminal datum \(G\), the viscosity solution \(u_y\) is locally Lipschitz on \(\mathbb{R}^{d \times n} \times [0, T]\) [@fleming2006controlled]. By Rademacher’s theorem, \(u_y\) is therefore differentiable on a Lebesgue-full subset of \(\mathbb{R}^{d \times n} \times [0, T]\), with \(\nabla_X u_y\) the classical gradient on that subset. We take the costate \(P(t) := \nabla_X u_y(X^\ast(t), t)\) in this sense: it is defined for a.e. \(t \in [0, T]\) along the OC-optimal trajectory \(X^\ast(t)\). The bang-bang direction of Theorem 1 and the costate-based stability bounds of Section 3.2 are understood at points where \(\nabla_X u_y\) is well-defined; all downstream estimates involve integrals against \(dt\), for which a.e. definition is sufficient.

3.1.0.5 Well-posed OC problem for stable training.

The bang-bang nature of the optimal velocity \(v^\ast\) (Theorem 1) places the OC optimum on the boundary of \(\mathcal{U}_{\mathrm{ad}}^q\). Engineering safeguards such as careful initialization may keep realized velocities away from this boundary during training, but the OC optimum itself saturates it (Figure 1 illustrates this contrast in a toy setting). At \(q > 1\), the boundary contains velocity fields whose forward trajectory blows up in finite time on \([0, T]\) (Section 2), and the optimum may select precisely those; blow-up is then a possibility the optimization cannot dodge. At \(q \le 1\), by contrast, every admissible velocity field — not just the optimum — yields a globally defined trajectory, so blow-up is ruled out by the class itself rather than by any property of the optimum. The stability threshold \(q \le 1\) is therefore the operational requirement for stable training of the OC problem. The coincident HJB well-posedness at \(q \le 1\) (Proposition 3), together with the Lipschitz a.e.-differentiability argument above, makes \(\nabla u_y\) well-defined a.e. on the optimal trajectory, so the costate \(P(t) := \nabla_X u_y(X^\ast(t), t)\) along the optimal trajectory is a well-defined object a.e.

The costate satisfies the HJB backward characteristic \[\dot{P}(t) = \nabla_X H_q(X^\ast(t), P(t)), \qquad P(T) = \nabla_X G(X^\ast(T), y), \label{eq:hjb-backward-char}\tag{24}\] which coincides with the adjoint equation 10 evaluated along the optimal trajectory \(X^\ast(\cdot)\) with optimal velocity \(v = v^\ast\), via Danskin’s envelope identity [@danskin1967] \[\nabla_X H_q(X, P) \;=\; -\bigl(\partial_X v^\ast(X, t)\bigr)^{\!\top} P, \label{eq:danskin}\tag{25}\] with \(\partial_X\) taken at fixed costate \(P = \nabla u_y(X, t)\). Section 3.2 uses the well-defined a.e. costate together with Danskin’s identity to derive forward and backward stability bounds on the discrete architecture, refining the Lipschitz behavior of \(v^\ast\) into quantitative depth-explicit bounds.

3.2 Forward and backward stability of the discrete architecture↩︎

This subsection brings the OC analysis of Section 3.1 to the discrete architecture 2 — the actual residual stack trained by backpropagation — and produces quantitative depth-explicit stability bounds. The two objects analyzed are the forward state \(X_k\) generated by the residual recursion \[X_{k+1} = X_k + \Delta t\, v(X_k;\theta_k), \qquad k = 0, 1, \dots, D-1, \label{eq:fwd-bwd-state}\tag{26}\] and the discrete costate (equivalently, the backpropagated adjoint) \(P_k\) generated by the chain-rule recursion \[P_k = \bigl(I + \Delta t\,\nabla_X v(X_k;\theta_k)\bigr)^{\!\top} P_{k+1}, \qquad P_D = \nabla_X G(X_D, y), \label{eq:fwd-bwd-costate}\tag{27}\] solved backward from depth \(D\). The cumulative Jacobian \[J_{i:D} \;:=\; \prod_{k=i+1}^{D} \bigl(I + \Delta t\,\nabla_X v(X_{k-1};\theta_{k-1})\bigr) \label{eq:fwd-bwd-jacobian}\tag{28}\] transports \(P_D\) to \(P_i\) as \(P_i = J_{i:D}^{\!\top} P_D\), so a bound on \(\left\|J_{i:D}\right\|_2\) is a bound on the backpropagated costate uniformly in the terminal datum.

3.2.0.1 Forward stability.

The forward analysis redoes the Grönwall argument on 26 (Theorem 5). The bound applies to every admissible trajectory, not only to the OC-optimal one.

Theorem 5 (Forward state bound). Let \(\left\{X_k\right\}_{k=0}^D\) be the discrete trajectory 2 with residual scaling \(\Delta t > 0\) and any admissible parameters \(\theta_k\) realizing \(v(\cdot;\theta_k) \in \mathcal{U}_{\mathrm{ad}}^q\) with scalar constants \(c, b\) uniform in \(\theta_k\). Then \[\left\|X_D\right\| \;\le\; \begin{cases} \left\|X_0\right\| + b\,D\Delta t, & q = 0 \text{ (linear in D)}, \\[4pt] \bigl(\left\|X_0\right\|^{1-q} + c(1-q)\,D\Delta t\bigr)^{1/(1-q)}, & q \in (0, 1),\;b = 0 \text{ (polynomial in D)}, \\[4pt] (1 + c\Delta t)^D\,\left\|X_0\right\| + \tfrac{b}{c}\bigl((1 + c\Delta t)^D - 1\bigr), & q = 1 \text{ (geometric in D)}. \end{cases} \label{eq:forward-bound}\tag{29}\] The \(q = 1\) bound is majorized by \(e^{cD\Delta t}\left\|X_0\right\| + \tfrac{b}{c}(e^{cD\Delta t} - 1)\); the \(q \in (0, 1)\) bound is polynomial in \(D\) of degree \(1/(1-q)\). For \(b > 0\) at \(q \in (0, 1)\), the shift \(r := (b/c)^{1/q}\) gives the analogous bound on \(\left\|X_D\right\| + r\) with the rate constant \(c\) replaced by \(2c\).

Theorem 5 applies both off the OC optimum and at it: the input-magnitude exponent \(q\) governs the qualitative growth regime, while the parameters \(c, b\) control the quantitative size. Figure 1 illustrates this parameter dependence: an architecture that appears safe at initialization can diverge as training approaches the OC optimum.

Proof. At \(q = 0\) the bound 11 reads \(\left\|v(X_k;\theta_k)\right\| \le b\); substituting into 2 gives \(\left\|X_{k+1}\right\| \le \left\|X_k\right\| + b\Delta t\), and iterating \(D\) times yields the \(q = 0\) bound. At \(q = 1\) the bound reads \(\left\|v(X_k;\theta_k)\right\| \le c\left\|X_k\right\| + b\); substituting gives \(\left\|X_{k+1}\right\| \le (1 + c\Delta t)\left\|X_k\right\| + b\Delta t\), and iterating \(D\) times with the geometric series yields the \(q = 1\) bound. The majorization by \(e^{cD\Delta t}\left\|X_0\right\| + \tfrac{b}{c}(e^{cD\Delta t} - 1)\) follows from \((1 + c\Delta t)^D \le e^{cD\Delta t}\), applied to both terms. For \(q \in (0, 1)\) at \(b = 0\), concavity of \(f(x) = x^{1-q}\) gives \((a + h)^{1-q} \le a^{1-q} + (1-q)\,a^{-q}\,h\) for \(a, h > 0\); applying this to \(a = \left\|X_k\right\|\) and \(h = c\Delta t\,\left\|X_k\right\|^q\) yields \(\left\|X_{k+1}\right\|^{1-q} \le \left\|X_k\right\|^{1-q} + c(1-q)\,\Delta t\), and telescoping from \(k = 0\) to \(D\) gives the \(q \in (0, 1)\) bound. For \(b > 0\), the substitution \(Y_k := \left\|X_k\right\| + (b/c)^{1/q}\) satisfies \(Y_{k+1} \le Y_k + \Delta t\,(c Y_k^q + b) \le Y_k + 2c\Delta t\,Y_k^q\), reducing to the \(b = 0\) case with \(2c\) in place of \(c\). ◻

Remark 6 (Practical depth limits at \(q = 1\)). The exponential factor \((1 + c\,\Delta t)^D\) in 29 is finite at every depth but grows rapidly. With \(\Delta t = 1\) and \(c = 1\) the factor is \(2^D\), exceeding the FP16 maximum (\(\approx 6.5 \times 10^4\)) at \(D = 17\) and the FP32/BF16 dynamic range (\(\approx 3.4 \times 10^{38}\)) at \(D = 128\). Even at \(q = 1\), therefore, the input-magnitude bound imposes a practical depth limit set by the numerical range.

3.2.0.2 Backward stability.

The backward dynamics of the learning problem are governed by the adjoint equation 8 during training (off the OC optimum) and by the HJB backward characteristic 24 at the OC optimum via Danskin’s identity. We treat the two regimes in turn.

Off-optimum: cumulative Jacobian under a Lipschitz hypothesis. Suppose the velocity field has, along the realized discrete trajectory \(\left\{X_k\right\}_{k=0}^D\), an operator-norm gradient bounded by a constant \(L\) uniform in \(k\) and the admissible parameters: \(\max_{0 \le k < D} \left\|\nabla_X v(X_k; \theta_k)\right\|_2 \le L\). By the triangle inequality and submultiplicativity of the operator norm, the cumulative backward Jacobian 28 satisfies the standard bound [@gouk2021lipschitz; @miyato2018spectral; @anil2019sorting] \[\left\|J_{i:D}\right\|_2 \;\le\; (1 + L\Delta t)^{D-i} \;\le\; e^{L(T - i\Delta t)}, \label{eq:fwd-bwd-jacobian-bound}\tag{30}\] uniformly in the admissible parameters \(\left\{\theta_k\right\}\). The constant \(L\) is determined by the architectural primitives: for a Peri-LN block at \(q = 0\) with Lipschitz sub-block \(F\) (attention or MLP), \(L \le L_{\mathrm{Peri}} = \mathrm{Lip}(\mathrm{LN})\mathrm{Lip}(F)\)1; for a spectrally normalized linear layer at \(q = 1\), \(L \le \left\|W\right\|_{\mathrm{op}}\).

At the OC optimum: costate bound via Danskin. At the OC optimum, the discrete adjoint \(\lambda_k\) of 8 evaluated along the optimal trajectory with \(v = v^\ast\) coincides with the discrete costate \[P_k := \nabla u_{k}(X_k^\ast),\]and the OC-optimal forward state bound coincides with Theorem 5 since \(v^\ast\) saturates 11 on the boundary of \(\mathcal{U}_{\mathrm{ad}}^q\). The asymmetry between forward and backward at the optimum lies in the costate, controlled by the input-magnitude exponent.

Theorem 7 (Lipschitz bound on the OC-optimal discrete costate). Fix \(q \in [0, 1]\). Let \(\left\{X_k^\ast\right\}_{k=0}^D\) be the discrete OC-optimal trajectory of 7 and \(P_k := \nabla u_{k}(X_k^\ast)\) the associated discrete costate, with terminal datum \(P_D = \nabla_X G(X_D^\ast, y)\) and \(G\) uniformly Lipschitz with constant \(\mathrm{Lip}(G)\). Then \(P_k\) is generated by the backward recursion \[P_k \;=\; \bigl(I + \Delta t\,\nabla_X v^\ast(X_k^\ast)\bigr)^{\!\top}\,P_{k+1}, \qquad k = D-1, \dots, 0, \label{eq:discrete-costate-recursion}\tag{31}\] and its norm obeys \[\left\|P_k\right\| \;\le\; \begin{cases} \mathrm{Lip}(G), & q = 0 \text{ (conserved costate)}, \\[4pt] \mathrm{Lip}(G)\,\bigl(1 + c\,\Delta t\,\left\|X^\ast\right\|_{L^\infty([0,T])}^{\,q-1}\bigr)^{D-k} & q \in (0, 1) \text{ (conditional geometric in D)}, \\[4pt] \mathrm{Lip}(G)\,\bigl(1 + c\,\Delta t\bigr)^{D-k}, & q = 1 \text{ (geometric in D)}, \end{cases} \label{eq:backward-optimal-bound}\tag{32}\] where \(\left\|X^\ast\right\|_{L^\infty([0,T])} := \sup_{k = 0, \dots, D} \left\|X_k^\ast\right\|\) is finite by Theorem 5.

Proof. Step 1: exact backward recursion. The Bellman recursion in Mayer form 14 is \(u_k(X) = u_{k+1}(X + \Delta t\, v_k^\ast(X))\) with terminal condition \(u_D = G\). Differentiating in \(X\) at \(X = X_k^\ast\) and using \(X_{k+1}^\ast = X_k^\ast + \Delta t\, v_k^\ast(X_k^\ast)\), \[\nabla u_k(X_k^\ast) = \bigl(I + \Delta t\, \nabla_X v_k^\ast(X_k^\ast)\bigr)^{\!\top} \nabla u_{k+1}(X_{k+1}^\ast),\] which is 31 under the definition \(P_k := \nabla u_k(X_k^\ast)\). At \(k = D\), the terminal condition \(u_D = G\) gives \(P_D = \nabla u_D(X_D^\ast) = \nabla_X G(X_D^\ast, y)\) as stated. The Jacobian \(\nabla_X v_k^\ast\) exists a.e. along the OC-optimal trajectory by the well-posedness of the OC problem (Proposition 3).

Step 2: Danskin’s envelope identity. The OC-optimal velocity \(v_k^\ast\) maximizes the Hamiltonian pointwise on \(\mathcal{U}_{\mathrm{ad}}^q(X_k^\ast)\) (Theorem 1). Danskin’s envelope identity 25 therefore gives \[\nabla_X v_k^\ast(X_k^\ast)^{\!\top} P_{k+1} \;=\; -\nabla_X H_q(X_k^\ast, P_{k+1}), \label{eq:danskin-envelope}\tag{33}\] where \(H_q(X, P) = (c\left\|X\right\|^q + b)\left\|P\right\|\) is the Hamiltonian of 16 . The right-hand side of 33 is a closed-form expression in \(X_k^\ast\) and \(P_{k+1}\) that bypasses realization-level derivatives of \(v_k^\ast\).

Step 3: costate-magnitude recursion. Combining 31 with 33 and taking norms, \[\left\|P_k\right\| \;\le\; \left\|P_{k+1}\right\| + \Delta t\, \left\|\nabla_X H_q(X_k^\ast, P_{k+1})\right\|.\] From the closed form \(H_q(X, P) = (c\left\|X\right\|^q + b)\left\|P\right\|\), we compute \(\nabla_X H_q(X, P) = c\,q\,\left\|X\right\|^{q-2}\, X\, \left\|P\right\|\) for \(q > 0\), so \[\left\|\nabla_X H_q(X_k^\ast, P_{k+1})\right\| \;\le\; c\, q\, \left\|X_k^\ast\right\|^{q-1}\, \left\|P_{k+1}\right\|. \label{eq:H-X-gradient-bound}\tag{34}\] At \(q = 0\), \(H_0(X, P) = b\left\|P\right\|\) is \(X\)-independent, so \(\nabla_X H_0 \equiv 0\) and 34 holds trivially with the right-hand side zero.

Step 4: discrete Grönwall. Combining the last two displays, \[\left\|P_k\right\| \;\le\; \bigl(1 + c\, q\, \Delta t\, \left\|X_k^\ast\right\|^{q-1}\bigr)\, \left\|P_{k+1}\right\|. \label{eq:costate-magnitude-recursion}\tag{35}\] Iterating backward from \(k = D\) with \(\left\|P_D\right\| \le \mathrm{Lip}(G)\), \[\left\|P_k\right\| \;\le\; \mathrm{Lip}(G)\,\prod_{j = k}^{D-1} \bigl(1 + c\, q\, \Delta t\, \left\|X_j^\ast\right\|^{q-1}\bigr). \label{eq:costate-product-bound}\tag{36}\] Bounding each factor by \(1 + c\, q\, \Delta t\, \left\|X^\ast\right\|_{L^\infty([0,T])}^{q-1}\), where the discrete sup-norm along the OC-optimal trajectory is finite by Theorem 5, gives the geometric bound 32 at \(q \in (0, 1]\). At \(q = 0\), 34 gives \(\left\|\nabla_X H_0\right\| \equiv 0\), and 35 becomes \(\left\|P_k\right\| \le \left\|P_{k+1}\right\|\), which iterates to \(\left\|P_k\right\| \le \mathrm{Lip}(G)\). ◻

Remark 8 (Sufficient conditions on the terminal cost \(G\)). Proposition 3 requires \(G\) bounded uniformly continuous (BUC) for HJB well-posedness of the value function \(u_y\); Theorem 7 requires \(G\) globally Lipschitz for the norm bound on the discrete costate \(P_k\). Both are sufficient conditions. Squared loss satisfies neither globally, but is continuous and locally Lipschitz on any bounded set. The forward state bound of Theorem 5 confines the trajectory to a ball \(B_R\), so \(G\) may be replaced without loss by any globally BUC extension of \(G|_{B_R}\): Proposition 3 then applies to the extension, and Theorem 7 applies with the local Lipschitz constant \(\mathrm{Lip}(G|_{B_R})\).

Theorem 7 establishes the discrete costate as well-defined along the OC-optimal trajectory, with a Lipschitz bound strictly sharper than 30 at the optimum: the conserved-costate result at \(q = 0\) and the exponential rate \(c\) at \(q = 1\) are obtained with no realization-level Lipschitz hypothesis on \(v\). Off the optimum, only 30 applies, with constant \(L\) supplied by the architecture.

4 Residual architecture design via input-magnitude exponents↩︎

Throughout Sections 2 and 3, we established the sublinear-growth condition \(q \le 1\) on the velocity fields of the discrete architecture — equivalently, on its residual blocks — through complementary ODE-existence and optimal-control arguments. Each block \(v_k\) is a composite of smaller architectural building blocks called architectural primitives or simply primitives, combined through operations such as sequential composition, parallel sum, and gating; verifying \(q_k \le 1\) directly on the composite is impractical when a single block can compose ten or more primitives. We therefore take a modular approach: Section 4.1 introduces the primitive set \(\mathcal{P}\), the input-magnitude exponent \(q(B)\) assigned to each \(B \in \mathcal{P}\), and five operations under which \(\mathcal{P}\) is closed and \(q\) transforms predictably, so that block-level \(q_k\) follows from per-primitive exponents (Table 1) propagated through these operations; Section 4.2 applies this machinery per-block and between-block to certify the sequence of exponents \(\left\{q_k\right\}\) across depth.

4.1 Input-magnitude exponents under operations on primitives↩︎

The construction applies within a single residual block: it computes the input-magnitude exponent \(q_k\) of the per-block velocity \(v_k\) in the residual block map 2 from the exponents of its constituent primitives. A primitive is a continuous map \(B\) between Euclidean spaces, with input and output dimensions determined by the architectural context: \(\mathbb{R}^d \to \mathbb{R}^d\) for primitives that act on a single token’s embedding (activations, normalizations), \(\mathbb{R}^{d \times n} \to \mathbb{R}^{d \times n}\) for primitives that act on the full sequence (self-attention, selective state-space models), and \(\mathbb{R}^d \to \mathbb{R}^{d'}\) for dimension-changing projections (Q/K/V projections in attention, up- and down-projections in an FFN). The polynomial-growth bound below applies uniformly via the Euclidean norm on the appropriate space.

Definition 2 (Architectural primitives, with the right function space at \(q \le 1\)). The set of primitives is \[\mathcal{P} := \left\{B : \mathbb{R}^{N_{\mathrm{in}}} \to \mathbb{R}^{N_{\mathrm{out}}} \,\big|\, B \text{ continuous, and } \exists\, q, c, b \ge 0 \text{ with } \left\|B(x)\right\| \le c\left\|x\right\|^q + b \text{ for all } x \in \mathbb{R}^{N_{\mathrm{in}}}\right\}, \label{eq:primitive-set}\tag{37}\] where \(N_{\mathrm{in}}, N_{\mathrm{out}}\) are the input and output dimensions determined by the architectural context, and \(\left\|\cdot\right\|\) is the Euclidean norm on the corresponding space (Frobenius norm when the dimension is a product \(d \times n\)). For \(B \in \mathcal{P}\), the input-magnitude exponent \(q(B)\) is the infimum of \(q \ge 0\) for which constants \(c, b \ge 0\) making the bound hold exist. The primitive \(B\) is sublinear-growth if \(q(B) \le 1\). We refer to the set of sublinear-growth primitives, \[\mathcal{P}^{\le 1} := \left\{B \in \mathcal{P} \,:\, q(B) \le 1\right\},\] as the sublinear-growth class, or equivalently the right function space, for residual architecture design.

Theorem 9 (Primitive operations and their arithmetic of exponents). Let \(B_1, B_2 \in \mathcal{P}\) with input-magnitude exponents \(q_1, q_2\), coefficients \(c_1, c_2\), and biases \(b_1, b_2\), and with compatible dimensions at the operation boundary. The set \(\mathcal{P}\) is closed under the following five operations, and the input-magnitude exponent of the result is given by:

  1. Sequential composition. \(q(B_1 \circ B_2) \le q_1 q_2\), with asymptotic coefficient \(c_1 c_2^{q_1}\) for large \(\left\|x\right\|\).

  2. Sum (parallel branches). \(q(B_1 + B_2) \le \max(q_1, q_2)\), with coefficient determined by the dominant exponent.

  3. Hadamard product (multiplicative coupling). \(q(B_1 \odot B_2) \le q_1 + q_2\) entrywise.

  4. Residual wrapping. \(q(I + \Delta t\, B) \le \max(1, q(B))\).

  5. Normalization wrapping. \(q(\mathrm{LN}\circ B) = 0\) regardless of \(q(B)\): an outer LN annihilates input-magnitude scaling.

Proof. (i) \(\left\|B_1(B_2(x))\right\| \le c_1\left\|B_2(x)\right\|^{q_1} + b_1 \le c_1(c_2\left\|x\right\|^{q_2} + b_2)^{q_1} + b_1\); for large \(\left\|x\right\|\) the leading term is \(c_1 c_2^{q_1} \left\|x\right\|^{q_1 q_2}\). (ii) \(\left\|B_1(x) + B_2(x)\right\| \le c_1\left\|x\right\|^{q_1} + c_2\left\|x\right\|^{q_2} + b_1 + b_2\); for large \(\left\|x\right\|\) the larger exponent dominates. (iii) Apply \(\left|(B_1 \odot B_2)(x)_i\right| \le c_1 c_2 \left\|x\right\|^{q_1 + q_2}\) entrywise, then aggregate to \(L^2\). (iv) \(\left\|(I + \Delta t B)(x)\right\| \le \left\|x\right\| + \Delta t(c\left\|x\right\|^q + b)\); for \(q \le 1\) this is \(O(\left\|x\right\|)\), for \(q > 1\) it is \(O(\left\|x\right\|^q)\). (v) Lemma 1 of [@kan2025stability]: \(\left\|\mathrm{LN}(z)\right\| \le \left|\beta\right| + \left\|\gamma\right\|\sqrt{d}\), independent of \(z\). ◻

The three operations that dominate practical design are (i), (ii), and (iii). Sequential composition multiplies exponents, so chains of \(q = 1\) primitives remain at \(q = 1\). Parallel addition takes the max, so adding a \(q = 0\) branch (skip connection) to a \(q = 1\) branch leaves the block at \(q = 1\). Multiplicative coupling adds exponents, so gating \(x\) by a function of \(x\) that itself has nonzero \(q\) pushes the block above the threshold; this is the mechanism by which selective state-space models such as Mamba acquire supercritical \(q\).

4.1.0.1 Examples of primitives.

Table 1 lists the input-magnitude exponent of architectural primitives commonly used in machine learning, derived either from Theorem 9 or by elementary calculation. Three primitives anchor the catalogue: linear layers and Lipschitz activations sit at \(q = 1\) and are the building blocks of standard transformers; layer normalization sits at \(q = 0\) and is the standard mechanism for collapsing \(q\); the native selective state-space model (Mamba) is the leading example of a supercritical primitive used in practice.

Table 1: Input-magnitude exponents of common architectural primitives. Coefficients are functions of the layer’s parameters; biases are omitted. Composition with these primitives follows Theorem [thm:composition]. Self-attention and multi-head attention [@vaswani2017attention] are multi-token primitives: with \(V = W_V X\), \(K = W_K X\), \(Q = W_Q X\) all linear in \(X\) and softmax saturating, the per-column output is controlled by the mixed norm \(\norm{X}_{\infty, 2} = \max_k \norm{X_{:,k}}_2\) rather than by \(\norm{X_{:,j}}_2\) alone, with \(\norm{B(X)_{:,j}}_2 \le \norm{W_V}_\mathrm{op} \norm{X}_{\infty, 2}\). The \(q = 1\) entry should be read in this worst-token sense; the implication \(\norm{X}_{\infty, 2} \le \norm{X}_F\) then puts attention in \(\Uad^1\) with a possibly looser Frobenius coefficient. The \(q = 0\) entries are realized by two structurally distinct mechanisms (magnitude stripping, saturation).
Primitive \(q\) Coefficient \(c\)
Saturating activation (tanh, sigmoid, ReLU6 [@howard2017mobilenets]) \(0\) \(M\sqrt{d}\) (\(M\) = saturation level)
Dynamic Tanh \(\mathrm{DyT}(x) = \gamma \cdot \tanh(\alpha x) + \beta\) [@zhu2025dyt] \(0\) \(\norm{\gamma}\sqrt{d} + \abs{\beta}\)
Bounded Hyperbolic Tanh (BHyT) [@bhyt2026] \(0\) data-driven bound
Softmax (row-stochastic, applied after a projection) [@bridle1990softmax] \(0\) \(1\)
LayerNorm \(\LN(x; \gamma, \beta)\) [@ba2016layernorm] \(0\) \(\abs{\beta} + \norm{\gamma}\sqrt{d}\)
RMSNorm \(\gamma \cdot x / \mathrm{RMS}(x)\) [@zhang2019rmsnorm] \(0\) \(\norm{\gamma}\sqrt{d}\)
L2/Scale normalization \(g \cdot x / \norm{x}\) [@nguyen2019scalenorm] \(0\) \(g\)
Affine layer \(x \mapsto Wx + b\) \(1\) \(\norm{W}_{\mathrm{op}}\)
Lipschitz activation (ReLU, GELU [@hendrycks2016gelu], SiLU [@elfwing2018silu]) \(1\) \(\Lip(\sigma)\)
MLP \(x \mapsto W_L\sigma_L(W_{L-1}\sigma_{L-1}(\cdots W_1 x))\) with Lipschitz \(\sigma_\ell\) \(1\) \(\prod_\ell \Lip(\sigma_\ell)\norm{W_\ell}_{\mathrm{op}}\)
Self-attention \(\mathrm{softmax}(Q^\top K / \sqrt{d_k})\,V\) [@vaswani2017attention] \(1\) \(\norm{W_V}_{\mathrm{op}}\)
Multi-head attention \(\sum_h W_O^{(h)} \mathrm{Attn}^{(h)}(X)\) [@vaswani2017attention] \(1\) \(\sum_h \norm{W_O^{(h)}}_{\mathrm{op}}\norm{W_V^{(h)}}_{\mathrm{op}}\)
Polynomial activation \(x \mapsto x^p\) (entrywise, \(p \ge 2\)) \(p\) \(1\)
Selective SSM (Mamba, native) [@gu2024mamba] \(5\) parameter-dependent

The \(q = 0\) rows in Table 1 are not equivalent at the inductive-bias level, even though they share the algebraic slot. LN and RMSNorm produce outputs on a fixed manifold (an ellipsoid or sphere); saturating activations like tanh produce outputs in a bounded hypercube. All these mechanisms yield \(\left\|B(x)\right\| \le c\) uniformly in input magnitude, but the geometric structure of the bounded output set differs across them. Theorem 9 treats them as interchangeable at the level of \(q\); distinguishing them requires geometric structure beyond the exponent.

4.1.0.2 The layer-normalization family.

Four entries of the \(q = 0\) block in Table 1 — LayerNorm [@ba2016layernorm], RMSNorm [@zhang2019rmsnorm], ScaleNorm [@nguyen2019scalenorm], and L2 normalization — form the layer-normalization family. They share a single algebraic mechanism: division of the input by some statistic of itself (a norm, an RMS, a standard deviation), which strips magnitude and places the output on a fixed sphere or ellipsoid determined by the affine wrapper. The four variants differ in which statistic divides and which affine wrapper is applied, not in their input-magnitude exponent: all four sit at \(q = 0\) with coefficient determined by the wrapper parameters alone (Table 1). Under the sublinear-growth principle they are therefore a single design slot: any one of them inserted into a residual block collapses the block exponent to \(q = 0\) by Theorem 9(v), independently of where inside the block it appears. The full \(q = 0\) analysis of Section 3.2 then transfers to the family wholesale: the forward state stays uniformly bounded (Theorem 5, \(q = 0\) branch) and the OC-optimal costate is Lipschitz with no depth-amplification factor (Theorem 7, \(q = 0\) branch).

Remark 10 (Scope of \(\mathcal{P}\)). Table 1 shows that \(\mathcal{P}\) covers standard architectural primitives used in machine learning across single-token, sequence-level, and dimension-changing categories. The polynomial-growth bound accommodates the entire range of practically relevant growth rates — bounded (\(q = 0\)), linear (\(q = 1\)), and supercritical polynomial (\(q > 1\)) — and the rules of Theorem 9 apply uniformly across them. Super-polynomial growth (exponential or worse) is rare in machine-learning practice (standard activations are Lipschitz or saturating; weight layers are linear); when such a mapping appears, composition with a \(q = 0\) growth controller (LN, RMSNorm, saturating activation) returns the composite to \(\mathcal{P}\) by Theorem 9(v). The continuity requirement in Definition 2 excludes discontinuous primitives such as vector quantization [@vandenoord2017vqvae]; in practice these are accompanied by a smooth surrogate (straight-through estimator, Gumbel-softmax) at training time, and the effective primitive seen by the velocity field is the surrogate, which is continuous and lies in \(\mathcal{P}\).

4.2 Per-block and between-block certification procedure↩︎

The analysis above suggests a two-scale certification procedure for residual architecture design. Per block, Theorem 9 and Table 1 certify \(q_k \le 1\) from the constituent primitives, inheriting the well-posedness and stability of Sections 23. Between blocks, Remark 6 motivates organizing the exponents \(\left\{q_k\right\}\) across depth into a mixed-exponent residual stack.

4.2.0.1 Per-block certification procedure.

The primitive operations of Theorem 9 and the catalogue of Table 1 together yield a four-step certification procedure for a candidate residual block:

  1. Decompose the candidate block into primitives and read their input-magnitude exponents from Table 1.

  2. Apply the arithmetic of exponents in Theorem 9 to compute the block-level \(q\).

  3. If \(q > 1\), identify the highest-\(q\) primitive and apply one of two remedies: (a) wrap it with normalization (Theorem 9(v) collapses to \(q = 0\)); (b) replace it with a Lipschitz substitute (drops to \(q = 1\)).

  4. Once \(q \le 1\), the block satisfies the sublinear-growth bound 1 with coefficient \(c\) determined by the products and sums of the primitive coefficients, and inherits the well-posedness and stability guarantees of Theorems 5 and 7.

The certification procedure is conservative: the operations of Theorem 9 give upper bounds on the block-level \(q\), so a block certified by this procedure satisfies the magnitude bound, but the converse can fail—a block whose composed \(q\) exceeds \(1\) may still admit a tighter direct bound. The mathematical content of certification at the block level is therefore the existence of a primitive decomposition under which the operations of Theorem 9 suffice; the design step is identifying such a decomposition, with primitives realizing the same \(q\) through different inductive biases providing an orthogonal degree of freedom.

4.2.0.2 Between-block certification procedure.

A mixed-exponent architecture has per-block exponent \(q_k\) varying with block index \(k\) within the admissible range \(q_k \in [0, 1]\). This freedom arises naturally when the residual stack is heterogeneous: attention blocks, MLPs, gated state-space modules, and normalization layers can have different block-level exponents even when each is sublinear-growth admissible. The between-block design step certifies that the sequence \(\left\{q_k\right\}\) satisfies \(\max_k q_k \le 1\), which is the condition for the per-block certification procedure to apply at every depth.

Between-block design exposes two degrees of freedom that uniform-\(q\) analysis collapses. First, the fraction of \(q = 1\) blocks, \(|S_1|/D\) with \(S_1 := \left\{k : q_k = 1\right\}\), controls forward amplification of the state magnitude. Increasing the share of \(q = 0\) blocks (e.g., inserting normalization layers between \(q = 1\) blocks) lowers this fraction. Second, the depth-averaged Lipschitz constant \(\bar L := (1/D)\sum_{k=0}^{D-1} L_k\), where \(L_k\) is the operator-norm bound on \(\nabla_X v_k\) supplied by block \(k\), controls backward amplification of the Jacobian. The two degrees of freedom are chosen separately at design time, subject to \(\max_k q_k \le 1\).

Remark 11 (Forward and backward stability on mixed-exponent stacks). The per-block bounds of Theorem 5 and 30 apply block-by-block to a mixed-exponent stack with \(q_k \in \left\{0, 1\right\}\) and per-block constants \((c_k, b_k)\). For the forward state, \[\left\|X_D\right\| \;\le\; \prod_{k \in S_1}(1 + c_k\Delta t)\,\left\|X_0\right\| + \sum_{k=0}^{D-1} b_k\,\Delta t \prod_{\substack{j > k \\ j \in S_1}} (1 + c_j\Delta t), \label{eq:mixed-forward-bound}\tag{38}\] and for the cumulative Jacobian under per-block Lipschitz hypotheses \(\left\|\nabla_X v_k(X_k;\theta_k)\right\|_2 \le L_k\), \[\left\|J_{i:D}\right\|_2 \;\le\; \prod_{k=i+1}^{D}(1 + L_k\Delta t) \;\le\; \exp\!\Bigl(\sum_{k=i+1}^{D} L_k\,\Delta t\Bigr). \label{eq:mixed-backward-bound}\tag{39}\] Under residual scaling \(\Delta t = T/D\), the forward multiplicative factor reduces to \(\exp(\bar c\, T \cdot |S_1|/D)\) with \(\bar c := (1/|S_1|)\sum_{k \in S_1} c_k\), and the backward factor reduces to \(\exp(\bar L\, T)\). Both bounds are independent of depth \(D\) under residual scaling, with depth-dependence absorbed into the two architectural parameters \(|S_1|/D\) and \(\bar L\) identified above. The bounds use the additive residual update 2 ; architectures whose block map departs from this update—canonically Post-LN, \(X_{k+1} = \mathrm{LN}(X_k + \Delta t\, f(X_k))\), where an outer normalization resets the state at each layer—lie outside the present analysis.

Remark 12 (Strict interior \(q \in (0, 1)\) as an underexplored design space). The admissible range \(q_k \in [0, 1]\) contains a strict interior \(q \in (0, 1)\) that no entry in Table 1 occupies: the catalogue’s \(q \le 1\) primitives sit at the boundary cases \(q = 0\) (saturating activations, normalizations) and \(q = 1\) (linear layers, Lipschitz activations, attention, FFN). The forward bound 29 at \(q \in (0, 1)\) is polynomial in \(T\) rather than exponential, strictly sharper than the \(q = 1\) bound; an architectural primitive at strict interior \(q\) would inherit this sharper bound. No standard primitive currently realizes this regime, and the question of whether a useful \(q \in (0, 1)\) primitive exists — one that interpolates between the boundedness of \(q = 0\) and the magnitude-preservation of \(q = 1\) without sacrificing expressivity — is an open design question.

5 Experiments↩︎

The theoretical analysis of Sections 23 identifies the input-magnitude exponent \(q\) as the criterion for training stability: every block with \(q \le 1\) admits a globally defined forward flow and OC-well-posed backward dynamics, while supercritical admissible classes (\(q > 1\)) contain velocity fields that blow up in finite time. This criterion is independent of how a block reaches \(q \le 1\) — whether by composition with a normalization layer or by structural design. We test the prediction empirically in this section.

5.1 Time-series forecasting (TSF) at shallow depths↩︎

Modern time-series forecasting (TSF) models share an architectural template: a residual stack of identical blocks, each composed from a small set of primitives — a normalization layer (RMSNorm), a sequence backbone (selective state-space or self-attention) that carries the task-specific inductive bias, and a position-wise MLP. Both backbones we evaluate adopt this template with RMSNorm as the conventional normalization choice. We use the standard TSF benchmark setting: Weather [@wu2021autoformer] and ETTm1 [@zhou2021informer], input window length \(96\), prediction horizons \(H \in \{96, 192, 336, 720\}\), residual depths \(L \in \{3, 8\}\), and three seeds per Weather configuration and six per ETTm1 configuration, in the small-model regime (\(\le 2\)M parameters).

The two backbones differ in the input-magnitude exponent of the un-normalized block. Mamba [@gu2024mamba; @cmamba2024] uses a selective state-space operator whose multiplicative coupling between the input and the discretized state-transition matrix drives the block-level exponent to \(q = 5\) (Table 1; Figure 2, left panel) — supercritical when no normalization is present. PatchTST [@nie2023patchtst] is based on self-attention; its primitives are all \(q = 1\) by Table 1, so the block sits at \(q = 1\) even with the RMSNorm wrapper removed. For Mamba we evaluate five variants spanning the three regimes \(q = 0\), \(q = 1\), and supercritical \(q = 5\): the conventional RMSNorm baseline, Pre-LN, and Peri-LN (two normalization choices that each collapse the block to \(q = 0\)); a structural linear-growth modification (described below) that reduces the block to \(q = 1\) without any normalization layer; and a no-normalization control at the native \(q = 5\). The no-normalization control is called free velocity because removing the normalization leaves the block outside any sublinear-growth class, no longer bounded to \(q \le 1\). For PatchTST we evaluate the same four normalization choices (RMSNorm baseline, Pre-LN, Peri-LN, free velocity); the linear-growth modification does not apply since the block is already \(q = 1\) by construction.

The Mamba block provides a concrete demonstration of how the arithmetic of Theorem 9 certifies (or fails to certify) a block-level exponent, primitive by primitive. The selective state-space operator updates a state \(h_t \in \mathbb{R}^d\) along the token axis via \[h_t = \bar A_t(x_t)\, h_{t-1} + \bar B_t(x_t)\, x_t, \qquad y_t = C_t(x_t)\, h_t + D_t(x_t)\, x_t, \label{eq:mamba-update}\tag{40}\] with selective parameters \(\bar A_t, \bar B_t, C_t, D_t\) computed from \(x_t\) through linear projections, SiLU gating, and discretization. Under the standing Hurwitz assumption on the state matrix \(A\), the discretized transition \(\bar A_t = \exp(\Delta_t A)\) is a contraction (\(\left\|\bar A_t\right\| \le 1\), \(q = 0\)); without this assumption \(\bar A_t\) is not polynomially bounded and the block falls outside \(\mathcal{P}\). Applying Theorem 9 to 40 then accumulates the exponent as follows: \(\Delta_t, B_t, C_t\) are each linear in \(x_t\) (\(q = 1\)); \(\bar B_t = \Delta_t B_t\) is a Hadamard product of two \(q = 1\) factors, so Theorem 9(iii) gives \(q(\bar B_t) = 2\); the recurrence contribution \(\bar B_t \odot x_t\) raises the propagated state to \(q = 3\); the output projection \(y_t = C_t h_t\) multiplies a \(q = 1\) factor with the \(q = 3\) state, giving \(q = 4\); and the SiLU-gated output \(\mathrm{SiLU}(z_t) \odot y_t\) multiplies a \(q = 1\) gate with the \(q = 4\) output, producing the block-level exponent \(q = 5\) (Table 1, Figure 2, left panel).

The linear-growth modification (Figure 2, right panel) targets two contributions in this accumulation. First, the input \(x_t\) to the selective-parameter branch is replaced with its \(L^2\)-normalized version \(x_t / \left\|x_t\right\|\), capping the parameter path at \(q = 0\). Second, the SiLU gate \(x \mapsto x \cdot \sigma(x)\) is replaced with sigmoid gating \(x \mapsto \sigma(x)\), removing a \(q = 2\) contribution. The residual stream is untouched and inherits \(q = 1\) from the linear contribution. The composite block thereby reaches \(q = 1\) without any normalization layer.

Figure 2: Computational graphs of the Mamba velocity field. Left: the Mamba block with the normalization removed (the free-velocity variant); SiLU gating and multiplicative coupling between x_t and the discretized selective parameters drive the block-level exponent to q = 5. Right: the linear-growth modification, with L^2-normalization of the input on the selective-parameter path and sigmoid gating in place of SiLU (blue boxes); both modifications are localized to the parameter-computation branch and reduce the block to q = 1 without any normalization layer in the residual stream.

We report two types of result. The first, in Table 2, is the blow-up count per configuration cell. A run is a blow-up if its validation loss is nonfinite (NaN or \(\pm\infty\)) by the end of training, or if it terminates with a finite MSE many orders of magnitude above the target’s dynamic range. For Mamba, the four \(q \le 1\) variants (RMSNorm baseline, Pre-LN, Peri-LN, and the linear-growth modification) have no blow-ups in any of the \(288\) seed-level runs across their \(64\) configuration cells. The free-velocity Mamba (\(q = 5\)) diverges in \(48\) of its \(72\) runs, with the failure rate rising sharply with depth: \(12/24\) on ETTm1 and \(2/12\) on Weather at \(L = 3\); \(24/24\) on ETTm1 and \(10/12\) on Weather at \(L = 8\). For PatchTST, all four variants are completely stable across all \(288\) seed-level runs in their \(64\) cells, including the free-velocity variant, which sits at \(q = 1\) by primitive composition (Table 1) rather than at the supercritical exponent that the Mamba free-velocity variant exhibits.

Table 2: Blow-up counts per configuration cell on ETTm1 (\(6\) seeds per cell) and Weather (\(3\) seeds per cell). The “\(\Sigma/24\)” column for ETTm1 aggregates over the four prediction horizons (out of \(24 = 6 \times 4\)); the “\(\Sigma/12\)” column for Weather aggregates correspondingly (out of \(12 = 3 \times 4\)). The four \(q \le 1\) Mamba variants are completely stable across all \(288\) runs in their \(64\) cells; all Mamba blow-ups occur in the free-velocity variant. All four PatchTST variants—including free velocity at \(q = 1\)—are completely stable across all \(288\) PatchTST runs. Two Mamba ETTm1 free-velocity runs at \(L = 3\), horizon \(192\) (seeds\(2024\) and \(2025\)) terminated with finite MSE values exceeding\(10^{12}\) — ten orders of magnitude above the standardized targetscale — and are counted as blow-ups.
ETTm1 Weather
5-9(lr)10-14 Model Variant \(q\) \(L\) \(96\) \(192\) \(336\) \(720\) \(\Sigma/24\) \(96\) \(192\) \(336\) \(720\) \(\Sigma/12\)
Mamba free velocity \(5\) 3 2 3 4 3 12 0 1 0 1 2
8 6 6 6 6 24 1 3 3 3 10
baseline (RMSNorm) \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
Pre-LN \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
Peri-LN \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
linear growth \(1\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
PatchTST free velocity \(1\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
baseline (RMSNorm) \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
Pre-LN \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0
Peri-LN \(0\) 3 0 0 0 0 0 0 0 0 0 0
8 0 0 0 0 0 0 0 0 0 0

4pt

The second result is forecasting accuracy. Figure 3 reports MSE for each Mamba variant, and Figure 4 reports the corresponding result for PatchTST. Each value is the signed percentage deviation of a variant’s mean MSE from the per-cell variant median, computed on the surviving (non-blow-up) seeds. For Mamba, the four \(q \le 1\) variants cluster within \(\le 1.0\%\) of the median on Weather and \(\le 1.5\%\) on ETTm1, with no variant dominating across cells; within this class, the linear-growth modification is neither systematically better nor worse than the two normalized variants at a margin that exceeds within-cell sampling variability. The free-velocity Mamba, on cells where at least one seed survives, is uniformly above the median, by as much as \(4.6\%\) MSE on the smallest Weather configuration. For PatchTST, the spread across all four variants is much tighter: within \(\le 0.05\%\) on Weather, where the four variants are effectively indistinguishable at the plot resolution, and \(\le 0.8\%\) on ETTm1, with the free-velocity variant matching the two normalized variants in both stability and accuracy.

a

b

Figure 3: Forecasting MSE on Weather (top) and ETTm1 (bottom) for the five Mamba variants. Each value is the signed percentage deviation of a variant’s mean MSE from the per-cell median across variants (\(0\% =\) median; negative \(=\) lower error). The “baseline” (blue) is the conventional Mamba \(+\) RMSNorm configuration. The four \(q \le 1\) variants cluster within \(\le 1.5\%\) of the median; the free-velocity Mamba (red) is marked \(\nabla\) at cells with at least one blown-up seed and \(\times\) at cells where all seeds blew up (three on Weather, six on ETTm1)..

a

b

Figure 4: Forecasting MSE on Weather (top) and ETTm1 (bottom) for the four PatchTST variants. Conventions follow Figure 3. The “baseline” (blue) is the conventional PatchTST \(+\) RMSNorm configuration; the free-velocity variant (red) sits at \(q = 1\) by primitive composition (Table 1) rather than at a supercritical exponent. All four variants are stable on every cell, and the spread of MSE across variants stays within \(\le 0.05\%\) of the median on Weather and \(\le 0.8\%\) on ETTm1..

These observations match the sublinear-growth principle’s central prediction in the shallow setting: \(q \le 1\) blocks train stably and \(q > 1\) blocks blow up, as Section 2.1 and Theorem 5 guarantee. The criterion is the input-magnitude exponent of the block, not the presence or placement of a normalization layer. Two manipulations on the two backbones support this. Removing the normalization from Mamba leaves the block supercritical (Table 1; Figure 2, left) and the free-velocity variant diverges, while the linear-growth modification on the same backbone restores \(q = 1\) without any normalization layer and trains stably. Removing the normalization from PatchTST leaves the block at \(q = 1\) by primitive composition rather than supercritical, so the principle predicts stability with or without normalization — which the PatchTST free-velocity runs confirm. The two backbones together dissociate the input-magnitude exponent from the presence of a normalization layer: Mamba isolates the case where removing normalization admits blow-up, and PatchTST isolates the case where it does not.

Together these results validate the stability prediction of Sections 23 in the shallow TSF setting.


  1. Throughout the paper, LN and RMSNorm are understood in their standard \(\varepsilon\)-regularized form, \(\mathrm{LN}(x) = \gamma \odot (x - \mu) / \sqrt{\sigma^2 + \varepsilon} + \beta\) for a small constant \(\varepsilon > 0\). The Lipschitz constant \(\mathrm{Lip}(\mathrm{LN})\) depends on \(\varepsilon\), but the \(q = 0\) magnitude bound in Table 1 holds uniformly in \(\varepsilon > 0\).↩︎