跳转至

Phaedra: Learning High-Fidelity Discrete Tokenization for the Physical Sciences

会议: NeurIPS2026(队列归属;全文为 arXiv v2,2026-09-29)
arXiv: 2602.03915
领域: 物理 / 科学计算
关键词: 离散分词器、形态–幅值分解、有限标量量化、偏微分方程、科学数据压缩

一句话总结

Phaedra 将物理场潜变量分为多维形态 token 与高精度一维幅值 token,再学习重组解码,使分布内 PDE 重建 nMAE 从 FSQ 的 2.603 降至 1.522,并把 65.0 GB 科学数据压缩为 3.44 GB。

研究背景与动机

把 PDE 解表示成离散 token,可以让科学数据接入交叉熵训练、掩码预测和 Transformer 序列建模,也能减少训练数据的存储与传输开销。但物理场不是普通灰度图:涡旋、激波的几何形状值得保留,其速度、密度、压力的绝对量级也不能随意改变。面向自然图像的感知损失允许重建出“看起来像”的纹理,而科学计算关心的是误差、梯度和频谱是否真实;高频细节的幻觉与过度平滑都可能破坏后续动力学预测。

困难还在于科学数据即便标准化,仍有较宽的动态范围和重尾分布。同一种局部形状可以出现在不同强度下,若用单个离散码同时刻画形状与幅值,就要把词表容量分给大量“同形不同强度”的组合。仅增加多尺度层级并不直接解决这个瓶颈:它区分粗细尺度,却没有专门为连续变化的物理幅值提供高密度坐标。

Phaedra 从经典形状–增益量化(shape-gain quantization)获得启发,但不指定解析基函数,也不强制把形状与增益相乘。它让量化结构本身诱导两条支路承担不同职责,再让卷积解码器学习两者如何协同。核心 idea:用多维离散词表保存可复用的局部形态,用密集的一维标量词表保存幅值,避免几何多样性与数值精度争夺同一个码的容量。

方法详解

整体框架

输入是一个二维网格上的单个物理变量;多变量场按通道独立处理,因此分词器实际输入通道数始终为 1。 共享卷积编码器先压缩空间分辨率,再依次进行双潜变量分解、分支式 FSQ 和学习式重组,最后由卷积解码器恢复原始网格。 标准配置把 \(128\times128\) 场压为 \(32\times32\) 潜网格,每个位置产生一个形态 token 和一个幅值 token,共 2048 个离散码。

形态分支有 8 个潜通道,幅值分支有 1 个潜通道;这不是把原始输入拆成两幅图,而是对编码后的潜特征做通道分割。 两条支路并行量化,解码时根据 token 查回各自量化向量,拼接后经通道混合卷积重组。 重建监督只在训练时使用,推理重建不需要原场标签,也不需要求解 PDE。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    X["单变量物理场"] --> E["卷积编码器"]
    E --> F["双潜变量分解"]
    F --> Q["分支式 FSQ<br/>形态 8D / 幅值 1D"]
    Q --> R["学习式重组"]
    R --> D["卷积解码器 / 重建场"]
    X -.->|仅训练:L1 重建监督| D
    F -.->|仅训练:潜变量承诺约束| Q

关键设计

1. 双潜变量分解:给幅值变化一条独立的表达路径

一个激波边界的位置与形状可以相近,但压力强度不同;若这些变化全部进入同一个多维码,量化会同时损伤边界结构和幅值精度。 Phaedra 把编码器输出按通道分成 8 维形态潜变量和 1 维幅值潜变量,分别服务于结构表达与密集数值刻度。 “形态像基函数、幅值像系数”只是理解模型的类比,并不存在预先选定的物理基底,也没有额外的解耦标签。

这一点区分了 Phaedra 与 VQ-VAE-2:后者把信息按空间尺度拆成粗细两层,本文则在同一潜网格上按语义角色拆通道。 它也不同于先重建低频部分、再量化剩余残差的串行方案;并行编码允许形态分支直接利用原场信息,而不必完全依赖第一阶段留下什么残差。 两条支路的职责主要由不同量化配置诱导,而不是由损失显式保证,所以应理解为软解耦。

2. 分支式 FSQ:用多维组合保存形态,用密集标量刻度保存幅值

形态分支采用有限标量量化(Finite Scalar Quantization,FSQ),8 个通道的级数为 \([5,4,4,3,3,3,2,2]\),笛卡尔积构成 8640 个可能的形态码。 每维级数不多,但组合后可以区分多种局部几何模式;这一配置与单流 FSQ 基线完全相同。 幅值分支则使用 1024 级的一维 FSQ,使强度可以沿一条有序数轴细密变化,而不需要学习 1024 个互不关联的高维模式。

两种词表并非一个 9664 类的联合分类器:每个位置保存两个索引,分别在 8640 类和 1024 类词表中选择。 这种结构把“强度要连续细分”和“形态要组合丰富”分别交给合适的量化几何。 等基数消融用 \([4,4,4,4,4]\) 的五维 FSQ 替代一维幅值词表,仍有 1024 个码,但重建变差,说明码的排列方式不等于码的数量。

FSQ 会先经有界非线性再取整,直接使用可能让大部分潜变量落入 tanh 饱和区。 本文在量化前乘一个缩放因子、量化后再除回同一因子,以避免前后潜变量的数值域错配。 形态与幅值分支的缩放因子分别为 10 和 0.1;幅值支路的较小因子使主体分布更多处于 tanh 的近线性区域。 这仍然是有限范围的表示,而非真正无界的标量编码;极端值受到软限制,模型并没有获得无限动态范围。

3. 学习式重组:让解码器学习形态与幅值的非线性交互

取回两组量化向量后,模型把它们拼接,经可学习的通道混合卷积送入解码器,而不是固定执行“形态乘幅值”。 固定乘法有清晰解释,却未必适合复杂的潜空间:不同局部结构可能需要不同的幅值调制,卷积重组可以从重建目标中学习这种交互。 编码器与解码器均为残差卷积骨干;潜网格不超过 \(32\times32\) 时加入一个注意力层,解码端通过最近邻上采样恢复分辨率。

这种灵活性带来的代价是幅值与形态不能被解释为完全独立的物理量。 交换幅值 token 后,重建能量与目标幅值来源的相关系数达到 0.99;只交换形态 token 时,能量仍与原场保持 0.97 的相关系数。 这些干预支持宏观能量主要由幅值支路控制,但不是严格分解定理。 附录还显示把幅值嵌入置零会造成约 300% 的高频能量保留率,即非物理伪影,而不是“更强的细节恢复”。

一个完整示例

以 \(128\times128\) 的密度快照为例,编码后得到 \(32\times32\) 网格,每个位置具有 8 维形态表示和一个幅值标量。 形态向量被映射到 8640 个组合码之一,幅值标量被映射到 1024 个刻度之一;最终保存的是两张 \(32\times32\) 索引图,而不是 9 通道浮点特征。 重建时两张索引图分别还原成量化嵌入,经学习式重组与解码器恢复密度场,幅值负责大尺度强度,形态补充局部边界与涡结构。

若只计算紧凑打包后的理论负载,每个潜位置的幅值码占 10 bit、形态码占 14 bit,总共 24 bit。 相对原始 \(128\times128\) 的 32-bit 单通道数组,理论节省约 95.3%;真实数据集从 65.0 GB 降至 3.44 GB,对应 94.7% 的实测节省。 这衡量的是数据存储,不代表模型参数缩小同样比例,也不代表 Transformer 的输入序列比 FSQ 更短。

损失函数 / 训练策略

分词器不使用感知损失或显式 PDE 残差损失,而是以逐点 \(L_1\) 重建误差为主,并约束编码器潜变量靠近其量化版本。

\[ \mathcal{L}=\|x-\hat{x}\|_1+0.25\|z_\mu-\operatorname{sg}[\hat z_\mu]\|_2^2+\|z_\alpha-\operatorname{sg}[\hat z_\alpha]\|_2^2. \]

其中 \(\operatorname{sg}\) 表示停止梯度;幅值量化较密,其承诺项本来较小,因此不另乘 0.25。 承诺损失不是用于告诉某支路“这是形态/这是能量”,而是防止编码输出远离量化词表的有效范围。 标准模型约 97M 参数,使用 AdEMAMix,基础学习率 \(10^{-4}\)、权重衰减 0.01、EMA 衰减 0.999,训练 1 个 epoch。 预训练数据来自 Compressible Euler 与 Incompressible Navier-Stokes 的六类数据集,含约 4.8M 训练样本;主重建比较使用共享骨干与训练方案。

下游实验冻结分词器,用约 38M 参数的序列到序列 Transformer 学习未来物理状态。 算子学习的两个输出头分别预测形态和幅值类别分布,使用交叉熵;all2all 训练把 7 个时间点组织为 28 个输入–输出组合,并加入时间差嵌入与二维 RoPE。 掩码自编码则使用单解码器处理两类 token,随机遮掉 75% 的空间位置,并在各物理变量上使用一致的空间掩码。 下游模型预测的 token 通过冻结解码器恢复为物理场;这与只评估分词器的重建误差是两个不同问题。

实验关键数据

主实验

以下摘录原文 Table 2 与 Table 7,按变量平均,保留原表报告尺度,不把 nMAE 误写为相对 \(L_1\) 误差。 ID 为分布内;OD1 改变初始条件或边界几何;OD2 包括 Poisson、Darcy、Allen-Cahn 与 Acoustic Wave 等未见 PDE。

模型 分布 nMAE ↓ nRMSE ↓ 局部方差误差 ↓ 最小谱相干性 ↑
Continuous AE ID 0.672 1.122 2.98 98.4%
FSQ ID 2.603 4.292 11.29 85.3%
Phaedra ID 1.522 2.489 5.96 93.6%
Continuous AE OD1 1.154 2.079 3.40 99.2%
FSQ OD1 1.878 4.331 20.55 93.9%
Phaedra OD1 1.224 2.435 6.47 98.1%
Continuous AE OD2 1.967 2.821 3.02 97.0%
FSQ OD2 4.865 6.675 11.66 64.3%
Phaedra OD2 3.147 4.237 5.83 77.6%

nMAE 是平均绝对残差除以相应变量在训练数据上的全局标准差;nRMSE 则把均方根残差除以同一标准差。 局部方差误差先在 \(7\times7\) 滑动窗口计算方差图,再比较两张图的最大差异:

\[ \Delta\sigma^2_{\mathrm{loc}}=\frac{\max|\sigma^2(y)-\sigma^2(\hat y)|}{\max|\sigma^2(y)|}\times100\%. \]

最小谱相干性不是能量谱误差,也不是拓扑正确率;它考察各频带中真实场与重建场的最低相干程度。 令 \(G\) 为自谱或交叉谱,原文定义为:

\[ \gamma^2(k)=\frac{|G_{y\hat y}(k)|^2}{G_{yy}(k)G_{\hat y\hat y}(k)},\qquad\gamma_{\min}=\min_k\gamma(k)\times100\%. \]

相对 \(L_1\) 则是绝对残差之和除以真实值绝对值之和;下游表用百分数报告,不能与 nMAE 混用。

消融实验

原文 Table 12、15、16 的 CEU RC 密度变量,均为 \(4\times4\) 下采样配置;这是单变量结果,不是上表跨变量均值。

配置 nMAE ↓ 局部方差误差 ↓ 最小谱相干性 ↑ 说明
FSQ 8.4078 20.8387 65.41% 单形态流,1024 个 token
Phaedra 5.5397 9.5770 84.27% 并行形态–幅值,2048 个 token
Codebook Ablation 7.8021 19.2576 72.37% 幅值改为等基数五维 FSQ
Residual Ablation 7.2957 20.0013 74.11% 先幅值,后形态残差

FSQ 与 Phaedra 的空间网格相同,但 token 数并不相同;更有针对性的证据来自两个仍有 2048 个 token、幅值词表仍为 1024 类的消融。 它们表明,一维幅值量化与并行分解各有贡献,不过某些平滑数据变量上残差版本更好,不能把这个单变量结论推广成逐项全胜。

下游摘录 Table 3,算子学习在最终时间点比较、每个模型取三种步进策略中最好的一种;MAE 则跨所有时间点平均。

任务 / 模型 KH 相对 L1 ↓ RC 相对 L1 ↓ RKH 相对 L1 ↓
算子学习 / CNO 10.15% 29.92% 11.82%
算子学习 / Continuous Transformer 9.28% 119.0% 11.56%
算子学习 / FSQ 11.28% 38.05% 17.99%
算子学习 / Phaedra 9.50% 27.21% 10.23%
MAE / FSQ 5.84% 29.16% 9.45%
MAE / Phaedra 4.75% 22.17% 7.27%

关键发现

  • ID nMAE 相对 FSQ 降低约 41.5%,但 Continuous AE 仍有更低的重建误差;“缩小量化损失”不等于“超过连续表示”。
  • 连续模型从 ID 到 OD1 的相对退化更大,但 OD1 的绝对 nMAE 仍为 1.154,低于 Phaedra 的 1.224,不能据此宣称离散模型全面胜出。
  • 高分辨率 PDE 的 \(16\times16\) 下采样比较中,Phaedra 的 nMAE 为 2.33,Cosmos 为 14.53;后者使用逐样本归一化,且两者不是参数匹配或相同预训练数据的实验。
  • 幅值词表即使只有 32 级也能优于 FSQ,在约 512 级后收益趋缓;提高潜网格 token 数同样有效,但 \(32\times32\) 附近已出现收益递减。
  • 原文多处把重建和下游结果指向 Table 4,但实际数值分别在 Table 2、3;Table 4 标题是 token scaling。正文称 ID 谱相干性“接近 100%”,本文保留实际均值 93.6%。
  • Table 3 的连续 Transformer RC 为 119.0%,附录 Table 27 给出 119.05%,属于报告精度差异;本表忠实保留 Table 3,不自行统一精度。

亮点与洞察

  • 离散化的几何比单纯扩大词表更重要。1024 个有序标量刻度与 1024 个五维组合码具有相同基数,却不具有相同的幅值表达能力。
  • 可解释性来自干预而非只看嵌入图。交换 token、置零支路、对照 Fourier 低通/高通让“幅值负责宏观、形态负责细节”有了可检验的含义,同时暴露非线性重组的伪影。
  • 数据压缩与下游准确率需要分开验证。该工作既报告真实存储节省,也检查 token 是否支持动力学学习,避免仅凭好看的重建图论证科学建模能力。

局限与展望

  • 各物理变量独立编码,忽略密度、速度、压力之间的耦合;未训练联合通道版本,因此耦合信息损失尚未量化。
  • 没有守恒律、对称性或 PDE 残差的硬约束,低局部方差误差与高谱相干性不构成物理守恒保证。
  • 主评测是二维网格数据;附录只有一个三维 CEU RC 初步实验,不支持宣称覆盖完整三维 PDE 家族或不规则网格。
  • 下游只是概念验证,算子学习比较采用各模型最优步进策略,且 Phaedra 每样本约 0.0398 秒,比 FNO 的 0.0126 秒慢;序列建模并非无代价收益。
  • 跨领域表现并非处处领先:Table 5 的 ImageNet 局部方差误差仍偏大,Table 23 的 NDVI 及部分 RGB 指标也不是 Phaedra 最优;不能把领域泛化改写成万能图像压缩优势。
  • 未来可检查联合变量编码、守恒约束和几何感知骨干;不同网格、维度与符号方程共享词表仍是动机,而非本文已完成的系统。

相关工作与启发

  • vs FSQ:共享形态量化配置,增加幅值支路及学习式重组;这是最直接的结构基线,但 Phaedra token 数翻倍,应结合等预算消融解释效果。
  • vs VQ-VAE-2 / VAR:它们按尺度或空间残差分层,本文按幅值与形态分解;文中的 VAR 使用 FSQ 改造,不能把结果等同于原始 VAR 系统表现。
  • vs Cosmos:预训练图像分词器重视自然图像统计,Phaedra 专门在物理场上训练并优化逐点精度;压缩率匹配不意味着训练、容量与 token 预算全部相同。
  • vs 连续神经算子 / 连续潜空间:连续 AE 重建更准确,离散表示提供压缩与类别分布建模接口;RC/RKH 下游优势说明表示会影响学习稳定性,但并未证明所有连续生成方法都较差。

评分

  • 新颖性: 4/5 — 将形状–增益思想具体落实为科学场双流量化,结构简洁且有针对性。
  • 实验充分度: 4/5 — 覆盖 PDE 重建、跨域、干预与下游,但联合变量和大规模基础模型验证尚缺。
  • 写作质量: 3/5 — 方法清晰,若干表号引用及泛化表述与实际表格边界不完全一致。
  • 价值: 4/5 — 为物理场接入离散 Transformer 提供高保真表示与实测压缩证据。