流形假设下的架构再思考： JiT 扩散模型的改进与机理探究

摘要：Li 和 He 提出的 JiT (Just image Transformers) 架构基于流形假设，通过直接预测干净图像 (x-prediction)，验证了简单的线性层配合 ViT 即可有效处理高维像素数据。然而， JiT 的极简线性 Patch Embedding 可能不足以充分捕捉自然图像高度卷曲的非线性流形结构。本文首先在 Embedding 层引入 SiLU 激活函数，构建非线性瓶颈以增强对低维流形嵌入的拟合能力。 进一步地，本文深入探讨了骨干网络 (Backbone) 中流形约束 (降维) 与计算容量 (升维) 的本质矛盾。通过将 Transformer Block 内部替换为瓶颈结构的对比实验，本文揭示了一个关键的精度-多样性权衡 (Precision-Recall Trade-off)：显式的降维压缩虽然能有效过滤非流形噪声，从而显著提升生成图像的保 真度 (Precision) 与 FID 指标；但这种严苛的流形约束同时也限制了模型对高熵随机偏差的建模能力，导致生成样本的多样性 (Recall) 下降。 此外，针对 JiT 缺乏语义约束的问题，本文引入了时间 (time) 与旋转 (rotation) 预测的自监督辅助损失。在 ImageNet 256 × 256 数据集上的实验表明，非线性 Embedding 与自监督信号有效提升了 FID 指标，而 Block 层的瓶颈化实验则从反面论证了“计算容量”在扩散模型骨干网络中的必要性。

关键词： 计算机视觉；扩散模型； JiT；非线性流形；瓶颈结构；高维数据拟合；自监督学习