跳转至

The Galaxy's Guide to the Tokenizer: A Benchmark for Scientific Foundation Models

作者: Sogol Sanjaripour, Michael J. Smith, Manuel Pérez-Carrasco, Juan Rafael Martínez-Galarza, Bahram Mobasher et al.
主题: 天体统计
相关性: 6/10
链接: https://arxiv.org/abs/2606.25610

一、子领域定位

  • 本文属于天文学的哪一支：本文属于 astroinformatics / 天文机器学习 这一交叉子领域，更具体地说是 天文基础模型 (astronomical foundation models) 的构建与评估。核心科学问题是：如何将深度学习（特别是 transformer 架构）有效应用于海量天文图像数据，以提取可用于科学推断（如星系物理参数估计）的表征。该领域目前处于快速发展期，大量工作集中在将 NLP 领域的预训练-微调范式迁移到天文数据上，但 tokenization（将图像转化为 transformer 可处理的序列）这一关键步骤的影响尚未被系统研究。

  • 本文在这个子领域里的位置：它针对的是“天文基础模型中 tokenization 策略的选择如何影响下游科学任务性能”这一具体问题。它不提出新的模型架构，而是系统比较四种现有 tokenization 方法（Affine, AIM, JetFormer, VQ-VAE）在统一骨干网络（AstroPT）下的表现，并揭示重建质量与表征质量之间的解耦关系。

二、关键术语扫盲

  1. DESI Legacy Survey (DESI-LS)：一个大型地面巡天项目，使用光学望远镜拍摄了北半球约 14000 平方度天空的 g、r、z 三个波段图像。它为 DESI 光谱巡天提供目标星系图像，数据均匀、覆盖广，是训练天文基础模型的常用数据集。
  2. Postage stamp：从巡天大图像中裁剪出的、以单个星系为中心的小图像块（本文中为 256×256 像素），就像邮票一样。这是天文学家分析单个天体的标准数据格式。
  3. g, r, z 波段：天文观测中使用的不同滤光片，分别对应光谱的绿、红、近红外部分。不同波段的图像组合可以反映星系的不同物理性质，例如 g-r 颜色与恒星形成率相关，r-z 颜色与恒星年龄/质量相关。
  4. Redshift (z)：由于宇宙膨胀，遥远天体的光谱会向红端移动。红移值 z 是衡量天体距离和宇宙时间的关键量。photometric redshift (photo-z) 是通过多波段测光数据估算的，精度较低但覆盖广；spectroscopic redshift (spec-z) 是通过光谱测量得到的，精度极高但成本高、样本小。
  5. Stellar mass (M*)：星系中所有恒星的总质量，是天体物理中最基本的物理量之一，与星系的形成和演化历史密切相关。
  6. Specific Star Formation Rate (sSFR)：恒星形成率除以恒星质量，衡量一个星系单位质量下形成新恒星的活跃程度。高 sSFR 意味着星系正在活跃地形成恒星，低 sSFR 意味着星系已经“熄灭”。
  7. Galaxy morphology：星系的视觉形态，如旋涡星系（有旋臂）、椭圆星系（光滑无结构）、不规则星系。形态与星系的物理演化阶段和环境有关。本文使用了来自 Galaxy Zoo 公民科学项目的形态指标（如 smoothness, disk fraction）。
  8. Tokenization：在 transformer 模型中，将输入数据（如图像）分割并映射成一系列离散或连续的“token”（标记）的过程。对于图像，通常将图像切成小块（patch），然后将每个 patch 映射成一个向量。不同的 tokenization 策略（如线性投影、离散码本）会极大影响模型学到的表征。
  9. VQ-VAE (Vector Quantised-Variational AutoEncoder)：一种生成模型，其核心是学习一个离散的“码本”（codebook）。它将图像 patch 编码成连续的潜在向量，然后将其量化为码本中最接近的离散码字。这种离散瓶颈迫使模型丢弃细节噪声，学习更语义化的表征。
  10. Normalising Flow：一种生成模型，通过一系列可逆变换将简单的概率分布（如高斯分布）映射到复杂的数据分布。其可逆性意味着编码和解码是同一个过程，理论上可以无损地保留输入信息。JetFormer 使用 flow 来生成连续的 token。
  11. Linear probe / MLP probe：一种评估表征质量的方法。在预训练好的模型上，冻结其参数，只训练一个简单的线性分类器/回归器（linear probe）或一个浅层 MLP（MLP probe），来预测某个目标属性（如红移）。性能越高，说明该属性在表征中越容易被“读出”，即表征组织得越好。
  12. SSIM / PSNR：图像重建质量的评价指标。SSIM（结构相似性）衡量两幅图像在亮度、对比度和结构上的相似度；PSNR（峰值信噪比）衡量像素级误差。两者越高，重建质量越好。

三、天文学家关心的问题

天文学家正面临一个“数据洪流”问题。新一代巡天项目（如 Euclid, LSST, DESI）将产生数十亿个天体的多波段图像和光谱数据。传统的人工或基于简单统计的分析方法已无法应对。因此，天文学界迫切需要能够自动、高效地从海量数据中提取物理信息的机器学习方法。Transformer 基础模型（如 AstroPT）被视为一个有前景的范式：先在大量无标签图像上预训练，再针对特定科学任务（如预测红移、恒星质量、形态分类）进行微调。

然而，一个关键但被忽视的问题是：如何将图像转化为 transformer 可以处理的 token 序列？ 不同的 tokenization 策略（线性投影、MLP、离散码本、可逆流）会如何影响模型学到的表征？这些表征是否保留了天文学家关心的物理信息？重建质量好的模型是否一定意味着其表征对科学推断更有用？本文正是为了回答这些问题。

当前领域的主流方法是直接采用 NLP 领域的 tokenization 策略（如将图像切块后线性投影），或使用 VQ-VAE 等预训练好的图像分词器。这些方法的局限在于：它们通常针对自然图像设计，其评估标准（如重建质量）可能与天文科学目标（如准确预测物理参数）不一致。本文通过系统比较，揭示了这种不一致性，并指出没有一种 tokenization 策略在所有任务上都是最优的，选择应取决于最终的科学目标。

四、数据问题

  • 数据来源：DESI Legacy Survey (DESI-LS) Data Release 8。这是一个地面光学巡天项目。
  • 数据形态：imaging。每个星系是一个 256×256 像素的“邮票”图像，包含 g、r、z 三个波段（即三通道图像）。数据量：训练集 64 万张，测试/验证集 16.7 万张。
  • 几何结构：图像是规则的欧几里得网格，但天体在图像中的位置、大小、朝向是随机的。从物理角度看，星系图像位于一个高维流形上，其内在维度远低于像素数。
  • noise model & 测量误差：天文图像噪声通常由泊松噪声（来自光子计数统计）和读出噪声（来自探测器电子设备）组成，且通常与信号相关（heteroskedastic）。本文未显式建模噪声，但预处理中的 z-score 归一化（对每个 patch 减去均值除以标准差）可以部分缓解光照不均和噪声水平差异。
  • selection effect / survey mask / Malmquist bias：DESI-LS 的观测策略和源提取算法会引入选择效应。例如，更亮、更大的星系更容易被探测到，而暗弱、弥散的星系可能被遗漏（Malmquist bias）。本文使用的星系样本来自一个公开数据集，其选择函数是已知的，但本文未深入讨论其对结果的影响。
  • 缺失 / censoring / truncation / 计算约束：数据是完整的，没有缺失值问题。主要计算约束来自训练 transformer 模型（89M 参数）和 VQ-VAE 分词器所需的 GPU 资源。
  • 哪些数据特性是“漂亮的统计学问题”，哪些是“纯工程难题”：
    • 漂亮的统计学问题：星系图像的高维流形结构、不同物理参数（如红移、恒星质量）与图像特征之间的复杂非线性关系、重建质量与表征质量的解耦现象——这些都是非参数统计、高维推断和表征学习可以切入的问题。
    • 纯工程难题：训练大规模 transformer 的工程优化（如分布式训练、混合精度）、数据加载和预处理管道的效率、不同 tokenizer 训练过程的差异（本文承认其比较并非完全消融实验）——这些是工程实现问题，而非统计挑战。

五、模型问题

  • 文章建立的模型/方法：本文建立了一个统一的评估框架，而不是提出新模型。核心是：固定一个 AstroPT transformer 骨干网络，然后替换其输入端的 tokenization 模块。四种 tokenization 策略是：
    1. Affine：一个简单的线性投影层，将每个 8×8 像素的 patch 映射到 768 维的 embedding 空间。这是最简基线。
    2. AIM：用一个 MLP 替换线性投影，提供非线性映射能力。
    3. JetFormer：使用一个可逆的 normalising flow 将图像 patch 映射为连续的 token。由于 flow 是可逆的，编码器和解码器是同一个网络。
    4. VQ-VAE：使用一个预训练好的 VQ-VAE 模型，将图像 patch 编码为离散的码本索引（codebook index）。码本大小为 512。
  • 模型的关键假设：
    • 来自物理学约束：假设星系图像的物理信息（如红移、恒星质量）可以被 transformer 模型从像素中学习到。这是整个天文基础模型领域的基本假设。
    • 为了计算可行性：固定骨干网络大小（89M 参数）和训练超参数，以便公平比较 tokenizer。将图像切分为 8×8 的 patch 是计算上的折中。VQ-VAE 需要预训练，这引入了额外的计算成本。
  • 推断手段：本文的推断是评估而非参数估计。评估手段是 linear probe 和 MLP probe，即训练一个简单的回归模型来预测物理参数。这本质上是一个监督学习任务，使用 R² 作为性能指标。不确定性通过 k=10 折交叉验证 来量化。
  • 核心数值结论 + uncertainty 量化方式：
    • VQ-VAE 在物理参数预测上表现最好（线性 probe R² 最高可达 0.85 用于 spec-z），且其线性 probe 性能接近甚至超过 MLP probe，说明其表征组织得更好。
    • JetFormer 在图像重建上表现最好（PSNR 31.11 dB, SSIM 0.762），但在物理参数预测上不如 VQ-VAE。
    • Affine 和 AIM 性能几乎相同，说明 MLP tokenizer 头在骨干网络足够强时没有带来额外好处。
    • 不确定性通过交叉验证的标准差给出（如 SSIM 为 0.762 ± 0.138），表明重建质量在不同星系间有较大波动。

六、对统计学家的判断

  1. 这篇文章作为入门读物质量如何？

    • 4/5 星。理由：文章对数据来源、预处理、模型架构和评估协议交代得非常清晰，术语解释（如 tokenization, probe）对非天文背景的读者友好。它很好地暴露了本子领域的核心思路（基础模型 + 下游任务评估）和关键挑战（tokenization 选择、重建 vs. 表征的权衡）。唯一的扣分点是，它假设读者对 transformer 和 VQ-VAE 等深度学习概念有一定了解，但这对统计学家来说通常不是障碍。

  2. 这个问题值不值得统计学家进入工作？

    • 值得。论证如下：
      • (i) 科学重要性：天文学界非常在乎这个问题。随着 Euclid 和 LSST 等巡天项目启动，如何高效、准确地从海量图像中提取物理信息是决定下一代科学发现的关键。基础模型是当前最受关注的解决方案，而 tokenization 是其核心组件。本文的发现（重建与表征解耦）直接挑战了“重建好 = 表征好”的常见假设，对模型设计和评估有重要指导意义。
      • (ii) 方法学空间：数据特性提出了真正的统计挑战。例如：
        • 表征质量的量化：本文使用的 linear/MLP probe 是一种简单方法，但可能低估了表征中的信息量。统计学家可以设计更严谨的信息论度量（如互信息估计）或非参数假设检验来比较不同表征。
        • 重建与表征解耦的理论理解：为什么 VQ-VAE 的离散瓶颈能产生更“线性可读”的表征？这背后是否有信息论或统计学习理论上的解释？这是一个有深度的理论问题。
        • 选择效应与因果推断：天文数据存在严重的选择效应。如果 tokenizer 学到的表征与选择函数相关，那么下游预测可能会产生偏差。统计学家可以引入因果推断的视角来分析和纠正这种偏差。
      • (iii) 社区开放性：作者群中没有统计学家（主要是天文学家和计算机科学家）。方法学讨论（如 probe 的局限性、比较的非消融性）是坦诚的，但深度有限。该领域非常欢迎方法学贡献，因为天文学家普遍意识到他们需要更好的统计工具。本文的 GitHub 仓库是公开的，社区协作氛围好。
      • (iv) 武器库匹配度：
        • 够用。研究者的 very_familiar 武器库（非参数统计、高维渐近、逆问题、软件开发）足以理解本文的评估框架，并可以立即用于提出改进方案。例如，nonparametric statistics 可用于设计更鲁棒的表征质量度量；high-dimensional asymptotics 可用于分析 embedding 空间的几何性质；software development 可用于实现新的评估工具。
        • 缺一块：研究者对深度学习（特别是 transformer 和 VQ-VAE）的熟悉程度是 moderately_familiar 或更低。要在这个方向做深入的 follow-up 工作（如设计新的 tokenizer 或改进训练过程），需要补上深度学习理论（如自注意力机制的统计性质、离散表征的优化）这一块。但这不构成障碍，因为统计学家可以专注于“评估”和“理解”问题，而非“设计”新模型。
      • 明确结论：值得。理由：这是一个科学重要性高、方法学空间大、社区开放的方向。研究者的现有武器库足以在“表征评估”和“理论理解”这两个子问题上立即做出贡献，而无需成为深度学习专家。

  3. 若值得进入，研究者能做的具体问题（最多 2 条）

    • 问题 1：设计一个基于非参数统计的表征质量度量，以替代 linear/MLP probe。 武器库：nonparametric statistics, minimax bounds for estimation problems。第一步动作：将每个星系的 embedding 向量视为一个随机变量，将物理参数（如红移）视为另一个随机变量。使用距离相关性 (distance correlation) 或 HSIC (Hilbert-Schmidt Independence Criterion) 来度量 embedding 与物理参数之间的依赖关系。这比 probe 更通用，因为它不假设任何函数形式，且可以捕捉非线性依赖。可以比较不同 tokenizer 的 HSIC 值，看是否与 probe 结果一致。
    • 问题 2：从高维渐近的角度分析 VQ-VAE 的离散瓶颈为何能产生更“线性可读”的表征。 武器库：high-dimensional asymptotics, inverse problems with random noise。第一步动作：将 VQ-VAE 的编码-量化-解码过程建模为一个高维逆问题。假设原始图像 patch 位于一个低维流形上，VQ-VAE 的码本相当于对这个流形进行离散化。可以证明，在一定的信噪比和码本大小条件下，这种离散化过程等价于一个去噪步骤，它去除了与流形正交的噪声维度，从而使得下游的线性回归问题（即 linear probe）的 minimax 风险降低。可以用随机矩阵理论来分析码本向量的谱分布。

  4. 下一步读什么？

    • 入门综述：
      • Smith & Geach (2023) "Astronomia ex machina: a history, primer and outlook on neural networks in astronomy" (被引文献 [13])。这是一篇关于天文深度学习的综述，涵盖了从早期 MLP 到当前基础模型的发展，是进入该领域的绝佳起点。
    • 方法学奠基论文：
      • Smith et al. (2024) "AstroPT: Scaling Large Observation Models for Astronomy" (被引文献 [21])。本文使用的骨干网络，了解其架构、训练细节和 scaling law 是理解本文工作的前提。
      • van den Oord et al. (2017) "Neural Discrete Representation Learning" (被引文献 [5])。VQ-VAE 的原始论文，是理解离散表征学习的基础。
    • 可动手的公开数据集：
      • The Multimodal Universe (2024) (被引文献 [18])。一个包含 100TB 多模态天文数据的大型公开数据集，包括图像、光谱、时序数据等。它附带了一系列基准任务，是测试新方法、进行方法学研究的理想平台。数据集和代码在 GitHub 上公开。

七、术语小抄

英文术语 中文 一句话解释
Foundation Model 基础模型 在大规模数据上预训练的通用模型，可通过微调适应多种下游任务。
Tokenization 分词/令牌化 将输入数据（如图像）转换为模型可处理的离散或连续序列的过程。
Autoregressive Model 自回归模型 一种按顺序生成数据的模型，每一步预测下一个 token，基于之前的所有 token。
Decoder-only Transformer 仅解码器 Transformer 一种 transformer 架构，只有解码器部分，常用于自回归生成任务。
Embedding 嵌入/表征 模型内部对输入数据的向量表示，通常位于一个低维连续空间中。
Linear Probe 线性探测 一种评估表征质量的方法，通过训练一个线性模型来预测目标属性。
VQ-VAE 向量量化变分自编码器 一种生成模型，通过离散码本将连续数据压缩为离散的 token。
Normalising Flow 归一化流 一种生成模型，通过一系列可逆变换将简单分布映射到复杂分布。
Codebook 码本 VQ-VAE 中一个可学习的离散向量集合，用于将连续向量量化为离散索引。
Postage Stamp 邮票图像 从巡天图像中裁剪出的、以单个天体为中心的小图像块。
Redshift (z) 红移 由于宇宙膨胀导致天体光谱向红端移动的量，用于衡量距离和宇宙时间。
Stellar Mass (M*) 恒星质量 星系中所有恒星的总质量。
Specific Star Formation Rate (sSFR) 比恒星形成率 单位恒星质量的恒星形成速率，衡量星系形成新恒星的活跃程度。
Malmquist Bias 马尔姆奎斯特偏差 由于观测极限导致更亮的天体更容易被探测到，从而引入的样本选择偏差。
SSIM / PSNR 结构相似性/峰值信噪比 衡量图像重建质量的指标，SSIM 更关注结构，PSNR 更关注像素级误差。

Maintained by 陈星宇 · Homepage · Source on GitHub

评论