One Architecture to Rule Them All: Separable Neural Networks as a Unified Primitive

Paper: 2603.12244 Authors: Reza T. Batley, Apurba Sarker, Rajib Mostakim, Andrew Klichine, Sourav Saha Categories: cs.LG, cs.AI

The Gap

Most real-world systems — fluid dynamics, language, robotic control, material microstructures — have a hidden gift: their high-dimensional behavior can often be factored into lower-dimensional interacting parts. A turbulent flow field isn’t a random soup; it has spatial modes that interact in structured ways. A sentence isn’t a bag of tokens; it has compositional dependencies with bounded interaction order.

The problem is that standard deep learning ignores this. MLPs, Transformers, and CNNs are monolithic: they throw all dimensions into a shared nonlinear blender and hope the network learns the structure implicitly. This works, but it’s wasteful — you’re paying the full parameter and compute cost of a general function approximator when the target function is actually much more structured.

Prior work has nibbled at this from different angles. Kolmogorov-Arnold Networks (KANs) exploit additive decompositions. Tensor networks (Tucker, CP decompositions) exploit low-rank multilinear structure. Quadratic networks capture pairwise interactions. But these are separate tools with separate literatures. Nobody had asked: is there a single representational class that contains all of these as special cases, and can that class be deployed uniformly across prediction and generation tasks?

Problem: High-dim systems have factorisable structure
         but monolithic nets ignore it
         |
         v
Assumption: Separability is a property of *coordinates*,
            not necessarily the system itself
         |
         v
Method: SNA -- constrain interaction order + tensor rank
        to unify additive / quadratic / tensor-decomposed models
         |
         v
Evidence: 4 domains (RL nav, microstructure gen,
          turbulence, language modeling)
         |
         v
Conclusion: SNA is a domain-agnostic primitive for
            both deterministic and distributional tasks

The Increment

One sentence: Before this paper, factorised neural models were domain-specific tools; after it, there’s a single architectural class that subsumes them and works across physics, language, and control.

Core Mechanism

An SNA is defined by two constraints on how a neural network computes its output: interaction order (how many input dimensions can jointly influence a single output component) and tensor rank (how many rank-1 terms approximate the full interaction tensor). Set interaction order to 1 and you get an additive model — each input dimension contributes independently. Set it to 2 and you get quadratic interactions — pairs of dimensions can interact. Allow higher orders with low-rank tensor decomposition and you get something like a Tucker or CP network.

The key insight is that separability is coordinate-dependent. A chaotic PDE like Navier-Stokes looks hopelessly entangled in physical space, but if you embed its states into a learned smooth latent space, the dynamics can become approximately separable. This is the move that lets SNAs handle turbulence: don’t model the raw vorticity field, model it in a coordinate system where the separability assumption actually holds.

For generation (language, microstructures), the paper draws an analogy between autoregressive language modeling and distributional modeling of chaotic trajectories. Both involve predicting the next state given a history, and both can be cast as learning a conditional distribution over a separable embedding. This lets the same SNA formalism handle discrete token sequences and continuous physical states without architectural surgery.

Input x (high-dim)
  |
  v
Learned embedding phi(x)  <-- smooth, separable coords
  |
  v
+---------------------------+
| SNA core                  |
|  interaction order k      |
|  tensor rank r            |
|                           |
|  f(x) = sum_i  w_i *      |
|         prod_{j in S_i}   |
|         g_j(phi_j(x))     |
|                           |
|  k=1: additive            |
|  k=2: quadratic           |
|  k>2 + low rank: tensor   |
+---------------------------+
  |
  v
Output: deterministic prediction OR
        distributional sample (autoregressive)

Think of it like a mixing board in a recording studio. A monolithic network is a single giant effects unit — everything goes in, something comes out, and you have no idea which knob does what. An SNA is a proper mixing board: each channel (input dimension) has its own strip, and you can route specific channels to interact with each other at controlled points (interaction order), while keeping the overall signal path sparse and interpretable (tensor rank). The embedding step is like tuning each instrument before it hits the board — you’re finding the representation where the mixing rules are actually clean. Autoregressive generation is just playing the board forward in time, one measure at a time.

Key Concepts

• Interaction order: Imagine you’re modeling house prices. Order-1 means each feature (size, location, age) contributes independently — you just add them up. Order-2 means pairs can interact — maybe size × location matters more than either alone. Order-3 means triples interact. Most real systems don’t need full interaction order equal to the input dimension; they’re well-approximated by low-order interactions. SNA makes this a first-class architectural choice rather than something you hope the network learns.

• Tensor rank in neural networks: A full interaction tensor for k-way interactions over d dimensions has d^k entries — exponentially expensive. Low-rank decomposition says: approximate this tensor as a sum of r simple rank-1 terms (outer products of vectors). This is the same idea as matrix factorization (SVD), just generalized to higher-order tensors. Rank r controls the expressivity vs. cost tradeoff. SNA uses this to make high-order interactions tractable.

• Coordinate-aware separability: A function might look non-separable in one coordinate system and perfectly separable in another. The classic example: rotation. The function f(x,y) = x² + y² is separable (sum of independent terms). Rotate the axes and it stays separable. But f(x,y) = xy is not separable in Cartesian coordinates — yet in rotated coordinates (u = x+y, v = x-y), it becomes (u²-v²)/4, which is a difference of separable terms. SNAs learn the coordinate system (the embedding) jointly with the separable model, so the architecture can adapt to whatever representation makes the problem factorisable.

Framework Shift

Before (mainstream approach):        After (this paper):

Input                                Input
  |                                    |
  v                                    v
+----------+                        Learned embedding
| Monolithic|                       (separable coords)
| net       |                          |
| (MLP /    |                          v
| Transformer|                      +------------------+
| / CNN)    |                       | SNA              |
|           |                       | - interaction    |
| implicit  |                       |   order k        |
| structure |                       | - tensor rank r  |
| learning  |                       | - explicit factor|
+----------+                        +------------------+
  |                                    |
  v                                    v
Output                             Output
(prediction OR generation,         (same model handles
 separate architectures)            both)

Structure: hoped for               Structure: enforced
Cost: full                         Cost: rank-controlled
Domains: separate models           Domains: one primitive

From monolithic implicit learning to structured explicit factorisation, the core shift is: stop hoping the network discovers separability and instead build it in as a hard constraint, then learn the coordinates where that constraint is valid.

Expert Assessment

Problem choice: This is a real gap, not a manufactured one. The fragmentation between tensor networks, quadratic nets, KANs, and autoregressive models is genuine, and the observation that separability is coordinate-dependent rather than system-dependent is a genuinely useful reframing. That said, the field has been circling this territory for a while — the novelty is in the synthesis and the cross-domain demonstration, not in any single technical component.

Method maturity: The core idea is elegant rather than brute-force. The unification of additive/quadratic/tensor models under one constraint framework is clean. The weaker point is the embedding learning — the paper doesn’t give a fully principled account of when and why the learned coordinates will actually be separable. It works empirically, but the theoretical grounding for the embedding step feels underdeveloped.

Experimental integrity: Four domains is ambitious, and the breadth is the paper’s main selling point. The concern is depth: each domain gets a relatively thin treatment. For turbulence in particular, the comparison against deterministic baselines (showing SNA avoids nonphysical drift) is compelling, but the baselines for the generative quality of the distributional model are not rigorously specified. Language modeling results are presented without comparison to standard Transformer baselines at matched parameter counts, which is a significant omission.

Writing quality: The abstract and introduction do a good job of motivating the unification angle. The weakest section is the theoretical development of the SNA class itself — it reads like a compressed version of a longer derivation, and a reader without prior exposure to tensor decomposition will struggle. Expanding that section with concrete worked examples (not just equations) would make the paper substantially more accessible and would strengthen the claim that SNA is a “primitive” rather than a collection of tricks.

Verdict: weak accept — the unification idea is genuinely useful and the cross-domain demonstration is impressive, but the experimental comparisons are too thin in each individual domain to fully support the strong claims.

Takeaways

The most transferable idea here is the coordinate-aware separability move: if your system looks non-separable, don’t give up on structured models — instead, learn an embedding where separability holds, then apply the structured model there. This is a concrete design pattern you can apply whenever you have domain knowledge that the underlying dynamics are low-order but the raw features aren’t.

The second steal: the structural analogy between chaotic trajectory modeling and autoregressive language modeling. If you’re working on scientific ML for dynamical systems, framing your problem as “distributional autoregression over a smooth embedding” rather than “deterministic rollout” is a practical way to avoid the drift problem without switching to a completely different architecture class.

The interaction order hyperparameter is also worth keeping in your toolkit — it’s a cleaner way to express “I think pairwise interactions are sufficient” than hoping an MLP learns that implicitly.