Learnable Nonlinear Superposition Model

• Learnable nonlinear superposition model is a representation method that replaces fixed gating laws with feature-specific nonlinear maps to generate activation coefficients.
• It improves feature utilization and reduces redundancy by enabling diverse activation regimes, as demonstrated in mechanistic weather modeling.
• The approach maintains a linear decoder for additive interpretability while employing trained nonlinear activations to capture regime-dependent behaviors.

A learnable nonlinear superposition model is a representation in which an object of interest is expressed as a superposition of components, while the coefficients or composition rule are themselves nonlinear and learned rather than fixed. In the mechanistic-interpretability formulation introduced for AI weather models, the decoder remains a linear superposition of feature directions, but each latent coefficient is produced by a learned per-feature nonlinear gate, yielding a “learnable nonlinear superposition model” of internal activations (Cheon, 17 May 2026). In adjacent literatures, closely related phrases denote Lie-type nonlinear superposition rules for differential systems, Kolmogorov-style superpositions of learned univariate functions, and learned nonlinear image-formation models, which suggests that the term names a family of structurally related ideas rather than a single canonical architecture (Treumova et al., 2024).

1. Definition and conceptual scope

In sparse representation learning, the standard linear superposition hypothesis assumes that an activation vector can be approximated as a sparse sum of feature directions plus a bias term,

$$\mathbf{x} \approx \sum_{j=1}^{n} \alpha_j \,\mathbf{d}_j + \mathbf{b}.$$

A baseline sparse autoencoder operationalizes this with a linear encoder projection, a shared fixed nonlinearity such as ReLU, and a linear decoder, so that reconstruction remains a linear combination of decoder columns (Cheon, 17 May 2026).

The learnable nonlinear variant keeps the additive decoder structure but relaxes the assumption that all features share the same gating law. In the KAN-SAE formulation, the representation is still

$$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$

but each coefficient is generated by a feature-specific learned nonlinear map,

$$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$

with $\phi_j$ a cubic B-spline. The superposition is therefore linear in representation space and nonlinear in the mapping from input activations to coefficient values (Cheon, 17 May 2026).

This distinction matters when the latent mechanisms are regime-dependent. The weather-transformer study motivating KAN-SAE argues that blocking anticyclones, jet structure, and tropical cyclones exhibit thresholding, saturation, and nonlinear regime transitions, whereas linear SAEs impose a shared ReLU gate, assume linear separability in pre-activation space, and often produce dead neurons, redundant feature clusters, or conflated mechanisms (Cheon, 17 May 2026).

A broader mathematical reading appears in the Lie-systems literature, where nonlinear superposition means that the general solution of a first-order ODE system can be written as a fixed nonlinear function of a finite set of particular solutions and constants. There the relevant class is characterized by the Lie–Vessiot–Guldberg theorem, and the “model” is learnable only in the sense that one may parameterize or verify dynamics within that finite-dimensional Lie class (Treumova et al., 2024). A plausible implication is that “learnable nonlinear superposition” has at least two non-equivalent technical meanings: one tied to learned feature coefficients, and another tied to exact nonlinear composition laws of trajectories.

2. Encoder-side nonlinear superposition in KAN-SAE

The baseline linear SAE used for comparison is

$$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$

with training loss

$$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$

This architecture assumes a shared fixed encoder nonlinearity and a purely linear decoder dictionary (Cheon, 17 May 2026).

KAN-SAE modifies only the encoder nonlinearity. Each latent dimension $j$ receives a learned B-spline activation: $$z_j = \phi_j\!\left([W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}]_j\right),$$ while the decoder is unchanged,

$$\hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}}.$$

Each $\phi_j$ is a cubic B-spline of order $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$0 with $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$1 control points,

$$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$2

where $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$3 is uniformly spaced over the empirical range of pre-activations and $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$4 are learnable coefficients (Cheon, 17 May 2026).

This architecture yields a precise notion of learnable nonlinear superposition. The decoder still expresses the reconstruction as an additive composition of directions $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$5, but the coefficients are no longer determined by a universal ReLU. Instead, each feature chooses its own activation regime: sharp or soft thresholds, saturation plateaus, asymmetric responses to positive and negative pre-activations, and even non-monotone profiles, although most observed profiles are convex or concave (Cheon, 17 May 2026).

The training objective is unchanged: $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$6 Sparsity is enforced purely through the $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$7 penalty; no top-$$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$8 or hard gating is used. Gradients through the B-splines use automatic differentiation over the de Boor recurrence (Cheon, 17 May 2026).

The specific implementation studied uses latent size $$\hat{\mathbf{x}} = \sum_j z_j\,\mathbf{d}_j + \mathbf{b}_{\mathrm{pre}},$$9 for $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$0, Adam with learning rate $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$1, $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$2, 100 epochs, and $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$3 annealed linearly from $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$4 to $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$5. Decoder columns are renormalized to unit $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$6 norm at each step; spline knots are initialized uniformly over the 1st–99th percentile of pre-activation, and control points are initialized to zero so that initial activations are near-linear or flat (Cheon, 17 May 2026).

3. Training behavior, sparsity, and empirical evidence

The weather-model application probes the residual stream at Layer 5 of the Step-2 path of Sonny, a hierarchical weather transformer trained on ERA5 for medium-range forecasting. Sonny operates at 1.5° resolution with 20.5M total parameters, and the probed activations are $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$7. SAE training uses 2016–2017 data and testing uses 2018–2022, totaling approximately 3.6 million tokens (Cheon, 17 May 2026).

The paper reports that KAN-SAE discovers 975 alive features out of 1024, whereas the linear baseline yields 566 alive features, corresponding to feature utilization of 95.2% versus 55.5% and dead feature rates of 4.8% versus 44.5% (Cheon, 17 May 2026). Alive status is defined through spline magnitude: $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$8 with $$z_j(\mathbf{x})=\phi_j\!\bigl(w_j^\top(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+b^{(j)}_{\mathrm{enc}}\bigr),$$9, chosen from the bimodal distribution of $\phi_j$0 (Cheon, 17 May 2026).

Redundancy is also reduced. Median pairwise $\phi_j$1 of EU-mean activation time series is 0.076 for KAN-SAE and 0.092 for the linear SAE, while the proportion of pairs with $\phi_j$2 is 4.4% versus 11.9% (Cheon, 17 May 2026). The interpretation given is that nonlinear gating allows features to specialize by activation regime instead of replicating shifted linear detectors.

Reconstruction fidelity remains comparable: explained variance on test is 84.2% for KAN-SAE and 86.6% for the linear baseline, while mean $\phi_j$3 norm is 111.4 versus 147.4 (Cheon, 17 May 2026). The paper therefore characterizes the nonlinear model as losing only about 2.4 EV points while using about 24% lower $\phi_j$4 activation and much more of its dictionary.

The learned activation shapes provide direct evidence that the encoder is not merely a linear SAE with a slightly altered threshold. Among the 975 alive features, 902 are convex and 73 concave, with median nonlinearity score 0.43. Many profiles deviate substantially from identity or ReLU (Cheon, 17 May 2026). This makes the nonlinear coefficient map itself an interpretable object: each $\phi_j$5 exposes the pre-activation ranges in which a feature is suppressed, amplified, or saturated.

4. Interpretability case studies and causal axes

A principal case study is feature F590, described as a European heatwave detector. It localizes strongly over western Europe ($\phi_j$6–$\phi_j$7N, $\phi_j$8–$\phi_j$9E), peaks on 27 June 2019 during a record European heatwave, and overlaps the positive Z500 anomaly of a blocking anticyclone. The best linear-SAE match is displaced by about $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$0, failing to track the heatwave center (Cheon, 17 May 2026).

The nonlinear aspect of F590 is explicit in its spline gate: the activation profile is strongly convex, stays near zero for typical pre-activations, and rises sharply only when pre-activations enter the blocking regime. The study argues that a fixed ReLU can encode only a single linear ramp after an offset threshold and cannot shape how fast activation grows or saturates within the positive domain (Cheon, 17 May 2026). This suggests that the learned nonlinear coefficient is functioning as a regime detector rather than a generic sparse code.

A second case study is feature F831, a western Pacific typhoon tracker. It activates over $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$1–$$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$2N, $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$3–$$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$4E and peaks on 14 Sept 2018 during Typhoon Mangkhut. Both KAN-SAE and the linear baseline find features within $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$5 of storm center, but KAN-SAE yields more spatially concentrated activation across typhoon days (Cheon, 17 May 2026). The paper interprets this as evidence that linear SAEs can capture strong localized vortices, whereas nonlinear gating becomes essential for subtler thresholded structures such as blocking and heatwaves.

The causal role of a learned feature is tested by steering along the decoder direction of F590: $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$6 with $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$7. Rolling Sonny forward after this intervention yields a warm T2m anomaly over western Europe, EU mean up to $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$8 K, local hotspots over $$\mathbf{z} = \mathrm{ReLU}\!\left(W_{\mathrm{enc}}(\mathbf{x}-\mathbf{b}_{\mathrm{pre}})+\mathbf{b}_{\mathrm{enc}}\right), \quad \hat{\mathbf{x}} = W_{\mathrm{dec}}\mathbf{z}+\mathbf{b}_{\mathrm{pre}},$$9 K, and an almost perfectly linear dose–response in $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$0 with $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$1 per day (Cheon, 17 May 2026). Z500 and T2m respond together, while MSLP decreases as expected under radiative heating. In the paper’s interpretation, the encoder’s learned nonlinear gate has identified a direction that functions as a causal intensity axis in the model’s dynamics.

The weather-transformer results motivate the paper’s strongest claim: nonlinear activations are essential for mechanistic interpretability of deep learning weather prediction models, because they recover physically meaningful climate features that remain invisible to linear baselines (Cheon, 17 May 2026).

5. Alternative formalizations in adjacent literatures

Outside sparse autoencoders, nonlinear superposition is formalized differently. In symbolic computation and Lie theory, a first-order ODE system admits a nonlinear superposition rule when its general solution can be written as

$$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$2

and the Lie–Vessiot–Guldberg theorem characterizes exactly those systems whose vector field is a time-dependent linear combination of vector fields spanning a finite-dimensional Lie algebra (Treumova et al., 2024). That paper develops an algorithm, based on Newton polytopes and iterated commutators, to decide whether a polynomial ODE lies in this class. A plausible implication is that one route to a learnable nonlinear superposition model is to restrict learned dynamics to polynomial Lie systems and verify finite-dimensional closure symbolically.

A second formalization comes from the Kolmogorov superposition theorem. The geometric KAN work recalls the classical decomposition

$$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$3

and interprets KANs as trainable models in which the inner and outer univariate functions are learned (Alesiani et al., 23 Feb 2025). That paper extends the construction to $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$4, $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$5, $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$6, and general $$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$7 settings, arguing that symmetry-aware physical models can still be written as finite superpositions of univariate nonlinearities applied to scalar invariants. Here, “learnable nonlinear superposition” refers not to nonlinear latent coefficients over a fixed dictionary, but to a network architecture whose multivariate function class is built from trainable univariate nonlinear maps and addition.

A third formulation appears in reflection separation, where the proposed Learnable Offset-Residual Superposition model writes an observed sRGB image as

$$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$8

$$\mathcal{L} = \|\mathbf{x}-\hat{\mathbf{x}}\|_2^2 + \lambda \|\mathbf{z}\|_1.$$9 is a learnable nonlinear residual capturing content-adaptive inter-layer coupling, and $j$0 is a learnable zero-order offset capturing ISP-induced bias (Hu et al., 1 Jun 2026). There the “superposition model” concerns the forward image-formation law rather than latent feature decomposition, but the same structural theme reappears: additive base terms supplemented by learned nonlinear interaction terms.

These usages are not interchangeable. The SAE/weather formulation is encoder-nonlinear and decoder-linear; the Lie-systems formulation concerns exact nonlinear composition of trajectories; the KAN/KST formulation concerns exact or approximate synthesis of multivariate functions from learned univariate nonlinearities; and the reflection-separation formulation concerns learnable nonlinear composition of physical layers (Cheon, 17 May 2026). This suggests that the most stable cross-domain characterization of the term is architectural rather than domain-specific: a superposition with learned nonlinear composition rules.

6. Limitations, misconceptions, and open directions

A common misconception is that “beyond linear superposition” means abandoning additive interpretability altogether. In KAN-SAE this is not the case: the decoder remains a linear dictionary, so feature contributions are still additive in representation space (Cheon, 17 May 2026). The change lies in how coefficients are generated. The distinction is operationally important because it preserves interpretable directions while expanding the class of admissible gating behaviors.

Another misconception is that nonlinear superposition automatically implies unrestricted expressivity. The weather paper keeps the decoder linear and the encoder nonlinear only through learned univariate spline gates (Cheon, 17 May 2026). The ODE verification paper is restricted to polynomial or Laurent-polynomial vector fields and does not compute the explicit superposition map $j$1 in general (Treumova et al., 2024). The reflection-separation paper argues that high-order fixed polynomials become numerically unstable even when they improve fit, motivating learned residual interaction models instead (Hu et al., 1 Jun 2026). These examples indicate that the point of nonlinear superposition is not maximal generality but structured relaxation of linear assumptions.

For weather-model interpretability, the central open question is whether nonlinear gating is merely useful or fundamentally necessary across architectures and layers. The empirical case studies, especially F590, are presented as evidence that certain atmospheric regimes require thresholded and convex activation laws that linear SAEs cannot recover (Cheon, 17 May 2026). A plausible implication is that future mechanistic-interpretability pipelines for scientific transformers may need to treat encoder nonlinearity as a first-class design variable rather than a minor architectural detail.

Across domains, current research points toward several directions already visible in the cited work: verifying nonlinear superposition structure in learned dynamical systems (Treumova et al., 2024), imposing geometric equivariance within KAN-style superposition networks (Alesiani et al., 23 Feb 2025), and coupling nonlinear forward models to interactive inverse architectures in image decomposition (Hu et al., 1 Jun 2026). Taken together, these developments frame the learnable nonlinear superposition model as a general strategy for retaining compositional structure while replacing fixed linear or fixed-form interaction laws with trainable nonlinear ones.