可分离神经架构（SNA）：结构化归纳偏置统一预测与生成智能

一句话总结

SNA 用张量分解的视角重新审视神经网络设计：通过约束交互阶数和张量秩，将高维映射分解为低元可组合组件，在 RL 控制、混沌系统建模和语言生成中统一适用。

背景：单块架构悄悄浪费的结构信息

强化学习里有个被忽视的问题：我们用 MLP 拟合 $Q(s, a)$ 时，实际上让网络从零学习状态和动作的联合表示。但很多环境里，这个函数天然有可分离结构——状态价值 $V(s)$ 和优势 $A(s, a)$ 本质独立。

Dueling DQN 是第一个利用这个直觉的实用方案：

\[Q(s, a) = V(s) + A(s, a) - \frac{1}{|\mathcal{A}|}\sum_{a'} A(s, a')\]

但这只是加法分解，是最简单的一阶可分离性。现实中的结构更丰富：

加法分离：$f(x, y) = g(x) + h(y)$，零交叉项
双线性分离：$f(x, y) = \mathbf{u}(x)^\top \mathbf{v}(y)$，二阶交互
张量分离：高阶交互，用 CP 分解约束秩

这篇来自 arxiv 的论文（2603.12244）把这三类统一到一个框架——SNA（Separable Neural Architecture），并跨 RL 导航、微结构生成、湍流建模、语言模型四个领域验证了这一原语的通用性。

算法原理

直觉解释

最直接的类比：注意力机制本质上就是双线性分离。

\[\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d}}\right)V\]

其中 $QK^\top$ 就是 $q^\top k$——一个低秩双线性型。Transformer 用的正是二阶可分离性，只不过没有把它放进这个框架里说。

可分离性不一定是系统本身的性质，而常常是系统在某种坐标/表示下的涌现性质。这正是 SNA 的核心 insight：找到让函数”看起来可分离”的嵌入空间。

数学推导

SNA 的核心定义：给定 $n$ 个输入 $\mathbf{x}_1, \ldots, \mathbf{x}_n$，秩为 $R$ 的 SNA 输出为：

\[f(\mathbf{x}_1, \ldots, \mathbf{x}_n) = \sum_{r=1}^{R} \prod_{i=1}^{n} \phi_r^{(i)}(\mathbf{x}_i)\]

其中 $\phi_r^{(i)}: \mathcal{X}_i \to \mathbb{R}$ 是标量嵌入函数。这正是 CP 张量分解的神经版本。

三种特殊情况（秩约束从松到紧）：

\[\text{加法（无穷秩，无交叉）：} \quad f(x,y) = g(x) + h(y)\] \[\text{双线性（秩 } R \text{，二阶）：} \quad f(x,y) = \mathbf{u}(x)^\top \mathbf{v}(y)\] \[\text{三阶 SNA：} \quad f(x,y,z) = \sum_{r=1}^{R} \phi_r(x)\cdot\psi_r(y)\cdot\xi_r(z)\]

与 Bellman 方程的关系：

\[Q(s, a) = r(s, a) + \gamma\,\mathbb{E}_{s' \sim P(\cdot \mid s, a)}\!\left[\max_{a'} Q(s', a')\right]\]

如果转移概率 $P(s’ \mid s, a)$ 和奖励 $r(s, a)$ 都有低秩双线性结构，最优 $Q^*$ 也继承这个结构。SNA 提供了归纳偏置来显式利用这一点。

与其他架构的关系

架构	分离类型	秩约束	典型应用
Dueling DQN	加法（1+1）	—	离散控制
线性 Attention	双线性	$d_{head}$	序列建模
Tucker Q-net	Tucker 分解	$(R_1, R_2)$	大动作空间
SNA	统一框架	可调	通用原语

实现

核心 SNA 模块

import torch
import torch.nn as nn

class SeparableLayer(nn.Module):
    """
    双线性可分离层：f(x, y) = u(x)^T v(y)
    本质是低秩双线性型，秩 = rank
    """
    def __init__(self, dim_x: int, dim_y: int, rank: int):
        super().__init__()
        self.rank  = rank
        self.phi_x = nn.Linear(dim_x, rank)
        self.phi_y = nn.Linear(dim_y, rank)
        self.out   = nn.Linear(rank, 1)
        # 正交初始化，防止双线性层梯度消失
        nn.init.orthogonal_(self.phi_x.weight)
        nn.init.orthogonal_(self.phi_y.weight)

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        u = torch.relu(self.phi_x(x))         # [B, rank]
        v = torch.relu(self.phi_y(y))         # [B, rank]
        interaction = (u * v) / self.rank**0.5 # 缩放防爆炸
        return self.out(interaction)           # [B, 1]


class SNA(nn.Module):
    """
    完整 SNA Q 网络：加法项（V + A）+ 双线性交互项
    """
    def __init__(self, state_dim: int, action_dim: int,
                 rank: int = 32, hidden: int = 128):
        super().__init__()
        # 加法分支：无交叉项
        self.value_net = nn.Sequential(
            nn.Linear(state_dim, hidden), nn.ReLU(), nn.Linear(hidden, 1))
        self.adv_net = nn.Sequential(
            nn.Linear(action_dim, hidden), nn.ReLU(), nn.Linear(hidden, 1))
        # 双线性分支：捕捉状态-动作交互
        self.bilinear = SeparableLayer(state_dim, action_dim, rank)

    def forward(self, state: torch.Tensor,
                action: torch.Tensor) -> torch.Tensor:
        additive = self.value_net(state) + self.adv_net(action)
        cross     = self.bilinear(state, action)
        return additive + cross

SAC + SNA 训练循环（核心部分）

import torch.optim as optim

class SNAQAgent:
    def __init__(self, state_dim, action_dim, rank=32):
        self.q1    = SNA(state_dim, action_dim, rank)
        self.q2    = SNA(state_dim, action_dim, rank)
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 256), nn.ReLU(),
            nn.Linear(256, action_dim * 2)  # mean + log_std
        )
        self.q_opt  = optim.Adam(
            list(self.q1.parameters()) + list(self.q2.parameters()), lr=3e-4)
        self.pi_opt = optim.Adam(self.actor.parameters(), lr=3e-4)
        self.alpha  = 0.2

    def select_action(self, state: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            mean, log_std = self.actor(state).chunk(2, dim=-1)
            std  = log_std.clamp(-20, 2).exp()
            dist = torch.distributions.Normal(mean, std)
            return torch.tanh(dist.rsample())

    def update(self, s, a, r, s_, done):
        # Critic 更新
        with torch.no_grad():
            a_ = self.select_action(s_)
            target_q = r + 0.99 * (1 - done) * torch.min(
                self.q1(s_, a_), self.q2(s_, a_))

        loss = ((self.q1(s, a) - target_q)**2 +
                (self.q2(s, a) - target_q)**2).mean()

        self.q_opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(
            list(self.q1.parameters()) + list(self.q2.parameters()), 1.0)
        self.q_opt.step()
        return loss.item()

# ... (经验回放、Actor 更新、完整训练循环省略)

SNA 用于语言建模（结构类比）

论文最有趣的 insight：混沌时间序列与自回归语言模型在结构上同构——连续物理状态可视为光滑的可分离嵌入，用分布式建模处理混沌动力学，与语言模型预测下一个 token 的概率分布完全类比。

class SNALanguageModel(nn.Module):
    """
    用双线性层替换标准 QK 点积，实现线性时间复杂度的低秩注意力
    """
    def __init__(self, vocab_size, d_model, rank, seq_len):
        super().__init__()
        self.embed    = nn.Embedding(vocab_size, d_model)
        self.phi_q    = nn.Linear(d_model, rank)
        self.phi_k    = nn.Linear(d_model, rank)
        self.phi_v    = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        h = self.embed(x)                              # [B, T, d]
        Q = torch.relu(self.phi_q(h))                  # [B, T, R]
        K = torch.relu(self.phi_k(h))                  # [B, T, R]
        V = self.phi_v(h)                              # [B, T, d]
        # 利用可分离性：softmax(QK^T)V ≈ Q(K^T V)   O(TR) 而非 O(T²)
        KtV  = torch.einsum('btr,btd->brd', K, V)      # [B, R, d]
        attn = torch.einsum('btr,brd->btd', Q, KtV)    # [B, T, d]
        return self.out_proj(attn)

实验

环境选择

用 HalfCheetah-v4（状态 17 维，动作 6 维）测试 SNA vs 等参数 MLP Q 网络。选它的原因：

状态-动作空间维度适中，能看出分解效果而不被维度灾难淹没
连续动作天然适合双线性型建模
业界标准 benchmark，可与公开数字对比

rank 消融实验

# 测试不同 rank 对性能的影响
configs = [
    {"rank": 8,   "label": "SNA-rank8"},
    {"rank": 32,  "label": "SNA-rank32"},   # 通常是甜点
    {"rank": 128, "label": "SNA-rank128"},
    {"rank": None,"label": "MLP-baseline"},  # None = 纯 MLP
]
# 建议：5 个随机种子，报告均值±标准差
# ... (多种子实验代码省略)

SNA 相比等参数 MLP 的典型优势在 sample-limited 场景（前 500k 步）最明显。训练数据充足时，MLP 能靠容量弥补没有归纳偏置的缺陷。

调试指南

常见问题

1. 双线性分支学不动（加法分支主导，交互项趋近于 0）

原因：phi_x 和 phi_y 初始化相关性低，Hadamard 积方差接近 0。

# 诊断：检查各分支的梯度范数
for name, p in agent.q1.named_parameters():
    if p.grad is not None:
        print(f"{name:40s}: {p.grad.norm():.6f}")
# 如果 bilinear.* 的梯度比 value_net.* 小 100x，就是这个问题
# 修复：已在 SeparableLayer 中用正交初始化 + rank 缩放处理

2. Q 值爆炸（loss NaN）

Hadamard 积 $u \cdot v$ 的方差 $\approx \text{Var}(u) \cdot \text{Var}(v)$，随 rank 增大会爆炸。核心修复在 SeparableLayer 里的 / self.rank**0.5，另外梯度裁剪必须加（裁剪值 1.0）。

3. 如何验证模型”学到了可分离结构”

# SVD 检验双线性权重的有效秩
W = (agent.q1.bilinear.phi_x.weight.T
     @ agent.q1.bilinear.phi_y.weight)        # [rank, rank]
s = torch.linalg.svdvals(W)
cum = (s / s.sum()).cumsum(0)
print(f"90% 能量所需秩: {(cum < 0.9).sum().item() + 1}")
# 若有效秩远小于设定 rank，说明任务确实有可分离结构
# 若有效秩 ≈ rank，说明可分离假设不成立，换回 MLP

超参数敏感度

参数	推荐值	敏感度	说明
rank	16–64	中	从 32 开始；太小会欠拟合
lr	3e-4	高	双线性层比 MLP 更敏感，别用 1e-3
梯度裁剪	1.0	中	必须加，不加容易 NaN
隐层宽度	128–256	低	不是主要因素
初始化	正交	高	默认 kaiming 可能让双线性分支死亡

什么时候用 / 不用？

适用场景	不适用场景
连续控制，动作维度 > 6	低维离散动作（Dueling DQN 更简单）
状态-动作有天然分离结构（多关节机器人）	密集像素 → 动作的纯感知任务
样本效率要求高，数据昂贵	训练数据极其充足（MLP 追得上）
需要可解释性（分解后可分析各分支）	已有调好的 MLP baseline 且不想引入复杂度
物理模拟中的混沌动力学分布建模	大离散动作空间的 Q 网络（用 Tucker 更好）

我的观点

坦率说：SNA 不是银弹，是有条件有效的归纳偏置。

论文的核心价值在于提供了一个统一的理论语言：以前 Dueling DQN、线性 Attention、Tucker Q-net 是各自独立提出的工程技巧，SNA 把它们放进同一框架，解释清楚了”为什么这样分解有效”——这比单独的工程技巧更有解释力。

但有几个现实问题需要说清楚：

rank 必须调：每个领域最优 rank 不同，没有通用默认值，增加了一个调参维度。
实现比 MLP 脆：双线性层的梯度问题、缩放问题、初始化问题，都需要额外处理。如果你的 MLP 已经好用，引入 SNA 有工程成本。
混沌动力学 ↔ 语言建模的类比更多是理论直觉，实际上 turbulent flow 建模能否跑赢 FNO、Mamba 这些专门方法，需要更多实验支撑。

什么时候真的值得一试：你有多维连续动作空间，当前 MLP Q 函数在 sample-limited 情况下训练不稳定，或者你需要可解释的状态-动作分解。RL 本来就玄学，加一个有理论支撑的归纳偏置，至少方向是对的。