Transformer as Optimal Control Theory

1. Introduction: Optimal Control Perspective

Transformers have achieved state-of-the-art performance across diverse domains including natural language processing, computer vision, program synthesis, and computational biology. However, practitioners often discover effective architectures through costly trial-and-error approaches, lacking a systematic, theory-driven foundation for design decisions.

This article presents a novel framework that studies Transformers through the lens of optimal control theory¹. The key insight is elegant:

Transformer训练可被形式化为最优控制问题，其中损失函数作为终端代价（terminal cost），模型深度对应时间（depth as time）。

这种视角带来了深远的影响：

从控制理论角度推导最优传输（Optimal Transport）正则化
提供可操作的训练和架构设计洞察
赋予泛化和鲁棒性理论保证

2. Continuous-Time Formulation

2.1 From Discrete to Continuous

传统Transformer通过离散的Transformer块序列处理输入：

$X_{i + 1} = f_{i + 1} (X_{i}), i = 0, 1, \dots, D - 1$

其中 $D$ 是模型深度， $f_{i}$ 是第 $i$ 个Transformer块的操作。

连续时间 formulation 将其重新解释为动力系统的离散化：

\frac{d X ( t )}{d t} = f (X (t), t), t \in [0, T], X (0) = X_{0} (1)

其中 $f$ 是Transformer块组成的复合函数 $f = f_{D} \circ \dots \circ f_{1}$ 。

2.2 Connection to Neural ODEs

这种 formulation 与 Neural ODEs 密切相关。关键区别在于：

特性	Neural ODE	OT-Transformer
速度场 $f$	任意神经网络	专用Transformer块
正则化	雅可比惩罚	最优传输正则化
应用场景	生成模型	分类、生成、序列建模

2.3 Physical Interpretation

我们可以将 Transformer 的信息流类比为物理系统：

$m \frac{d ^{2} x}{d t ^{2}} + γ \frac{d x}{d t} = f (x, t)$

其中：

$X (t)$ 是时间 $t$ 的隐状态
$f (X (t), t)$ 是注意力驱动的”力场”
$T$ 是终端时间（对应模型深度 $D$ ）

3. Loss as Terminal Cost, Depth as Time

3.1 Optimal Control Problem Formulation

在最优控制框架中，Transformer训练对应以下连续时间问题：

f min E_{(X_{0}, y)} [G (X (T), y) + \frac{λ}{2} \int_{0}^{T} ∥ f (X (t), t) ∥_{F}^{2} d t] (2)

其中：

$G (\cdot, y)$ 是 终端代价（terminal cost），即传统损失函数
$∥ f (X (t), t) ∥_{F}^{2}$ 是 运行代价（running cost），控制速度场的光滑性
$λ > 0$ 是正则化强度

3.2 Depth-Time Duality

这种 formulation 建立了 深度-时间对偶性：

优化控制视角	Transformer视角
时间 $t$	模型深度层
终端时间 $T$	总深度 $D$
轨迹 $X (t)$	逐层隐状态
速度场 $f$	Transformer块操作

3.3 Optimality Conditions

最优控制问题的KKT条件给出Hamilton-Jacobi-Bellman (HJB) 方程：

\frac{\partial V}{\partial t} + f min {\frac{λ}{2} ∥ f ∥_{F}^{2} + \nabla_{X} V \cdot f} = 0 (3)

其中 $V (X, t)$ 是值函数（value function）。

4. Optimal Transport as Regularization

4.1 Why Optimal Transport Emerges

检查最优性条件揭示了一个关键发现：最优传输（Optimal Transport）作为自然正则化项出现。

从Benamou-Brenier formulation¹，问题(2)等价格式为：

γ \in Γ (μ_{0}, μ_{1}) in f \int_{0}^{1} \int ∥ v (t, x) ∥^{2} d μ_{t} (x) d t (4)

其中 $μ_{t}$ 是 $X (t)$ 的分布， $Γ (μ_{0}, μ_{1})$ 是所有从 $μ_{0}$ 到 $μ_{1}$ 的传输计划。

4.2 OT Regularization Intuition

为什么 OT 正则化有效？

避免病态解：无正则化时，存在无穷多个最优速度场，其中一些高度不规则
控制轨迹光滑性：惩罚速度场范数确保隐状态轨迹平滑
数值稳定性：正则化使离散化问题更稳定

4.3 The Plug-and-Play Training Objective

OT-Transformer 的训练目标为：

θ, γ min E_{(X_{0}, y)} [G (X_{M}, y; γ) + \frac{λ Δ t}{2} m = 0 \sum M - 1 ∥ f (X_{m}, t_{m}; θ) ∥_{F}^{2}] (5)

约束于离散动力学：

X_{m + 1} = X_{m} + Δ t \cdot f (X_{m}, t_{m}; θ) (6)

其中 $Δ t = T / M$ 是时间步长。

5. Generalization Guarantees

5.1 Stable Forward Propagation

核心定理（稳定前向传播）¹：

对于输入-目标对 $(X_{1} (0), y_{1})$ 和 $(X_{2} (0), y_{2})$ ，模型输出 $\tilde{y}_{1}$ 和 $\tilde{y}_{2}$ 满足：

∥ \tilde{y}_{1} - \tilde{y}_{2} ∥_{2} \leq (1 - \frac{T L ^{2}}{λ})^{- 1} (L ∥ X_{1} (0) - X_{2} (0) ∥_{F} + \frac{T L ^{2}}{λ} ∥ y_{1} - y_{2} ∥_{2}) (7)

意义：模型输出是输入的 Lipschitz 连续函数。这意味着：

相似的输入 $\Rightarrow$ 相似的输出
输入扰动被均匀有界地放大

5.2 Distributional Robustness

定理3（分布鲁棒性）：

令 $T_{♯}$ 为训练后Transformer的推前算子， $W_{p} (\cdot, \cdot)$ 为 $p$ -Wasserstein距离。存在常数 $\hat{C} (p) > 0$ 使得：

W_{p} (T_{♯} μ, T_{♯} ν) \leq \hat{C} (p) L (1 - \frac{T L ^{2}}{λ})^{- 1} W_{p} (μ, ν) (8)

5.3 Generalization Bounds

结合分布鲁棒优化（DRO），我们获得 非渐近泛化界：

E_{test} [G] \leq E_{train} [G] + O (\frac{1}{n} + \frac{ϵ}{λ}) (9)

其中 $ϵ$ 是分布扰动半径。

6. Robustness Properties

6.1 Input Perturbation Robustness

OT正则化直接增强对输入扰动的鲁棒性：

扰动类型	机制	效果
噪声注入	抑制速度场范数	抑制扰动放大
对抗攻击	Lipschitz约束	有界对抗敏感性
缺失数据	平滑轨迹	插值能力强

6.2 Training Data Perturbation

定理3确保：在扰动数据上训练的模型在原始数据上保持良好性能：

Error_{original} \leq Error_{perturbed} + C \cdot ϵ (10)

6.3 Hyperparameter Selection

正则化参数 $λ$ 的选择原则：

下界： $λ > T L^{2}$ 确保稳定性
上界： $λ$ 过大导致输出对输入不敏感
实践：通过调整 $λ$ 或输出层权重衰减控制

7. Plug-and-Play Framework

7.1 Architecture Overview

Input X(0)
    ↓
┌─────────────────────────────────────────┐
│  Dynamical System                       │
│  dX/dt = f(X(t), t)                     │
│  where f = f_D ○ ... ○ f_1 (Transformer) │
│                                         │
│  + OT Regularization:                   │
│  λ ∫₀ᵀ ‖f(X(t), t)‖² dt                │
└─────────────────────────────────────────┘
    ↓
Output X(T) → ŷ

7.2 Implementation

# OT-Transformer: Plug-and-Play Implementation
import torch
import torch.nn as nn
 
class OTTransformer(nn.Module):
    """
    OT-Transformer: Optimal Transport Regularized Transformer
    """
    def __init__(self, base_transformer, lambda_ot=1.0, T=1.0, M=10):
        super().__init__()
        self.base_transformer = base_transformer
        self.lambda_ot = lambda_ot  # OT regularization strength
        self.T = T                   # Terminal time
        self.M = M                   # Number of integration steps
        self.dt = T / M
 
    def forward(self, x0):
        """
        Continuous-time forward pass
        """
        X = x0
        trajectory = [X]
        
        for m in range(self.M):
            # Compute velocity field using Transformer
            velocity = self.base_transformer(X)
            
            # Euler integration step
            X = X + self.dt * velocity
            trajectory.append(X)
        
        return X, trajectory
    
    def ot_regularization(self, trajectory):
        """
        Compute OT regularization term
        """
        reg = 0.0
        for m in range(len(trajectory) - 1):
            velocity = (trajectory[m+1] - trajectory[m]) / self.dt
            reg += torch.sum(velocity ** 2)
        return self.lambda_ot * self.dt * reg / 2

7.3 Training Loop

def train_ot_transformer(model, train_loader, optimizer, num_epochs):
    """
    Training loop for OT-Transformer
    """
    for epoch in range(num_epochs):
        total_loss = 0.0
        total_reg = 0.0
        
        for batch in train_loader:
            x, y = batch
            optimizer.zero_grad()
            
            # Forward pass
            output, trajectory = model(x)
            
            # Task loss (terminal cost)
            task_loss = nn.functional.cross_entropy(output, y)
            
            # OT regularization
            ot_reg = model.ot_regularization(trajectory)
            
            # Combined loss
            loss = task_loss + ot_reg
            
            # Backward pass
            loss.backward()
            optimizer.step()
            
            total_loss += task_loss.item()
            total_reg += ot_reg.item()
        
        print(f"Epoch {epoch}: Loss={total_loss:.4f}, OT-Reg={total_reg:.4f}")

7.4 Minimal Code Modifications

将现有Transformer转换为OT-Transformer只需：

# Before (Vanilla Transformer)
model = TransformerEncoder(vocab_size, d_model, nhead, nlayer)
trainer = Trainer(model, train_loader)
trainer.train()
 
# After (OT-Transformer) - minimal changes
from ot_transformer import OTTransformer
 
base_transformer = TransformerEncoder(vocab_size, d_model, nhead, nlayer)
model = OTTransformer(base_transformer, lambda_ot=1.0, T=1.0, M=10)
trainer = Trainer(model, train_loader)
trainer.train()  # Same training loop!

8. Experimental Results

8.1 nanoGPT on Shakespeare (Character-level)

任务：字符级文本生成

配置	参数量	测试损失
Baseline	10.65M	2.68 ± 0.006
OT-Transformer	6.16M	1.44 ± 0.005

关键发现：

46% 测试损失降低
42% 参数量减少
更好的泛化能力，更少过拟合

8.2 GPT-2 on Shakespeare (Word-level)

任务：词级文本生成

配置	参数量	测试损失
Baseline	123.7M	5.18 ± 0.032
OT-Transformer	123.7M	4.96 ± 0.012

改进：9.3% 测试损失降低

8.3 GPT-2 on OpenWebText (9B tokens)

大规模实验：验证OT-Transformer的可扩展性

配置	参数量	测试损失
Baseline	123.7M	3.21
OT-Transformer	123.7M	2.91

8.4 Additional Experiments

实验	Baseline	OT-Transformer	改进
Point Cloud (ModelNet40)	87.4%	89.9%	+2.5%
MNIST Classification	93.0%	97.1%	+4.1%
Cats & Dogs	77.6%	79.0%	+1.4%
Sentiment Analysis	83.9%	84.6%	+0.7%

8.5 Robustness Tests

噪声注入实验：

噪声水平	Baseline 损失	OT-Transformer 损失
0.0	2.68	1.44
0.05	3.65	1.95
0.1	4.60	2.42

观察：噪声水平越高，OT-Transformer的优势越明显。

9. Code Examples

9.1 Complete Training Script

#!/usr/bin/env python3
"""
OT-Transformer Training Script
Based on arXiv:2505.13499
"""
 
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from ot_transformer import OTTransformer
from transformer_model import TransformerLM
 
def create_ot_transformer(config):
    """
    Create an OT-Transformer from a base Transformer
    """
    # Load base transformer
    base_transformer = TransformerLM(
        vocab_size=config['vocab_size'],
        d_model=config['d_model'],
        nhead=config['nhead'],
        nlayer=config['nlayer']
    )
    
    # Wrap with OT framework
    model = OTTransformer(
        base_transformer=base_transformer,
        lambda_ot=config.get('lambda_ot', 1.0),
        T=config.get('T', 1.0),
        M=config.get('M', 10)  # Number of integration steps
    )
    
    return model
 
def train_step(model, batch, optimizer, device='cuda'):
    """
    Single training step with OT regularization
    """
    x, y = batch
    x, y = x.to(device), y.to(device)
    
    optimizer.zero_grad()
    
    # Forward pass through OT-Transformer
    output, trajectory = model(x)
    
    # Task loss (cross-entropy for classification/generation)
    task_loss = nn.functional.cross_entropy(
        output.view(-1, output.size(-1)), 
        y.view(-1)
    )
    
    # OT regularization term
    ot_reg = model.ot_regularization(trajectory)
    
    # Total loss
    loss = task_loss + ot_reg
    
    # Backward pass
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    
    return {
        'total_loss': loss.item(),
        'task_loss': task_loss.item(),
        'ot_reg': ot_reg.item()
    }
 
# Example configuration
config = {
    'vocab_size': 50257,      # GPT-2 vocab size
    'd_model': 768,          # Embedding dimension
    'nhead': 12,             # Number of attention heads
    'nlayer': 12,            # Number of layers
    'lambda_ot': 0.1,        # OT regularization strength
    'T': 1.0,                # Terminal time
    'M': 10                  # Integration steps
}

9.2 Sensitivity Analysis

def sensitivity_analysis():
    """
    Study effect of lambda_ot on performance
    """
    lambda_values = [0.01, 0.1, 1.0, 10.0, 100.0]
    results = []
    
    for lambda_ot in lambda_values:
        config['lambda_ot'] = lambda_ot
        model = create_ot_transformer(config)
        trainer = Trainer(model, train_loader, val_loader)
        
        history = trainer.train(epochs=50)
        val_loss = min(history['val_loss'])
        
        results.append({
            'lambda_ot': lambda_ot,
            'val_loss': val_loss,
            'test_loss': evaluate(model, test_loader)
        })
    
    return results

10. Theoretical Insights

10.1 Well-posedness of Training

定理1：优化问题(2)存在唯一解 当且仅当 $λ > 0$ 。

这揭示了：

无正则化（ $λ = 0$ ）训练问题是不适定的
存在无穷多最优速度场，其中一些高度不规则
OT正则化是保证训练良定义的关键

10.2 Connection to HJB Equations

核心数学创新在于利用 HJB偏微分方程的正则性理论 证明Transformer映射的稳定前向传播。

HJB方程：

λ Δ v - ∥\nabla v ∥^{2} = 0

其正则性理论直接给出Lipschitz常数的界。

10.3 Why Decoder-only Models Excel

理论分析表明，凸性假设对泛化保证至关重要：

架构类型	损失函数凸性	理论保证
Encoder-only	✓	✓
Decoder-only	✓	✓
Encoder-Decoder	✗	经验观察

这与工业界趋向decoder-only LLM的趋势一致！

11.1 Transformers as Continuous Dynamics

OT-Transformer扩展了将Transformer视为微分方程的视角：

方面	早期工作	OT-Transformer
视角	数学理论	最优控制
正则化	层归一化投影	OT正则化
保证	表达能力	泛化+鲁棒性

11.2 Relationship to Optimal Transport Theory

OT-Transformer与OT理论的深层联系：

Benamou-Brenier公式：OT距离作为动态规划问题
Wasserstein梯度流：训练动态理解为分布空间的梯度流
Sinkhorn算法：高效计算OT距离的近似方法

11.3 Connection to Energy-Based Models

OT正则化可视为能量函数的正则化：

E (X, t) = G (X (T), y) + \frac{λ}{2} \int_{0}^{T} ∥ f (X (t), t) ∥_{F}^{2} d t

12. Future Directions

12.1 Theoretical Extensions

Layer Normalization分析：将其纳入最优控制框架
注意力机制理论：解释softmax注意力的控制理论意义
更深层理论：超越凸性假设的分析

12.2 Practical Improvements

自适应时间步长：变步长ODE求解器
二阶方法：利用Hessian信息的训练算法
分布式训练：大规模OT-Transformer的高效实现

12.3 Applications

RL中的决策Transformer
Diffusion Models的连续时间统一
世界模型的控制理论视角

References

Last updated: 2026-05-03

Kan, K., Li, X., Zhang, B. J., Sahai, T., Osher, S., & Katsoulakis, M. A. (2025). Optimal Control for Transformer Architectures: Enhancing Generalization, Robustness and Efficiency. arXiv:2505.13499. [https://arxiv.org/abs/2505.13499] [OpenReview] ↩ ↩² ↩³

Metaphor

探索

Transformer as Optimal Control