跳转至

ControlJEPA: Principled Trajectory Regularization via Lyapunov Tube Loss

会议: NeurIPS 2026
arXiv: 2610.24951
代码: https://github.com/t-khanusa/ControlJEPA
领域: 其他
关键词: 隐状态轨迹正则化, 离散李雅普诺夫函数, 输入到状态稳定性 (ISS), 测地线假说, 自回归生成稳定性

一句话总结

本文提出基于离散时间输入到状态稳定性(ISS)理论的 ControlJEPA 框架,通过李雅普诺夫管道损失(Lyapunov Tube Loss)对规范化横向偏差施加有界的收缩约束,在不改变模型架构的前提下为语言模型隐状态轨迹提供显式几何管道控制,显著提升了推理性能与样本效率。

研究背景与动机

自回归语言模型在自回归下一词预测(Next-Token Prediction, NTP)训练范式下,针对每个 token 位置独立计算交叉熵损失,缺乏对隐状态序列全局动态几何演化的显式约束。基于测地线假说(Geodesic Hypothesis)的研究表明,表征空间中更平直、平滑的隐状态轨迹有助于降低一阶预测误差并增强动态可预测性。为此,近期涌现了诸如语义管道预测(Semantic Tube Prediction, STP)以及时间平直化(Temporal Straightening, TS)等轨迹正则化方法,试图通过惩罚随机采样三元组或连续潜速度之间的夹角余弦偏差来拉直隐状态轨迹。

然而,这类基于角度的正则化方案存在根本性缺陷:余弦相似度惩罚仅约束了状态位移的方向,却对横向偏离参考路径的实际欧氏幅度完全无约束。模型表征在训练或开环自回归生成过程中,横向偏差依然可能随时间步无限漂移发散。更为严峻的是,当前方法缺乏完备的控制理论与稳定性框架支撑,无法给出轨迹发散程度随超参数变化的显式数学边界,使得所谓“管道宽度”和“正则化强度”停留在启发式经验调参阶段。

针对这一核心矛盾,本文的切入角度是将语言模型逐层、逐 token 的离散步进动态建模为离散时间控制系统,并引入非线性控制中的实用输入到状态稳定性(Input-to-State Stability, ISS)理论。核心 idea:将隐状态相对于语义参考割线的横向位移构造成无量纲、尺度不变的李雅普诺夫候选函数,并设计单边松弛惩罚的李雅普诺夫管道损失(Lyapunov Tube Loss),在保留模型局部表示容量的偏 ISS(Partial-ISS)稳态区间内将隐状态严格约束在显式几何管道内。

方法详解

整体框架

ControlJEPA 无需修改底层 Transformer 架构,仅在训练阶段的标准 NTP 损失上增添李雅普诺夫管道正则项 \(\mathcal{L}_{\mathrm{Lyap}}\)。针对输入序列中的连续语义片段,系统首先通过起始语义锚点计算参考割线向量,将后续各 token 位置的隐藏状态位移正交分解为沿语义轴推进的纵向分量和偏离轴向的横向误差;随后利用割线长度对横向误差能量归一化,构建无量纲李雅普诺夫函数;最后通过单边软门控函数对违反离散 ISS 收缩条件的时间步施加梯度惩罚。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列与语义锚点<br/>h_start, h_end"] --> B["1. 正交横向分解<br/>提取参考割线与横向误差 e_t"]
    B --> C["2. 无量纲李雅普诺夫候选函数<br/>V_t = ||e_t||² / ||v_geo||² 尺度归一化"]
    C --> D["3. 偏 ISS 稳态收缩条件<br/>计算瞬时增量 ε_t = V_{t+1} - γ V_t - τ"]
    D --> E["4. 单边李雅普诺夫管道损失<br/>Softplus/ReLU 惩罚违约项"]
    E --> F["联合优化目标<br/>L_ControlJEPA = L_NTP + λ L_Lyap"]

关键设计

1. 正交横向分解:解耦语义推进与扰动漂移 为了建立轨迹稳定性的基准坐标,模型在输入文本的连续语义片段内选取起始隐状态 \(h_{t_{\mathrm{start}}}\) 和终止隐状态 \(h_{t_{\mathrm{end}}}\)(均为 Transformer 最后一层未经过输出反嵌入矩阵投射的向量)。两者之间的连线定义为参考割线 \(v_{\mathrm{geo}} = h_{t_{\mathrm{end}}} - h_{t_{\mathrm{start}}}\),其单位方向为 \(\hat{v} = v_{\mathrm{geo}} / \|v_{\mathrm{geo}}\|_{2}\)。对于片段内的任意中间时间步 \(t\),隐状态相对于起点的总位移为 \(d_{t} = h_{t} - h_{t_{\mathrm{start}}}\)。通过向 \(\hat{v}\) 投影,位移被严格解耦: $\(d_{t} = \langle d_{t}, \hat{v}\rangle \hat{v} + e_{t}, \quad \text{其中 } e_{t} = (I - \hat{v}\hat{v}^{\top}) d_{t}, \quad e_{t} \perp \hat{v}\)$ 其中纵向分量 \(\langle d_{t}, \hat{v}\rangle \hat{v}\) 承载了自回归推进任务的语义进展,而正交横向误差 \(e_{t}\) 纯粹反映轨迹远离理想语义路径的横向扰动漂移,其欧氏范数 \(\|e_{t}\|_{2}\) 直观对应轨迹管道的物理半径。

2. 无量纲李雅普诺夫候选函数:避免表征空间尺度坍缩与异方差失稳 若直接将未经缩放的绝对横向误差能量 \(\|e_{t}\|_{2}^{2}\) 作为李雅普诺夫函数,优化过程极易落入全局表征收缩的平庸解(即 \(h_{t} \to 0\) 且 \(\|v_{\mathrm{geo}}\|_{2} \to 0\)),导致表征秩退化并严重损害模型预测容量;而若按局部瞬时位移 \(\|d_{t}\|_{2}^{2}\) 归一化,由于位移随时间单调累积,分母动态漂移会引入严重的异方差性,导致训练动态剧烈振荡。为此,ControlJEPA 提出全局尺度归一化的无量纲李雅普诺夫候选函数: $\(V_{t} = \frac{\|e_{t}\|_{2}^{2}}{\|v_{\mathrm{geo}}\|_{2}^{2}}\)$ 该候选值非负且对全局线性缩放具备严格的不变性。当且仅当隐藏状态精确位于参考割线所在的仿射线上时 \(V_{t}=0\)。更重要的是,该度量允许模型在抑制相对横向偏离的同时,通过扩大端到端语义割线长度 \(\|v_{\mathrm{geo}}\|_{2}\) 来扩展表征流形的整体几何跨度。

3. 偏 ISS 稳态收缩:平衡测地线约束与分类容量 经典的完全输入到状态稳定性要求系统在每一步都满足指数收缩不等式 \(V_{t+1} \leq \gamma V_{t} + \tau\)(其中 \(\gamma \in (0, 1)\) 为收缩率,\(\tau \geq 0\) 为容许横向正交裕度)。然而,论文通过谱容量上界证明(Proposition 4.1),若将横向偏差硬性压缩至极限(\(\delta \to 0\)),反嵌入前的隐藏状态将被限制在二维仿射流形内,使输出 logits 矩阵的有效秩不超过 2,极度压缩词表预测概率分布的表达自由度。因此,ControlJEPA 明确主张工作在偏 ISS(Partial-ISS)体制,不强求条件在所有步百分之百满足,而是将其作为软性收缩惩罚,保留必要的正交表征容量。

损失函数 / 训练策略

定义瞬时符号扰动增量为 \(\varepsilon_{t} = V_{t+1} - \gamma V_{t} - \tau\)。当满足 ISS 条件时 \(\varepsilon_{t} \leq 0\),违反时 \(\varepsilon_{t} > 0\)。模型采用平滑单边激活函数 \(\sigma(x) = \log(1 + e^{x})\)(Softplus,也可取 ReLU)对有效语义区间 \(\mathcal{T} \subseteq \{t_{\mathrm{start}}+1, \ldots, t_{\mathrm{end}}-1\}\) 内的违约步进行平均惩罚: $\(\mathcal{L}_{\mathrm{Lyap}}(\theta) = \frac{1}{|\mathcal{T}|} \sum_{t \in \mathcal{T}} \log\left(1 + \exp\left(\frac{\|e_{t+1}\|_{2}^{2}}{\|v_{\mathrm{geo}}\|_{2}^{2}} - \gamma \frac{\|e_{t}\|_{2}^{2}}{\|v_{\mathrm{geo}}\|_{2}^{2}} - \tau\right)\right)\)$ 最终端到端联合微调目标为: $\(\mathcal{L}_{\mathrm{ControlJEPA}}(\theta) = \mathcal{L}_{\mathrm{NTP}}(\theta) + \lambda \mathcal{L}_{\mathrm{Lyap}}(\theta)\)$ 在理论分析上,定理 4.2 证明在偏 ISS 体制下,时间平均李雅普诺夫能量受到显式约束 \(\bar{V}_{T} \leq \frac{\tau + \mathcal{L}_{\mathrm{Lyap}}(\theta)}{1 - \gamma}\),时间方差 \(\mathrm{Var}_{t}[V_{t}]\) 在最大单步违约有界时呈现时域长度无关性(Horizon-independent),其实际横向管道宽度上界满足 \(\mathbb{E}_{t}[\|e_{t}\|_{2}] \leq \|v_{\mathrm{geo}}\|_{2} \sqrt{\frac{\tau + \mathcal{L}_{\mathrm{Lyap}}(\theta)}{1 - \gamma}}\)。

实验关键数据

主实验

实验在 NL-RX-SYNTH, NL-RX-TURK, GSM8K, Spider, NQ-Open 和 HellaSwag 六大基准数据集上进行,覆盖 Llama-3.2 (1B/3B/8B)、Gemma-2-2b-it、OpenELM-1.1B、OLMo-2-1B、Qwen3-1.7B 和 DeepSeek-R1-Distill-Qwen-1.5B 六大模型家族,基于双路 NVIDIA H100 GPU 完成测试。下表展示代表性基准及多模型族的主要对比数据:

评估基准 / 架构对比 评估指标 标准微调 (\(\mathcal{L}_{\mathrm{NTP}}\)) LLM-JEPA 语义管道预测 (STP) ControlJEPA (本文) 相对 STP 增益
GSM8K (数学推理) Accuracy (%) ~42.0 ~44.5 ~46.2 52.3 +6.1%
NQ-Open (开放问答) Exact Match (%) ~24.5 ~27.0 ~30.8 34.7 +3.9%
NL-RX-SYNTH (规则生成) Accuracy (%) ~76.2 ~78.0 ~81.5 86.8 +5.3%
NL-RX-TURK (自然语言正则) Accuracy (%) ~61.0 ~63.2 ~66.0 70.4 +4.4%
Spider (Text-to-SQL) Exact Match (%) ~58.2 ~60.1 ~63.4 67.1 +3.7%
HellaSwag (常识推理) Accuracy (%) ~68.5 ~70.2 ~72.0 74.6 +2.6%

消融实验

针对横向归一化方式、收缩率 \(\gamma\) 以及严格收缩对表征容量影响的消融结果(以 NL-RX-SYNTH 为例):

实验配置 / 归一化与超参消融 准确率 (%) 稳态横向误差 \(\|e_{t}\|\) 割线长度 \(\|v_{\mathrm{geo}}\|\) 现象与机理解释
ControlJEPA (最佳软收缩 \(\gamma=0.95, \lambda=0.05\)) 86.78 正向平衡稳定 显著扩张增长 维持适当正交裕度,流形语义跨度扩大,兼顾稳定与容量
激进强收缩 (\(\gamma=0.80\)) 28.70 趋近于 0 受到过度压制 逼近严格测地线,隐藏状态秩塌缩,严重破坏词表表达力
无归一化横向惩罚 (\(V_{t}=\|e_{t}\|_{2}^{2}\)) 6.30 趋近于 0 严重塌缩 (\(\to 0\)) 整体表征空间尺度坍缩,丢失特征区分度
瞬时位移归一化 (\(V_{t}=\|e_{t}\|_{2}^{2}/\|d_{t}\|_{2}^{2}\)) 64.20 波动剧烈 演化不稳定 分母随时间单调递增引入异方差不稳定,导致优化震荡

关键发现

  • 复杂推理任务增益最为显著:在 GSM8K 和 NQ-Open 等长上下文复杂推理任务中,ControlJEPA 相比标准 NTP 提升约 10 个百分点,显著高于短序列任务。原因在于长距离多步推理中,累积横向漂移是造成逻辑偏航的致命因素,显式管道约束在此场景发挥了最大效用。
  • 开环推理时隐式管道记忆留存:在测试阶段,李雅普诺夫损失完全关闭、模型处于开环自回归生成状态下,ControlJEPA 依然表现出极低且极平稳的横向误差方差,证明模型权重学习到的向量场内生性地重塑了表征流形吸引子,而非机械记忆训练路径。
  • 宽阔的超参数稳定区间:在正则化权重 \(\lambda \in [0.01, 0.08]\) 范围内,模型准确率保持稳定高位;收缩率 \(\gamma\) 展现出与任务语义复杂度的深刻关联:强句法规则任务(SYNTH)在较紧收缩(\(\gamma \in [0.84, 0.90]\))下最优,而高语义多样性任务(NQ-Open)则在较宽松管道(\(\gamma \geq 0.90\) 且较大 \(\tau\))下表现更佳。

亮点与洞察

  • 将现代控制理论引入隐状态轨迹约束:首次将离散时间输入到状态稳定性(ISS)与李雅普诺夫函数系统性地应用于自回归语言模型,摆脱了以往仅约束速度余弦夹角的启发式几何修补,建立了具有严格数学上界与可解释性的轨迹管道控制机制。
  • 洞察并避免表征维度塌缩的偏 ISS 设计:从理论上证明了将横向位移无脑降至零会导致反嵌入输出退化为秩-2 矩阵,清晰阐明了“为何不能过度拉直测地线”,从而确立了以松弛量 \(\tau\) 和偏 ISS 体制保留表征容量的方法论。
  • 无量纲全局归一化的优雅实用性:采用全局端到端割线长度 \(\|v_{\mathrm{geo}}\|_{2}\) 消除量纲,不仅杜绝了空间缩放塌缩和局部动态异方差波动,还促使模型在训练中自发拉长语义割线跨度以获取更高信噪比(SNR)。

局限与展望

  • 依赖预先给定的语义锚点:当前框架中的参考割线端点依赖数据集自带的语义片段切分,对于无明确断句标记的任意长文本流,如何动态、自适应地识别语义锚点尚待探索。
  • 理论保证暂限于训练流形内:现有的 ISS 收敛和方差控制理论主要建立在训练数据分布上,尽管测试期在经验上观测到了优异的流形重塑与管道保持,但分布外(OOD)或开环生成的严格泛化理论保证仍是开放问题。
  • 跨模态与解码期控制拓展:未来可进一步将 ISS 轨迹管道正则化推广到视觉自回归、视频生成及具身多模态模型中,甚至探索在推理阶段引入轻量李雅普诺夫引导机制以进一步纠正逻辑幻觉。

相关工作与启发

  • vs 语义管道预测 (Semantic Tube Prediction, STP):STP 仅惩罚隐状态位移向量之间的余弦夹角,只控方向不控幅度,隐状态容易在垂直方向漂移发散;ControlJEPA 采用李雅普诺夫收缩条件,直接对横向偏离物理幅度建立数学上界,并在推理时保持单调平稳。
  • vs 时间平直化 (Temporal Straightening, TS):TS 聚焦于连续潜速度之间的局部曲率平滑,属于局部无约束优化;ControlJEPA 引入端到端全局参考割线,结合 ISS 稳态理论实现了兼顾局部收缩与全局时空方差约束的轨迹控制。
  • vs 架构级控制论模型 (如 PIDformer, Stable Neural ODEs):传统控制论方法多通过对 Transformer 残差连接或权重正交性进行物理修改来实现稳定性;ControlJEPA 纯粹作为即插即用的辅助损失函数,零侵入模型架构,完全兼容现有预训练权重与高效微调管线。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将离散时间 ISS 理论与李雅普诺夫收缩引入 Transformer 隐状态轨迹控制,理论建构精巧。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 个主流基准、6 大模型家族(从 1B 到 8B)以及开环推理流形分析,消融与机理验证扎实。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨,控制理论概念到语言模型流形动态的映射清晰自洽,实验分析深刻透彻。
  • 价值: ⭐⭐⭐⭐⭐ 为大语言模型高效微调和表征几何规整提供了坚实的理论新支点,在提升推理能力与样本效率上潜力巨大。