基于K-d Tree与Keras的测光红移估计：解决训练样本偏差的机器学习实践

1. 项目概述与核心挑战

在当代天文学，尤其是宇宙学研究中，红移是衡量天体距离和宇宙膨胀历史的核心物理量。光谱红移精度高，但获取成本巨大，无法应对如暗能量巡天（DES）这类产生数亿乃至数十亿天体测光数据的现代大型巡天项目。因此，测光红移（Photometric Redshift, photo-z）技术——即仅通过几个宽波段测光数据来估算红移——成为了不可或缺的工具。然而，测光红移的精度和可靠性始终是制约后续科学分析的瓶颈，其误差主要来源于测光误差、模板不匹配以及训练样本与目标样本之间的分布差异。

我最近深入实践了一项研究，核心目标是利用机器学习，特别是深度学习技术，来提升对DES数据中星系测光红移的估计质量。这项工作的特殊之处在于，我们并非简单地将数据扔进一个黑箱模型，而是直面一个关键的现实问题：我们的训练样本来自VIPERS光谱巡天，它只覆盖了DES极小一部分天区且有其特定的颜色选择函数。这意味着VIPERS星系并不能完全代表整个DES星系群体。直接用一个有偏的训练集去预测全体，结果必然在某些区域不可靠。因此，我们的方法可以概括为“两步走”：第一步，利用K-d Tree和多种分类器构建一个“相似度过滤器”，量化每个DES星系与VIPERS训练集的相似程度（即λ参数）；第二步，基于Keras框架构建一个专为回归问题优化的神经网络（我们称之为KeraZ），但仅对那些与VIPERS高度相似的DES星系（λ值高）的预测结果给予高置信度。这套组合拳的目的，是在利用深度学习强大拟合能力的同时，通过数据本身告诉我们预测结果在何处可信，在何处存疑，从而为宇宙学分析提供一个更干净、更可靠的红移样本。

2. 核心思路与方案选型解析

2.1 问题拆解：为何不能直接训练？

面对DES DR2超过2.47亿个星系和VIPERS仅约5万个光谱红移样本的悬殊对比，最直接的冲动是用VIPERS训练一个模型去预测所有DES星系。但这会忽略一个根本性的“域适应”问题。VIPERS样本并非DES的随机子集，它通过一系列颜色截断（color cuts）刻意选择了中高红移（z > 0.5）的星系，并且包含了DES所没有的u波段信息。这导致在颜色空间（例如g-r, r-i, i-z, z-Y）中，DES星系点云的分布与VIPERS点云的分布并不完全重叠。有些DES星系可能落在VIPERS覆盖区域之外，模型对这些“陌生”区域的预测纯属外推，可靠性极低。

因此，我们的核心思路从单纯的“预测”转变为“筛选+预测”。首先，我们需要一个机制来识别哪些DES星系在测光属性上与VIPERS星系足够相似，从而其红移预测是可信的。这个相似度指标，我们定义为λ，范围在0到1之间。λ接近1表示该DES星系在颜色空间中有很多VIPERS“近邻”，预测可信；λ接近0则表示它是VIPERS样本中的“异类”，预测结果参考价值低。

2.2 技术选型：为什么是K-d Tree与Keras的组合？

为了实现上述思路，我们在两个关键环节做出了技术选择：

1. 相似度度量与高效计算：K-d Tree的优势计算λ的核心是衡量DES星系与所有VIPERS星系在颜色空间中的“距离”。最笨的方法是计算每个DES星系与所有VIPERS星系的距离，这在海量数据下是计算灾难。因此，我们需要一种能高效进行近邻搜索的数据结构。K-d Tree（K-dimensional Tree，K维树）正是为此而生。它是一种空间划分数据结构，特别适用于多维空间（我们的颜色空间是4维或5维）的最近邻搜索。其原理是递归地将空间沿坐标轴的中位数进行划分，形成一棵二叉树。构建一次K-d Tree后，对于任意一个查询点（DES星系），寻找其最近邻VIPERS星系的时间复杂度可以从O(N)降至平均O(log N)，这对于数亿量级的计算是至关重要的效率提升。

我们选择了闵可夫斯基距离（Minkowski distance）作为距离度量，并设定参数p=2，此时它退化为最常用的欧几里得距离。在对比了欧氏距离、曼哈顿距离等多种度量后，闵可夫斯基距离在此问题上取得了最好的分类效果。我们构建了两个K-d Tree：一个基于4种颜色（g-r, r-i, i-z, z-Y），另一个在此基础上增加了i波段星等（magnitude），构成5维空间。增加星等信息是为了引入亮度维度，可能帮助区分由于亮度差异导致的颜色相似但本质不同的星系。

2. 红移预测模型：为何从ANNz转向自定义Keras模型？在红移预测本身，已有成熟工具如ANNz和GPz。我们选择基于Keras自建模型（KeraZ），主要出于灵活性和透明度的考虑。ANNz虽然经典，但其内部架构和正则化选项相对固定。Keras作为一个高层深度学习API，允许我们精细地控制网络每一层。

我们设计了一个相对简单的全连接神经网络：一个输入层（对应颜色/星等特征），一个包含30个节点的隐藏层，以及一个输出层（预测红移值）。隐藏层使用tanh激活函数，之后接了一个批归一化（Batch Normalization）层。这里有两个关键设计点：

• 激活函数选择tanh：相比更流行的ReLU，tanh的输出范围是(-1, 1)，是以零为中心的，这在回归问题中有时能使梯度下降更平稳。我们通过实验对比发现，在此问题上tanh略优于ReLU。
• 批归一化与正则化：批归一化层通过对每一批数据进行归一化处理，加速训练并缓解内部协变量偏移。更重要的是，我们为隐藏层的核（kernel）添加了L2正则化。天文数据量虽大，但特征维度低，模型很容易过拟合训练集中的噪声。L2正则化通过对权重施加惩罚，迫使网络学习更平滑、更通用的函数，这是提升模型泛化能力、防止在验证集上表现骤降的关键。

优化器我们选择了Adam，它自适应地调整每个参数的学习率，对于这种规模的数据和模型非常高效且稳定。损失函数直接采用均方误差（MSE），因为红移预测是典型的回归问题。

注意：处理异方差噪声天文测光数据的一个关键特性是噪声并非恒定。亮星的测光误差小，暗星的测光误差大，这就是异方差噪声。GPz模型原生支持异方差噪声建模，这是它的一个优势。在我们的KeraZ实现中，虽然没有像GPz那样在损失函数中显式建模噪声方差，但我们通过将测光误差作为额外的输入特征提供给网络，让网络在训练过程中自行学习不同噪声水平下的映射关系，这是一种隐式处理异方差性的方式。实验表明，这种方式同样有效。

3. 实操流程：从数据到可靠红移目录

3.1 数据准备与特征工程

我们使用的数据是DES DR2与VIPERS PDR-2的交集天区内的匹配星系。对于每个星系，核心特征如下：

• 颜色：计算相邻波段的星等差，如 g-r, r-i, i-z, z-Y。颜色能消除距离的影响，反映星系本身的恒星组成、年龄和尘埃消光等物理特性，是红移估计最主要的依据。
• 星等：我们主要使用i波段自动测光星等（MAG_AUTO），并进行了银河系内消光改正（DERED）。星等信息提供了亮度维度。
• 测光误差：每个波段的测光误差（MAGERR_AUTO）被保留，作为衡量特征测量不确定性的关键信息。

所有特征在输入模型前都进行了标准化处理（减去均值，除以标准差），这是加速神经网络收敛的标准操作。

3.2 构建相似度指标λ：K-d Tree与分类器的协同

这是本项目最具创���性的环节。步骤如下：

1. 构建标签数据集：将VIPERS星系标记为1（相似），将DES中不在VIPERS区域的星系（确保无光谱红移）标记为0（不相似）。这样就构成了一个二分类问题的训练集。
2. 特征空间投影：将所有星系（包括标记的和待预测的DES星系）投影到之前定义的4维或5维颜色-星等空间。
3. K-d Tree近邻搜索：为VIPERS星系点集构建K-d Tree。对于每一个DES星系，通过K-d Tree快速找到其在特征空间中的k个最近邻VIPERS星系（我们实验了k=5, 10, 20等）。
4. 训练分类器：我们尝试了多种分类算法来根据“近邻分布”预测λ值：
   • K近邻分类器：最直观的方法，一个DES星系的λ值可以直接定义为其k个最近邻中VIPERS星系（标签为1）所占的比例。
   • 随机森林、AdaBoost等：我们将每个DES星系的特征，以及其到最近邻VIPERS星系的距离统计量（如平均距离、距离方差等）作为新特征，输入到更复杂的集成分类器中训练。
5. 评估与选择：我们使用ROC曲线和AUC面积来评估不同分类器的性能。如图3所示，K近邻分类器在本问题上取得了最高的AUC分数（0.868），说明这种基于局部近邻比例的直接估计方法非常有效且稳定。因此，我们最终选择K近邻分类器的输出作为每个DES星系的λ值。
6. K-d Tree叶子节点聚合：为了更稳健，我们没有孤立地看待每个星系的λ，而是利用K-d Tree的结构。我们将整个特征空间划分成的最终子区域（叶子节点）内的所有DES星系的λ值取平均，得到该叶子节点的平均λ。如图4所示，这样处理后的λ值在叶子节点编号上有更清晰的排序和分离度，减少了随机波动。

3.3 训练KeraZ红移预测模型

在获得λ值之后，我们开始训练核心的红移预测模型：

1. 数据集划分：将VIPERS星系（拥有真实光谱红移z_spec）随机划分为训练集（70%）、验证集（15%）和测试集（15%）。验证集用于在训练过程中监控模型是否过拟合。
2. 模型构建：使用Keras Sequential API顺序构建网络。输入层节点数等于特征数（4或5）。隐藏层为30个节点，使用tanh激活，紧接着是BatchNormalization()层，并在该层的kernel_regularizer参数中添加l2(0.01)的正则化。输出层为1个节点，使用线性激活函数（因为红移是连续值）。
3. 编译与训练：

   model.compile(optimizer=Adam(learning_rate=0.001), loss='mse', metrics=['mae']) history = model.fit(X_train, z_spec_train, validation_data=(X_val, z_spec_val), epochs=200, batch_size=32, callbacks=[EarlyStopping(patience=10, restore_best_weights=True)])

   这里使用了早停法（EarlyStopping），当验证集损失连续10个epoch不再下降时，自动终止训练并恢复到验证损失最低的模型权重，这是防止过拟合的另一个重要技巧。
4. 预测：用训练好的KeraZ模型预测所有DES星系的测光红移z_phot。

3.4 生成最终可靠红移目录

至此，我们拥有了两样东西：所有DES星系的预测红移z_phot，以及每个星系的可信度指标λ。最终的科学样本需要根据λ进行筛选。我们通过分析λ值与红移估计精度的关系（如图5、6所示），确定了一个经验阈值λ > 0.45。选择这个阈值意味着我们只保留那些在颜色-星等空间中与VIPERS样本高度相似的DES星系，其红移预测结果才被采纳。

如图5所示，在颜色-颜色图中，如果不加λ截断，高、低红移星系点云混杂严重。应用λ>0.45的截断后，大部分低红移（z_phot < 0.59）的星系被清晰地限制在特定的颜色区域内，高、低红移星系的分离开显著提升。这证明λ有效过滤掉了那些因与训练集不相似而导致红移估计混乱的星系。

4. 结果分析与性能评估

4.1 整体性能对比

我们将KeraZ与ANNz2、GPz（仅用颜色、颜色+星等两种输入）在相同的VIPERS测试集上进行了对比。图7展示了红移分布直方图，图8是预测红移与光谱红移的散点图。

从视觉上，几种方法在z < 1.0的范围内都与1:1对角线贴合得较好。KeraZ在z ~ 0.3附近出现了一个微弱的聚集，这可能是由于训练数据中宇宙大尺度结构（即星系在特定红移处成团）导致的模型偏差，这是基于真实观测数据训练难以完全避免的。GPz在z > 1.0的高红移端表现出轻微的偏差。

4.2 量化指标深度解读

我们采用了一套天文学界评估测光红移的标准指标（公式6-9），并在红移区间[0.3, 1.1]内以0.1为宽度分bin统计，结果如图9所示。

• 偏差（Bias）：衡量系统误差。在所有方法中，KeraZ在大部分红移区间表现出与ANNz2、GPz可比甚至更低的偏差。特别是在0.4-0.6这个关键区间，我们的方法偏差显著低于DES Y3 GOLD样本的公开结果。
• 散射（Scatter, σ）与σ68：衡量预测值的离散程度。σ68是指68%的样本所满足的最小误差范围，是更稳健的离散度估计。GPz（尤其是加入星等后）和KeraZ在整个红移范围内都给出了约0.035的散射值，这是一个非常优秀的水平，意味着预测误差的典型值在3.5%左右。
• ** catastrophic outlier率（FRe=0.15）**：衡量严重异常值的比例，定义为|Δz|/(1+z) > 0.15的星系所占百分比。所有方法都将异常率控制在了百分之几的水平，其中GPz-Magnitude在高红移端表现最稳定。

实操心得：理解指标背后的物理不要只看整体指标。分红移区间查看指标至关重要。例如，我们的方法在中间红移（0.5-0.9）表现最佳，这正是VIPERS样本最丰富的区域。在低红移（z<0.4）和高红移（z>1.0），由于训练样本少或特征退化，性能自然下降。λ截断本质上就是帮助我们剔除在这些边缘区域最不可靠的预测，从而提升整体样本的纯净度。

4.3 与DES官方结果的交叉验证

我们将我们的VIPERS匹配样本与DES第三年数据（Y3 GOLD）的公开photo-z结果进行了比较。一个重要的发现是：经过λ > 0.45筛选后的、与VIPERS相似的DES子样本，其红移估计的整体散射和异常率优于完整的DES Y3 GOLD样本。这说明，虽然我们的训练集（VIPERS）很小且有偏，但通过严格的相似性筛选，我们可以从这个有偏训练集中“萃取”出一个预测质量极高的子样本。这对于需要极高红移精度的宇宙学应用（如弱引力透镜）具有重要价值。

5. 常见问题、挑战与解决方案实录

在实际操作中，我们遇到了诸多挑战，以下是其中关键的几个及其解决方案：

1. 类别不平衡问题在训练λ分类器时，标签为1（VIPERS星系）的样本数量远少于标签为0（非VIPERS星系）的样本。这会导致分类器倾向于将所有样本都预测为0，从而AUC虚高但实际分类能力差。

• 我们的应对：我们采用了分层采样（Stratified Sampling）来确保训练集和测试集中正负样本比例一致。同时，在评估时主要依赖ROC-AUC指标，它对类别不平衡相对不敏感。我们也尝试了为少数类样本加权，但发现对于K近邻算法，简单的分层采样已��够。

2. 高维空间距离度量失效在4维或5维空间中，欧氏距离会随着维度升高而变得不直观，所有点对之间的距离都趋于相似，这就是所谓的“维度灾难”。

• 我们的应对：这正是我们系统比较不同距离度量的原因。最终选择闵可夫斯基距离并调整参数p，实际上是在寻找最适合本数据几何结构的度量。此外，K-d Tree本身在高维空间效率也会下降，但由于我们的维度仅为4或5，尚未达到严重影响的临界点。

3. 神经网络训练不稳定初期训练KeraZ时，损失函数震荡剧烈，难以收敛。

• 我们的应对：我们引入了“学习率调度器”（ReduceLROnPlateau），当验证损失停滞时自动降低学习率。结合早停法和批归一化层，最终使训练过程非常平稳。将优化器从SGD改为Adam是解决此问题的第一步，也是最重要的一步。

4. λ阈值的选择依据选择λ > 0.45作为阈值看似主观，需要有客观依据。

• 我们的做法：我们绘制了λ值与红移估计误差|Δz|/(1+z)的关系图。可以清晰地看到，当λ低于0.4-0.5时，误差的分布明显变宽，异常值比例急剧上升。同时，我们参考了ROC曲线上靠近左上角的点所对应的λ值。0.45是一个在保证样本纯度（高精度）和保留样本数量（高效率）之间取得良好平衡的折中点。

5. 计算资源与效率对2.47亿个DES星系进行K-d Tree近邻搜索和神经网络预测，即使算法高效，计算量也极其庞大。

• 我们的优化：
  • 向量化操作：使用NumPy和SciPy的向量化函数替代循环。
  • 分批处理：将DES星系划分为多个批次，依次进行预测，避免内存溢出。
  • 并行化：对K-d Tree的查询过程，利用joblib或multiprocessing进行多进程并行，充分利用多核CPU。
  • GPU加速：Keras模型训练和预测在支持GPU的环境下进行，可提速数十倍。

6. 总结与展望

回顾整个项目，其价值不仅在于实现了一个精度不错的测光红移估计流程，更在于提出并实践了一套“可信度评估”框架。K-d Tree与机器学习分类器的结合，为我们提供了一把度量数据相似性的尺子（λ），让我们能明确知道预测结果在何处站稳了脚跟，在何处是沙上筑塔。这对于将机器学习应用于存在样本选择偏差的真实科学问题，是一个具有普适参考意义的范例。

从技术细节上看，基于Keras的自定义网络（KeraZ）给予了我们极大的灵活性去调整架构和正则化策略，以适配天文数据小样本、低特征维度的特点。简单的单隐藏层网络配合恰当的归一化、正则化，其表现不亚于甚至优于一些更复杂的现成工具。

如果未来要继续推进这项工作，我个人会从以下几个方向入手：首先是探索更先进的深度网络架构，例如引入残差连接或注意力机制，看是否能更好地捕捉颜色与红移之间的非线性关系；其次是尝试将λ的计算与红移预测融合进一个端到端的模型中，例如设计一个多任务学习网络，一个头输出λ，另一个头输出z_phot，让两个任务共享特征提取层并相互促进；最后，可以将这套方法应用到其他巡天数据（如LSST）的交叉认证中，验证其泛化能力。天文数据洪流已至，类似这样将严谨的数据筛选与强大的机器学习模型相结合的思路，或许是从中淘出真金的关键。