非官方 DiffMath 复现结果：我按照你们论文里面的思路写了一个代码，扩散/CTC 指标有所改善，但字形可读性依然较差。整体效果似乎不理想。

[hxypqr]

您好，感谢您发布 DiffInk 以及非常有意思的 DiffMath 研究工作。

我之前尝试基于 DiffMath 进行非官方的 PyTorch 复现，并以已发布的 DiffInk 实现及预训练权重作为起点。

在项目启动时，我未能找到 DiffMath 的官方实现或模型检查点（checkpoint），因此系统中的数学表达式处理部分是根据论文及补充材料重新构建的。因此，很有可能我在某个关键实现细节上存在误解。

当前流程

针对数学表达式，处理流程如下：

LaTeX
-> MathML
-> RelAST
[符号标识、空间关系、嵌套深度]
-> MathDiT
-> MathVAE 解码器
-> 在线笔迹轨迹

目前的 RelAST 实现支持以下关系：

• RIGHT（右侧）
• SUP / SUB（上标/下标）
• ABOVE / BELOW（上方/下方）
• INSIDE（内部）
• 分数结构
• 平方根
• 求和与积分
• 矩阵行（ROW）/ 单元格（CELL）结构

非神经网络的排版器（compositor）仅作为诊断或备用系统存在。下文讨论的结果均来自纯神经网络生成，在采样或候选结果选择过程中，未将排版器生成的轨迹用作参考。

数据

初始数学训练集包含：

• 来自 MathWriting 的 253,064 个真实在线数学表达式；
• 6,178 条真实的孤立符号轨迹；
• 60,000 个由真实孤立符号轨迹组合而成的平衡合成表达式；
• 合并索引中的 313,064 个条目。

合成数据集涵盖了：

• 线性表达式；
• 上标和下标；
• 分数；
• 根式；
• 嵌套表达式；
• 矩阵。

我还准备了约 6,000 条采用 Graves 格式（text + [pen_up, dx, dy]）的对齐英语行轨迹。这些数据未包含在初始数学模型检查点中。后续的 JOINT/BALANCED 实验旨在改进英语及英数混合内容的生成效果。

评估时采用以下方式：

• 包含 384 个真实 MathWriting 表达式的固定集合（使用固定扩散噪声）；
• 包含 100 个英语/LaTeX 输入的固定基准测试集（涵盖十种结构类别）；
• 对 OLD、JOINT10K 和 BALANCED 检查点生成的样本进行直观的视觉对比。

架构

为了获得一个可运行且易于调试的基线模型，我保留了与 DiffInk 兼容的维度设置，而不是直接训练论文中原始的架构：

MathVAE

• 一维卷积编码器/解码器；
• 潜在空间维度：384；
• 时间维度压缩率：8倍；
• 参数量：约 1300 万；
• GMM 坐标解码器；
• 笔触状态预测；
• 符号/关系 CTC 预测头。

MathDiT

• 16 层 Transformer；
• 隐藏层维度：896；
• 注意力头数：14；
• 参数量：约 2.42 亿；
• 余弦扩散调度（cosine diffusion schedule）；
• x_0 预测；
• DDIM 采样；
• 无分类器引导（classifier-free guidance）；
• 基于 RelAST 的结构条件输入。

选择这些维度主要是为了能够严格加载并映射已发布的 DiffInk 权重。我并不认为 384 维、8 倍压缩的配置对于数学手写生成而言是最优的。该论文所采用的配置似乎使用了较短的潜在通道维度和 4 倍的时间压缩。我推测，这种差异对于小字形、下标及矩阵元素可能至关重要。

MathVAE 目标函数

当前的 VAE 目标函数如下：

[
L_{\mathrm{VAE}}

0.1 L_{\mathrm{GMM}}
+1.0 L_{\mathrm{pen}}
+1.0 L_{\mathrm{CTC}}
+10^{-7} L_{\mathrm{KL}}.
]

其中：

• L_GMM 是坐标负对数似然（coordinate negative log-likelihood）；
• L_pen 是一种 Focal 风格的笔画状态损失（pen-state loss）；
• L_CTC 用于监督符号/关系识别；
• L_KL 用于对潜在分布进行正则化。

选择这些权重主要是为了与 DiffInk 的配置及 DiffMath 相关资料中报告的数值保持兼容；目前尚未通过全面的超参数搜索来确定其最优性。

一个难点在于 GMM 负对数似然的值可能会变为负数。因此，标量形式的总损失无法作为衡量视觉重建质量的可靠指标。

DiT 模型的训练采用掩码 MSE（masked MSE）策略，针对扩散退化过程（diffusion corruption process）下的纯净潜在目标（clean latent target）进行优化。

此外，模型还包含一个长度预测头，针对 log(1 + length) 使用较小的 SmoothL1 损失进行训练。由于其表现并未优于基于统计的长度先验（statistical length prior），该预测结果目前仅用于诊断分析，在推理阶段不予采信。

已修复的重要实现缺陷

在早期版本中，系统会将每个批次（batch）中的样本填充（pad）至最长公式的长度，并沿该填充长度对结构条件进行插值处理。

这导致同一表达式的结构表示会随着批次中包含的其他表达式的不同而发生变化。

该问题已得到修复：在进行批处理之前，每个样本仅根据其自身的有效 RelAST token 进行独立对齐。

修复该问题后，模型对正确的 RelAST 条件变得更加敏感。

训练

经修复后的训练序列如下：

1. 10,000 步平衡的合成结构训练；
2. 26,000 步使用真实 MathWriting 轨迹数据的重放（replay）训练。

这相当于大约 755 万次样本曝光。

后续实验包括：

• JOINT10K：一个 10,000 步的联合训练阶段；
• BALANCED：一个简短的平衡几何重放阶段，旨在恢复字形和轨迹质量，同时不损失 JOINT10K 阶段在内容方面取得的成效。

在推理阶段，我目前采用以下设置：

• 20 步 DDIM；
• 每个表达式生成 4 个候选结果；
• 模型常驻显存，并采用批处理方式生成候选结果。

定量改进

与修复条件对齐（conditioning-alignment）问题之前的模型相比，表现最好的修复后检查点（checkpoint）实现了以下改进：

• t = 999 时 MSE 降低约 11.0%；
• t = 900 时 MSE 降低约 23.1%；
• t = 500 时 MSE 降低约 23.5%；
• 正确 RelAST 条件与故意错误的 RelAST 条件之间的区分度（separation）提升了 78.5%。

对于最终的 BALANCED 实验，在高噪声评估设置下的指标变化如下：

• 总体：1.1552 -> 1.0661；
• 英语子集：1.0799 -> 0.9721；
• 英语/LaTeX 混合子集：1.2317 -> 1.0256。

与纯联合训练模型相比，BALANCED 模型恢复了约 5.9% 的几何得分，而内部 CTC 得分仅下降了约 0.7%。

全部 100 个基准输入均能生成有效的 PNG、InkML 和 JSON 输出。

在模型常驻 GPU 的情况下，生成 4 个候选结果（使用 20 步 DDIM）平均每个输入耗时约 0.233 秒。

遗留的视觉问题

尽管数值指标有所提升，但输出结果往往仍难以辨认。

典型的失败案例包括：

• 尽管整体单词长度合理，但局部字符形状畸变；
• 数字和变量发生扭曲；
• 根号符号包含正确的大尺度外框，但内部符号错误；
• 微小的上标和下标丢失；
• 同一表达式内笔画尺度不一致；
• 公式整体结构合理，但局部字形无法辨认；
• CTC 分数提升，但人类可读性并未相应改善。

我附上了一张对比图，包含以下几列：

1. OLD 检查点（checkpoint）；
2. JOINT10K 检查点；
3. BALANCED 检查点。

即使对于范数、根号、距离和短英文单词等相对简单的例子，局部字形质量仍远差于论文中展示的效果。

问题

如果能就以下几点得到您的建议，我将不胜感激。

1. 4倍 VAE 压缩是必要的吗？

继承自 DiffInk 的 8倍时间轴压缩（temporal compression）是否是导致小字符、索引和局部笔画细节丢失的主要原因？

在您的实验中，针对数学手写内容，是需要从头训练一个独立的 4倍 MathVAE，还是可以直接迁移 DiffInk 的 VAE？

2. RelAST token 是如何与潜在轨迹（latent trajectory）对齐的？

这是我最不确定的部分。

关于符号（symbol）、关系（relation）和深度（depth）的嵌入（embeddings）：

• 是在最终潜在序列上均匀重复；
• 分配给符号级别的轨迹区间；
• 根据笔画数或点数进行对齐；
• 在没有显式时间对齐的情况下进行交叉注意力（cross-attention）计算；
• 还是通过其他规则进行扩展？

此处微小的差异很可能解释了为什么模型能学会整体布局，却无法正确还原局部符号。

3. 使用了什么样的轨迹预处理？

能否明确说明以下方面的具体选择：

• 绝对坐标与相对坐标；
• 坐标归一化；
• 长宽比归一化；
• 点重采样；
• 重复点去除；
• 抬笔（pen-up）动作的表示；
• 最大轨迹长度；
• 填充（padding）与掩码（masking）；
• 极短笔画的处理方式？我担心预处理环节的不匹配可能会导致 VAE 将过多的容量分配给采样密度细节，而非字形本身的特征。

4. 在训练 MathDiT 之前，MathVAE 的重构精度应达到何种水平？

在训练 MathDiT 之前，您是否已将 MathVAE 训练至其独立重构结果与输入轨迹几乎无法区分的程度？

在开始扩散模型（diffusion）训练之前，是否存在您认为必须达到的重构或识别阈值？

我目前的推测是，DiT 学习建模的那个潜在空间，可能已经丢失了过多的原子字形信息。

5. 使用了什么样的扩散目标（diffusion target）和时间步加权（timestep weighting）？

发布的模型（或论文中的模型）是基于以下哪种方式训练的：

• x_0 预测；
• 噪声预测；
• 速度预测；
• SNR 或 min-SNR 加权；
• 均匀时间步采样；
• 还是侧重于高噪声时间步的偏置分布？

目前的“高噪声 MSE”指标持续改善，但这种改善并未可靠地转化为可读的轨迹。

6. VAE 的各项损失（losses）是如何平衡的？

在 GMM、笔画（pen）、CTC 和 KL 损失之间，你是使用固定的权重，还是采用了某种调度策略，例如：

• 延迟 CTC（delayed CTC）；
• KL 预热（warm-up）；
• 梯度平衡；
• 将 VAE 重建（reconstruction）与识别（recognition）阶段分开？

强大的 CTC 头部（head）是否可能促使模型学习到一种潜在表示——这种表示虽有利于模型自身的识别器，却不足以支持生成人类可读的解码结果？

7. 在此场景下，CTC 是可靠的模型选择指标吗？

我的 CTC 评估器与生成系统共享组件和表示，因此它可能会给那些“自洽但视觉上畸形”的输出打高分。

你使用了哪种独立指标来选择检查点（checkpoint）：外部手写识别器、ExpRate、符号准确率、结构准确率、人工评估，还是其他衡量标准？

8. 训练规模和课程学习（curriculum）有多重要？

目前的模型在主要的训练阶段仅处理了约 755 万个样本，这远低于我根据已公布的全局批次大小（global batch size）和训练步数所估算的论文级训练规模。

根据您的经验，在这一规模下，目前的视觉质量是预料之中的吗？还是说这暗示了某种更根本的实现错误？

您是否采用了某种课程学习策略，例如：

孤立字形
-> 简短的线性表达式
-> 上标/下标
-> 分数与根式
-> 长嵌套表达式
-> 矩阵

而不是从一开始就对所有结构进行采样？

9. 书写者/风格条件（writer/style conditioning）有多重要？

目前的复现版本尚未包含可靠的书写者/风格编码器。

即便结构条件正确，缺少风格条件是否会导致模型将不兼容的笔画形状进行平均，从而产生局部歧义的字形？

我目前的假设

针对可能的原因，我目前的排序如下：

1. 8x MathVAE 导致的信息丢失过多；
2. RelAST 与潜在空间（latent space）的对齐不正确；
3. 独立 VAE 的重构质量不足；
4. 轨迹归一化或重采样（resampling）不匹配；
5. 训练规模远小于论文中的设置；
6. CTC 和高噪声 MSE 无法很好地反映视觉质量；
7. 合成器（synthetic-compositor）存在分布偏移；
8. 缺少书写者/风格条件（writer/style conditioning）。

如果需要，我可以提供相关的配置、预处理代码、模型定义、训练日志以及少量生成的轨迹样本。

非常感谢您抽出时间，也感谢您在 DiffInk/DiffMath 方面所做的工作。