1. 项目概述当宇宙学遇上机器学习最近几年我身边不少做天体物理和宇宙学研究的同行聊天时的话题都开始从“哪个巡天项目的数据快释放了”转向了“你试过用哪个神经网络架构处理这个数据吗”。这背后反映的正是观测宇宙学领域一场静悄悄的革命。我们手头的数据量已经从过去的“精雕细琢”变成了现在的“洪水猛兽”。以弱引力透镜为例新一代的宽视场巡天比如欧几里得Euclid、薇拉·鲁宾天文台Vera C. Rubin Observatory的LSST它们产生的数据不再是几张精美的星系图像而是数以十亿计的星系形态数据。传统基于似然函数和马尔可夫链蒙特卡洛MCMC的参数推断方法在面对如此高维、复杂的似然曲面时计算成本已经高到令人望而却步。这就好比以前我们是在一条清晰的小路上找最高点现在却是在一片云雾缭绕、沟壑纵横的群山里寻找而且这座山还在不断变大。“弱引力透镜宇宙学挑战”这个项目正是要直面这个新常态下的核心痛点。它的目标很明确利用机器学习特别是深度学习的方法来革新宇宙学参数推断的整个流程并解决一个随之而来的、更隐蔽的问题——分布外检测。简单来说我们要做两件事第一训练一个模型让它能像经验丰富的老天文学家一样从海量的弱透镜数据中“看”出背后的宇宙学参数比如物质密度参数Ω_m、物质功率谱振幅σ_8等第二也是更关键的一步给这个模型装上“警报器”当它遇到训练数据中从未出现过的、超出其认知范围的宇宙模型或系统效应时它能主动“举手”说“这个我不确定可能有问题。”这不仅仅是把机器学习当做一个更快的计算器来用。传统MCMC虽然慢但它的整个采样过程是透明的我们能清晰地看到参数空间的探索路径和后验分布的形状。而深度学习模型就像一个黑盒它给出的预测值很快但我们很难知道这个预测的置信度有多高更不知道当输入数据来自一个完全不同的宇宙时它会给出多么荒谬而又自信的答案。因此这个项目的真正挑战在于构建一个既快又准同时还“自知之明”的智能推断系统。它不仅要会算还要知道自己什么时候可能算错这对于将机器学习可靠地应用于真实的科学发现至关重要。2. 核心思路与架构设计面对这个挑战一个鲁棒的解决方案不能只靠一个单一的模型。我们需要设计一套组合拳将不同的机器学习范式串联起来形成一个完整的分析流水线。经过多次迭代和尝试我总结出一套相对稳定有效的架构其核心可以分为三个层次仿真与数据引擎、核心推断模型、以及不确定性量化与分布外检测模块。2.1 仿真与数据引擎构建“数字宇宙实验室”一切机器学习项目都始于数据而对于宇宙学我们无法进行物理实验只能依靠数值模拟。我们的数据引擎需要生成两大类数据训练与验证集基于当前主流的ΛCDM宇宙学模型及其合理的参数变化范围生成大量的弱引力透镜观测数据。这通常通过N体模拟结合光线追踪算法来实现生成的是二维的宇宙剪切场shear field或收敛场convergence field图。分布外测试集这是项目的关键。我们需要故意生成一些“非标准”的宇宙学数据。例如非标准宇宙模型引入早期暗能量、中微子质量有较大偏差、或者修改引力理论如f(R)引力下的模拟数据。极端系统效应在模拟数据中加入远超预期的点扩散函数PSF误差、光度红移误差、或者天体物理前景如星系内禀对齐的污染。这个数据工厂的质量直接决定了整个项目的上限。我们使用的工具链通常包括CAMEL、TRIANGLE等宇宙学模拟代码或者基于PyCosmo、CCL等理论计算工具快速生成理论预言。数据预处理环节我们会将剪切场转换为多尺度的统计量最常用的是两点相关函数2PCF或峰统计Peak Counts有时也会直接使用像素化的收敛图作为神经网络的输入。这一步的降维和特征工程至关重要它决定了模型学习效率的高低。注意仿真数据的真实性是生命线。必须尽可能真实地模拟观测效应如噪声、掩模、选择函数等。一个在“干净”模拟数据上表现完美的模型在真实数据面前可能不堪一击。我们通常采用“向前建模”的范式即用模拟数据训练模型然后将其直接应用于经过相同系统效应模拟的真实数据上。2.2 核心推断模型从数据到参数的“翻译官”这是架构的心脏。我们的目标是将高维的观测数据如2PCF的多个bin值映射到低维的宇宙学参数空间。这里有几个主流选择深度神经网络回归器最直接的思路。构建一个多层全连接网络或卷积网络如果输入是图像以观测统计量为输入直接输出宇宙学参数的点估计值如Ω_m, σ_8。这种方法简单快捷但缺乏对预测不确定性的估计。似然函数估计器这是更科学的方法。我们训练一个模型来估计给定观测数据下参数的后验概率分布P(θ|d)。常用方法有归一化流这是一种强大的概率生成模型。我们训练一个可逆的神经网络将一个简单的基分布如高斯分布变换成复杂的后验分布。一旦训练完成我们可以通过从基分布采样并经过网络变换快速生成后验样本。sbi模拟基于推断工具箱是这方面的利器。密度估计网络例如掩码自回归流MAF或实值非体积保持变换Real NVP。它们直接对后验密度进行建模。仿真器与贝叶斯推断另一种思路是训练一个仿真器它学习从宇宙学参数θ到观测数据d的快速前向映射即替代昂贵的数值模拟。然后将这个快速仿真器嵌入到传统的MCMC或嵌套采样框架中。虽然推断速度仍受采样算法限制但每次似然计算的速度得到了百万倍的提升。在实际项目中我通常会先训练一个深度回归网络作为基线模型和快速检查工具因为它训练快能快速验证数据管道和特征的有效性。然后重点投入资源构建基于归一化流的似然/后验估计器因为它能提供完整的概率推断是进行严谨科学分析的基础。2.3 不确定性量化与分布外检测模型的“自知之明”这是本项目区别于普通机器学习应用的核心。一个只会给出点估计的模型在科学上是危险的。我们需要量化两种不确定性认知不确定性源于模型自身对数据认知的不足。对于分布内的数据这种不确定性应该较小对于分布外数据应该急剧增大。偶然不确定性数据中固有的噪声。为了实现分布外检测我们通常在模型架构或训练策略上做文章集成学习训练多个结构相同但初始化不同的模型深度集成。对于同一个输入如果所有模型的预测结果高度一致则认知不确定性低如果分歧很大则认知不确定性高可能遇到了分布外样本。计算预测结果的方差即可作为不确定性指标。贝叶斯神经网络将网络权重视为概率分布而非固定值。在预测时通过对权重进行采样如使用MC Dropout来得到预测分布。预测分布的离散程度反映了认知不确定性。专门的不确定性估计模块在回归网络的末端除了输出参数预测值额外增加一个分支来输出预测的方差异方差不确定性。或者训练一个独立的“新奇性检测”分类器判断输入数据是否来自训练分布。在我们的流水线中深度集成因其实现简单、效果稳定而成为首选。我们将5-10个归一化流模型组成一个委员会。当遇到新的观测数据时委员会成员会给出各自的后验分布。如果这些后验分布的均值相近且形状紧凑说明模型有信心如果均值分散、形状各异甚至完全偏离物理合理范围这就是一个强烈的分布外警报。3. 关键实现步骤与技术细节下面我将以构建一个基于归一化流和深度集成的完整系统为例拆解其中的关键实现步骤。我们假设数据已经预处理为包含N个bin的两点相关函数向量。3.1 数据准备与仿真流水线首先我们需要一个可重复、可扩展的数据生成流程。这里我倾向于使用JAX或PyTorch来实现一个向量化的模拟器即使它比高精度的N体模拟简化很多。import jax import jax.numpy as jnp import numpyro.distributions as dist from jax import random def simulator(theta, key, n_bins10): 一个简化的弱透镜两点相关函数模拟器。 theta: 宇宙学参数例如 [Omega_m, sigma_8] key: JAX随机数密钥 n_bins: 相关函数bin的数量 返回: 模拟的观测数据向量 omega_m, sigma_8 theta # 1. 基于理论模型计算线性功率谱这里极度简化 # 实际中会调用像CCL、CLASS这样的理论代码 ell jnp.logspace(1, 3, n_bins) # 角多极矩 # 一个简单的参数化形式仅用于演示 theoretical_signal sigma_8**2 * omega_m**0.8 * ell**(-1.2) # 2. 加入观测效应高斯噪声和系统误差偏置 noise_key, sys_key random.split(key) # 观测噪声与信号幅度相关 noise random.normal(noise_key, (n_bins,)) * 0.05 * theoretical_signal # 简单的乘性系统误差如校准误差 sys_bias 1.0 random.normal(sys_key, (n_bins,)) * 0.02 # 3. 生成“观测到的”数据点 observed_data theoretical_signal * sys_bias noise return observed_data # 生成一批训练数据 prior dist.Uniform(lowjnp.array([0.1, 0.6]), highjnp.array([0.5, 1.0])) # Omega_m, sigma_8的先验 num_simulations 100000 theta_train prior.sample(random.PRNGKey(0), sample_shape(num_simulations,)) # 向量化模拟一次生成所有数据 keys random.split(random.PRNGKey(1), num_simulations) x_train jax.vmap(simulator)(theta_train, keys) # x_train.shape: (100000, 10)这个模拟器虽然简单但包含了信号、噪声和系统误差的核心要素。在生产环境中simulator函数会被替换为与真实观测条件匹配的高保真模拟接口。3.2 归一化流模型的构建与训练接下来我们使用sbi库来构建一个后验估计器。我们选择神经后验估计NPE的方法。import torch import sbi from sbi import analysis, utils from sbi.inference import SNPE_C # 1. 将JAX数组转换为PyTorch张量sbi基于PyTorch theta_tensor torch.as_tensor(theta_train, dtypetorch.float32) x_tensor torch.as_tensor(x_train, dtypetorch.float32) # 2. 定义先验分布需与仿真先验一致 from torch.distributions import Uniform prior_low torch.tensor([0.1, 0.6]) prior_high torch.tensor([0.5, 1.0]) prior Uniform(lowprior_low, highprior_high) # 3. 构建神经网络一个条件归一化流 # 我们使用嵌入网络embedding_net先对高维观测数据x进行编码 embedding_net torch.nn.Sequential( torch.nn.Linear(x_tensor.shape[1], 50), torch.nn.ReLU(), torch.nn.Linear(50, 20), torch.nn.ReLU(), ) # 4. 初始化推理算法SNPE-C inference SNPE_C(priorprior, density_estimatormaf, embedding_netembedding_net) # 5. 训练归一化流模型 # 注意这里我们只训练一个模型深度集成需要重复此步骤多次 num_rounds 2 # 多轮训练可以提升性能 posteriors [] proposal prior for round_num in range(num_rounds): # 如果是第一轮使用原始先验采样后续轮次使用上一轮的后验作为提议分布 if round_num 0: theta theta_tensor x x_tensor else: # 从上一轮的后验中采样新的参数进行仿真此例中我们已有一大批数据故跳过 # 在实际主动学习中这里会调用仿真器生成新数据 pass # 训练密度估计器 density_estimator inference.append_simulations(theta, x).train(training_batch_size100, max_num_epochs100) # 基于训练好的估计器构建后验 posterior inference.build_posterior(density_estimator) posteriors.append(posterior) proposal posterior.set_default_x(x_tensor[0:1]) # 为下一轮设置一个参考观测 final_posterior posteriors[-1]训练完成后final_posterior就是一个可以快速计算后验分布的对象。给定一个新的观测数据x_o我们可以通过posterior.sample((10000,), xx_o)快速得到一万个后验样本用于计算均值、置信区间等。3.3 深度集成与不确定性量化单一模型的不确定性估计可能不可靠。我们需要构建一个模型委员会。num_models 5 ensemble_posteriors [] for i in range(num_models): print(f训练集成模型 {i1}/{num_models}) # 每次使用不同的随机种子初始化网络和打乱数据 torch.manual_seed(i) # 重新初始化推理对象和网络 embedding_net torch.nn.Sequential( torch.nn.Linear(x_tensor.shape[1], 50), torch.nn.ReLU(), torch.nn.Linear(50, 20), torch.nn.ReLU(), ) inference SNPE_C(priorprior, density_estimatormaf, embedding_netembedding_net) # 可以随机打乱数据顺序 indices torch.randperm(len(theta_tensor)) theta_shuffled theta_tensor[indices] x_shuffled x_tensor[indices] density_estimator inference.append_simulations(theta_shuffled, x_shuffled).train(training_batch_size100, max_num_epochs100) posterior inference.build_posterior(density_estimator) ensemble_posteriors.append(posterior) # 定义集成预测函数 def ensemble_predict(x_observation, num_samples5000): 返回集成模型对给定观测的预测结果。 all_samples [] for i, post in enumerate(ensemble_posteriors): samples post.sample((num_samples,), xx_observation, show_progress_barsFalse) all_samples.append(samples.numpy()) # all_samples 是一个列表包含5个数组每个形状为 (5000, 2) return np.array(all_samples) # 形状: (5, 5000, 2)现在对于任何一个新观测x_oensemble_predict(x_o)会返回5个后验样本集。我们可以通过计算这5个后验分布的均值之间的标准差来量化模型委员会的“分歧度”这直接反映了认知不确定性。3.4 分布外检测的逻辑实现有了集成模型输出的后验样本我们就可以设计分布外检测的指标了。一个简单有效的指标是后验预测p值或分歧度评分。import numpy as np def compute_disagreement_score(x_observation): 计算集成模型对给定观测的预测分歧度。 返回一个标量分数分数越高越可能是分布外样本。 ensemble_samples ensemble_predict(x_observation, num_samples2000) # (5, 2000, 2) # 计算每个模型预测的参数均值 model_means np.mean(ensemble_samples, axis1) # (5, 2) # 计算所有模型均值之间的标准差跨模型维度 std_across_models np.std(model_means, axis0) # (2,) # 取两个参数不确定性的均方根作为综合分歧度分数 disagreement_score np.sqrt(np.mean(std_across_models**2)) return disagreement_score def is_out_of_distribution(x_observation, threshold0.05): 基于分歧度判断是否为分布外样本。 threshold阈值需要通过验证集上的分布内数据来确定。 score compute_disagreement_score(x_observation) # 假设我们通过验证集计算了分布内数据分歧度的95%分位数作为阈值 if score threshold: return True, score else: return False, score在实际操作中threshold的确定至关重要。我们需要在一个干净的、分布内的验证集上运行compute_disagreement_score然后取其结果的第95或99百分位数作为警报阈值。当处理真实观测数据时如果计算出的分歧度分数超过该阈值分析流程就会自动标记该数据点提醒科学家需要人工介入检查可能是遇到了新物理也可能是未被模拟的系统误差。4. 实战挑战与调优心得将这套架构应用于实际项目时会遇到许多在理论设计时未曾预料的挑战。下面分享几个关键的“踩坑”经验和调优心得。4.1 仿真与现实的差距系统误差的注入艺术最大的挑战来自于仿真数据与真实数据之间的“域鸿沟”。一个在完美模拟数据上准确率99%的模型在真实数据面前可能毫无用处。心得1正向建模必须贯穿始终。不要在“干净”的信号上训练模型然后期望它能处理充满噪声和缺陷的真实数据。你的仿真流水线必须包含所有已知的主要系统效应并且要以一种可微的、物理上合理的方式注入。例如模拟PSF效应时不要只是加个高斯模糊而要基于仪器光学模型生成真实的PSF图案并进行卷积。心得2为未知的误差留出空间。我们可以在仿真中引入一些“ nuisance parameters ”干扰参数比如一个全局的校准偏移因子、一个红移误差的幅度参数等。在训练时将这些参数与宇宙学参数一起作为神经网络的输入或条件。这样模型就学会了在存在这些干扰的情况下进行推断。在应用时这些干扰参数可以被边缘化掉。实操技巧建立一个系统误差的“武器库”并以模块化的方式集成到仿真器中。这样你可以轻松地开启或关闭某些效应来评估它们对模型性能的影响。使用JAX或PyTorch的自动微分特性甚至可以尝试让模型学习如何校正某些系统误差。4.2 模型校准与后验检验黑盒模型给出的后验分布是否可靠我们需要严格校准。诊断工具模拟-模拟检验。这是最有力的工具。从先验中抽取一组新的参数θ*用仿真器生成对应的数据d*。然后用训练好的模型对d*进行推断得到后验分布。如果模型是校准良好的那么真实的参数θ*应该以一定的概率例如68%落在推断出的后验分布的1σ置信区间内。重复这个过程成百上千次统计覆盖率。理想情况下68%的置信区间应该包含大约68%的真实参数值。常见问题过度自信的后验。如果模型给出的后验分布过于狭窄导致覆盖率过低说明模型低估了不确定性。这可能是因为训练数据不足模型过拟合了。神经网络表达能力过强记住了数据而非学习了映射关系。仿真中缺失了某些关键噪声或系统误差。调优方法增加训练数据这是最直接的方法。使用更简单的网络减少层数或神经元数量。在训练中注入噪声在输入数据或网络激活中加入随机噪声作为一种正则化手段。采用贝叶斯神经网络或深度集成如前所述它们天生能提供不确定性估计。4.3 计算资源与效率的平衡高保真宇宙学模拟极其昂贵生成十万级别的训练样本可能消耗数百万CPU小时。如何在有限资源下取得平衡策略1代理模拟器与迁移学习。先使用一个快速但精度较低的模拟器如基于扰动理论生成海量数据预训练一个模型。然后用少量高精度模拟数据对这个预训练模型进行微调。这通常比直接用少量高精度数据训练效果更好。策略2主动学习。不是随机生成训练参数而是让模型自己决定下一步在哪里仿真。基本流程是1) 用现有数据训练一个初始模型2) 用这个模型在当前先验空间内寻找那些模型预测不确定性最高的参数点3) 对这些点进行仿真将新数据加入训练集4) 重复。这能极大提高数据利用效率。实操心得在项目初期不要追求最高精度的模拟。先用一个中等精度的模拟器快速构建起整个机器学习流水线验证想法的可行性并进行超参数搜索。待整个流程跑通后再将计算资源投入到最终的高精度数据生成上。使用JAX的jit编译和向量化可以极大加速仿真和训练过程。4.4 分布外检测阈值的设定如何设定那个决定“是否报警”的阈值这没有黄金标准。方法基于验证集的ROC曲线。我们需要一个既包含分布内数据ID也包含分布外数据OOD的验证集。OOD数据可以是我们故意构造的非标准宇宙学模拟。然后我们计算验证集中每个样本的分歧度分数并绘制接收者操作特征曲线。通过调整阈值在误报率将ID数据误判为OOD和检出率正确识别OOD数据之间取得平衡。注意阈值的选取与科学目标强相关。如果你在进行盲分析寻找新物理的蛛丝马迹你可能愿意承受较高的误报率以换取更高的检出率宁错杀不放过。如果你只是用模型进行标准参数的常规推断那么可能会设定一个严格的阈值只对极端异常值报警。最终建议不要只依赖一个单一的阈值或指标。将分歧度分数、后验预测p值、以及模型输出的后验分布形态是否多峰、是否超出物理合理范围等多个指标结合起来形成一个综合的“异常评分”。同时永远保留人工检查的环节最终的判断权应交由领域专家。
机器学习在宇宙学参数推断中的应用：从归一化流到分布外检测
1. 项目概述当宇宙学遇上机器学习最近几年我身边不少做天体物理和宇宙学研究的同行聊天时的话题都开始从“哪个巡天项目的数据快释放了”转向了“你试过用哪个神经网络架构处理这个数据吗”。这背后反映的正是观测宇宙学领域一场静悄悄的革命。我们手头的数据量已经从过去的“精雕细琢”变成了现在的“洪水猛兽”。以弱引力透镜为例新一代的宽视场巡天比如欧几里得Euclid、薇拉·鲁宾天文台Vera C. Rubin Observatory的LSST它们产生的数据不再是几张精美的星系图像而是数以十亿计的星系形态数据。传统基于似然函数和马尔可夫链蒙特卡洛MCMC的参数推断方法在面对如此高维、复杂的似然曲面时计算成本已经高到令人望而却步。这就好比以前我们是在一条清晰的小路上找最高点现在却是在一片云雾缭绕、沟壑纵横的群山里寻找而且这座山还在不断变大。“弱引力透镜宇宙学挑战”这个项目正是要直面这个新常态下的核心痛点。它的目标很明确利用机器学习特别是深度学习的方法来革新宇宙学参数推断的整个流程并解决一个随之而来的、更隐蔽的问题——分布外检测。简单来说我们要做两件事第一训练一个模型让它能像经验丰富的老天文学家一样从海量的弱透镜数据中“看”出背后的宇宙学参数比如物质密度参数Ω_m、物质功率谱振幅σ_8等第二也是更关键的一步给这个模型装上“警报器”当它遇到训练数据中从未出现过的、超出其认知范围的宇宙模型或系统效应时它能主动“举手”说“这个我不确定可能有问题。”这不仅仅是把机器学习当做一个更快的计算器来用。传统MCMC虽然慢但它的整个采样过程是透明的我们能清晰地看到参数空间的探索路径和后验分布的形状。而深度学习模型就像一个黑盒它给出的预测值很快但我们很难知道这个预测的置信度有多高更不知道当输入数据来自一个完全不同的宇宙时它会给出多么荒谬而又自信的答案。因此这个项目的真正挑战在于构建一个既快又准同时还“自知之明”的智能推断系统。它不仅要会算还要知道自己什么时候可能算错这对于将机器学习可靠地应用于真实的科学发现至关重要。2. 核心思路与架构设计面对这个挑战一个鲁棒的解决方案不能只靠一个单一的模型。我们需要设计一套组合拳将不同的机器学习范式串联起来形成一个完整的分析流水线。经过多次迭代和尝试我总结出一套相对稳定有效的架构其核心可以分为三个层次仿真与数据引擎、核心推断模型、以及不确定性量化与分布外检测模块。2.1 仿真与数据引擎构建“数字宇宙实验室”一切机器学习项目都始于数据而对于宇宙学我们无法进行物理实验只能依靠数值模拟。我们的数据引擎需要生成两大类数据训练与验证集基于当前主流的ΛCDM宇宙学模型及其合理的参数变化范围生成大量的弱引力透镜观测数据。这通常通过N体模拟结合光线追踪算法来实现生成的是二维的宇宙剪切场shear field或收敛场convergence field图。分布外测试集这是项目的关键。我们需要故意生成一些“非标准”的宇宙学数据。例如非标准宇宙模型引入早期暗能量、中微子质量有较大偏差、或者修改引力理论如f(R)引力下的模拟数据。极端系统效应在模拟数据中加入远超预期的点扩散函数PSF误差、光度红移误差、或者天体物理前景如星系内禀对齐的污染。这个数据工厂的质量直接决定了整个项目的上限。我们使用的工具链通常包括CAMEL、TRIANGLE等宇宙学模拟代码或者基于PyCosmo、CCL等理论计算工具快速生成理论预言。数据预处理环节我们会将剪切场转换为多尺度的统计量最常用的是两点相关函数2PCF或峰统计Peak Counts有时也会直接使用像素化的收敛图作为神经网络的输入。这一步的降维和特征工程至关重要它决定了模型学习效率的高低。注意仿真数据的真实性是生命线。必须尽可能真实地模拟观测效应如噪声、掩模、选择函数等。一个在“干净”模拟数据上表现完美的模型在真实数据面前可能不堪一击。我们通常采用“向前建模”的范式即用模拟数据训练模型然后将其直接应用于经过相同系统效应模拟的真实数据上。2.2 核心推断模型从数据到参数的“翻译官”这是架构的心脏。我们的目标是将高维的观测数据如2PCF的多个bin值映射到低维的宇宙学参数空间。这里有几个主流选择深度神经网络回归器最直接的思路。构建一个多层全连接网络或卷积网络如果输入是图像以观测统计量为输入直接输出宇宙学参数的点估计值如Ω_m, σ_8。这种方法简单快捷但缺乏对预测不确定性的估计。似然函数估计器这是更科学的方法。我们训练一个模型来估计给定观测数据下参数的后验概率分布P(θ|d)。常用方法有归一化流这是一种强大的概率生成模型。我们训练一个可逆的神经网络将一个简单的基分布如高斯分布变换成复杂的后验分布。一旦训练完成我们可以通过从基分布采样并经过网络变换快速生成后验样本。sbi模拟基于推断工具箱是这方面的利器。密度估计网络例如掩码自回归流MAF或实值非体积保持变换Real NVP。它们直接对后验密度进行建模。仿真器与贝叶斯推断另一种思路是训练一个仿真器它学习从宇宙学参数θ到观测数据d的快速前向映射即替代昂贵的数值模拟。然后将这个快速仿真器嵌入到传统的MCMC或嵌套采样框架中。虽然推断速度仍受采样算法限制但每次似然计算的速度得到了百万倍的提升。在实际项目中我通常会先训练一个深度回归网络作为基线模型和快速检查工具因为它训练快能快速验证数据管道和特征的有效性。然后重点投入资源构建基于归一化流的似然/后验估计器因为它能提供完整的概率推断是进行严谨科学分析的基础。2.3 不确定性量化与分布外检测模型的“自知之明”这是本项目区别于普通机器学习应用的核心。一个只会给出点估计的模型在科学上是危险的。我们需要量化两种不确定性认知不确定性源于模型自身对数据认知的不足。对于分布内的数据这种不确定性应该较小对于分布外数据应该急剧增大。偶然不确定性数据中固有的噪声。为了实现分布外检测我们通常在模型架构或训练策略上做文章集成学习训练多个结构相同但初始化不同的模型深度集成。对于同一个输入如果所有模型的预测结果高度一致则认知不确定性低如果分歧很大则认知不确定性高可能遇到了分布外样本。计算预测结果的方差即可作为不确定性指标。贝叶斯神经网络将网络权重视为概率分布而非固定值。在预测时通过对权重进行采样如使用MC Dropout来得到预测分布。预测分布的离散程度反映了认知不确定性。专门的不确定性估计模块在回归网络的末端除了输出参数预测值额外增加一个分支来输出预测的方差异方差不确定性。或者训练一个独立的“新奇性检测”分类器判断输入数据是否来自训练分布。在我们的流水线中深度集成因其实现简单、效果稳定而成为首选。我们将5-10个归一化流模型组成一个委员会。当遇到新的观测数据时委员会成员会给出各自的后验分布。如果这些后验分布的均值相近且形状紧凑说明模型有信心如果均值分散、形状各异甚至完全偏离物理合理范围这就是一个强烈的分布外警报。3. 关键实现步骤与技术细节下面我将以构建一个基于归一化流和深度集成的完整系统为例拆解其中的关键实现步骤。我们假设数据已经预处理为包含N个bin的两点相关函数向量。3.1 数据准备与仿真流水线首先我们需要一个可重复、可扩展的数据生成流程。这里我倾向于使用JAX或PyTorch来实现一个向量化的模拟器即使它比高精度的N体模拟简化很多。import jax import jax.numpy as jnp import numpyro.distributions as dist from jax import random def simulator(theta, key, n_bins10): 一个简化的弱透镜两点相关函数模拟器。 theta: 宇宙学参数例如 [Omega_m, sigma_8] key: JAX随机数密钥 n_bins: 相关函数bin的数量 返回: 模拟的观测数据向量 omega_m, sigma_8 theta # 1. 基于理论模型计算线性功率谱这里极度简化 # 实际中会调用像CCL、CLASS这样的理论代码 ell jnp.logspace(1, 3, n_bins) # 角多极矩 # 一个简单的参数化形式仅用于演示 theoretical_signal sigma_8**2 * omega_m**0.8 * ell**(-1.2) # 2. 加入观测效应高斯噪声和系统误差偏置 noise_key, sys_key random.split(key) # 观测噪声与信号幅度相关 noise random.normal(noise_key, (n_bins,)) * 0.05 * theoretical_signal # 简单的乘性系统误差如校准误差 sys_bias 1.0 random.normal(sys_key, (n_bins,)) * 0.02 # 3. 生成“观测到的”数据点 observed_data theoretical_signal * sys_bias noise return observed_data # 生成一批训练数据 prior dist.Uniform(lowjnp.array([0.1, 0.6]), highjnp.array([0.5, 1.0])) # Omega_m, sigma_8的先验 num_simulations 100000 theta_train prior.sample(random.PRNGKey(0), sample_shape(num_simulations,)) # 向量化模拟一次生成所有数据 keys random.split(random.PRNGKey(1), num_simulations) x_train jax.vmap(simulator)(theta_train, keys) # x_train.shape: (100000, 10)这个模拟器虽然简单但包含了信号、噪声和系统误差的核心要素。在生产环境中simulator函数会被替换为与真实观测条件匹配的高保真模拟接口。3.2 归一化流模型的构建与训练接下来我们使用sbi库来构建一个后验估计器。我们选择神经后验估计NPE的方法。import torch import sbi from sbi import analysis, utils from sbi.inference import SNPE_C # 1. 将JAX数组转换为PyTorch张量sbi基于PyTorch theta_tensor torch.as_tensor(theta_train, dtypetorch.float32) x_tensor torch.as_tensor(x_train, dtypetorch.float32) # 2. 定义先验分布需与仿真先验一致 from torch.distributions import Uniform prior_low torch.tensor([0.1, 0.6]) prior_high torch.tensor([0.5, 1.0]) prior Uniform(lowprior_low, highprior_high) # 3. 构建神经网络一个条件归一化流 # 我们使用嵌入网络embedding_net先对高维观测数据x进行编码 embedding_net torch.nn.Sequential( torch.nn.Linear(x_tensor.shape[1], 50), torch.nn.ReLU(), torch.nn.Linear(50, 20), torch.nn.ReLU(), ) # 4. 初始化推理算法SNPE-C inference SNPE_C(priorprior, density_estimatormaf, embedding_netembedding_net) # 5. 训练归一化流模型 # 注意这里我们只训练一个模型深度集成需要重复此步骤多次 num_rounds 2 # 多轮训练可以提升性能 posteriors [] proposal prior for round_num in range(num_rounds): # 如果是第一轮使用原始先验采样后续轮次使用上一轮的后验作为提议分布 if round_num 0: theta theta_tensor x x_tensor else: # 从上一轮的后验中采样新的参数进行仿真此例中我们已有一大批数据故跳过 # 在实际主动学习中这里会调用仿真器生成新数据 pass # 训练密度估计器 density_estimator inference.append_simulations(theta, x).train(training_batch_size100, max_num_epochs100) # 基于训练好的估计器构建后验 posterior inference.build_posterior(density_estimator) posteriors.append(posterior) proposal posterior.set_default_x(x_tensor[0:1]) # 为下一轮设置一个参考观测 final_posterior posteriors[-1]训练完成后final_posterior就是一个可以快速计算后验分布的对象。给定一个新的观测数据x_o我们可以通过posterior.sample((10000,), xx_o)快速得到一万个后验样本用于计算均值、置信区间等。3.3 深度集成与不确定性量化单一模型的不确定性估计可能不可靠。我们需要构建一个模型委员会。num_models 5 ensemble_posteriors [] for i in range(num_models): print(f训练集成模型 {i1}/{num_models}) # 每次使用不同的随机种子初始化网络和打乱数据 torch.manual_seed(i) # 重新初始化推理对象和网络 embedding_net torch.nn.Sequential( torch.nn.Linear(x_tensor.shape[1], 50), torch.nn.ReLU(), torch.nn.Linear(50, 20), torch.nn.ReLU(), ) inference SNPE_C(priorprior, density_estimatormaf, embedding_netembedding_net) # 可以随机打乱数据顺序 indices torch.randperm(len(theta_tensor)) theta_shuffled theta_tensor[indices] x_shuffled x_tensor[indices] density_estimator inference.append_simulations(theta_shuffled, x_shuffled).train(training_batch_size100, max_num_epochs100) posterior inference.build_posterior(density_estimator) ensemble_posteriors.append(posterior) # 定义集成预测函数 def ensemble_predict(x_observation, num_samples5000): 返回集成模型对给定观测的预测结果。 all_samples [] for i, post in enumerate(ensemble_posteriors): samples post.sample((num_samples,), xx_observation, show_progress_barsFalse) all_samples.append(samples.numpy()) # all_samples 是一个列表包含5个数组每个形状为 (5000, 2) return np.array(all_samples) # 形状: (5, 5000, 2)现在对于任何一个新观测x_oensemble_predict(x_o)会返回5个后验样本集。我们可以通过计算这5个后验分布的均值之间的标准差来量化模型委员会的“分歧度”这直接反映了认知不确定性。3.4 分布外检测的逻辑实现有了集成模型输出的后验样本我们就可以设计分布外检测的指标了。一个简单有效的指标是后验预测p值或分歧度评分。import numpy as np def compute_disagreement_score(x_observation): 计算集成模型对给定观测的预测分歧度。 返回一个标量分数分数越高越可能是分布外样本。 ensemble_samples ensemble_predict(x_observation, num_samples2000) # (5, 2000, 2) # 计算每个模型预测的参数均值 model_means np.mean(ensemble_samples, axis1) # (5, 2) # 计算所有模型均值之间的标准差跨模型维度 std_across_models np.std(model_means, axis0) # (2,) # 取两个参数不确定性的均方根作为综合分歧度分数 disagreement_score np.sqrt(np.mean(std_across_models**2)) return disagreement_score def is_out_of_distribution(x_observation, threshold0.05): 基于分歧度判断是否为分布外样本。 threshold阈值需要通过验证集上的分布内数据来确定。 score compute_disagreement_score(x_observation) # 假设我们通过验证集计算了分布内数据分歧度的95%分位数作为阈值 if score threshold: return True, score else: return False, score在实际操作中threshold的确定至关重要。我们需要在一个干净的、分布内的验证集上运行compute_disagreement_score然后取其结果的第95或99百分位数作为警报阈值。当处理真实观测数据时如果计算出的分歧度分数超过该阈值分析流程就会自动标记该数据点提醒科学家需要人工介入检查可能是遇到了新物理也可能是未被模拟的系统误差。4. 实战挑战与调优心得将这套架构应用于实际项目时会遇到许多在理论设计时未曾预料的挑战。下面分享几个关键的“踩坑”经验和调优心得。4.1 仿真与现实的差距系统误差的注入艺术最大的挑战来自于仿真数据与真实数据之间的“域鸿沟”。一个在完美模拟数据上准确率99%的模型在真实数据面前可能毫无用处。心得1正向建模必须贯穿始终。不要在“干净”的信号上训练模型然后期望它能处理充满噪声和缺陷的真实数据。你的仿真流水线必须包含所有已知的主要系统效应并且要以一种可微的、物理上合理的方式注入。例如模拟PSF效应时不要只是加个高斯模糊而要基于仪器光学模型生成真实的PSF图案并进行卷积。心得2为未知的误差留出空间。我们可以在仿真中引入一些“ nuisance parameters ”干扰参数比如一个全局的校准偏移因子、一个红移误差的幅度参数等。在训练时将这些参数与宇宙学参数一起作为神经网络的输入或条件。这样模型就学会了在存在这些干扰的情况下进行推断。在应用时这些干扰参数可以被边缘化掉。实操技巧建立一个系统误差的“武器库”并以模块化的方式集成到仿真器中。这样你可以轻松地开启或关闭某些效应来评估它们对模型性能的影响。使用JAX或PyTorch的自动微分特性甚至可以尝试让模型学习如何校正某些系统误差。4.2 模型校准与后验检验黑盒模型给出的后验分布是否可靠我们需要严格校准。诊断工具模拟-模拟检验。这是最有力的工具。从先验中抽取一组新的参数θ*用仿真器生成对应的数据d*。然后用训练好的模型对d*进行推断得到后验分布。如果模型是校准良好的那么真实的参数θ*应该以一定的概率例如68%落在推断出的后验分布的1σ置信区间内。重复这个过程成百上千次统计覆盖率。理想情况下68%的置信区间应该包含大约68%的真实参数值。常见问题过度自信的后验。如果模型给出的后验分布过于狭窄导致覆盖率过低说明模型低估了不确定性。这可能是因为训练数据不足模型过拟合了。神经网络表达能力过强记住了数据而非学习了映射关系。仿真中缺失了某些关键噪声或系统误差。调优方法增加训练数据这是最直接的方法。使用更简单的网络减少层数或神经元数量。在训练中注入噪声在输入数据或网络激活中加入随机噪声作为一种正则化手段。采用贝叶斯神经网络或深度集成如前所述它们天生能提供不确定性估计。4.3 计算资源与效率的平衡高保真宇宙学模拟极其昂贵生成十万级别的训练样本可能消耗数百万CPU小时。如何在有限资源下取得平衡策略1代理模拟器与迁移学习。先使用一个快速但精度较低的模拟器如基于扰动理论生成海量数据预训练一个模型。然后用少量高精度模拟数据对这个预训练模型进行微调。这通常比直接用少量高精度数据训练效果更好。策略2主动学习。不是随机生成训练参数而是让模型自己决定下一步在哪里仿真。基本流程是1) 用现有数据训练一个初始模型2) 用这个模型在当前先验空间内寻找那些模型预测不确定性最高的参数点3) 对这些点进行仿真将新数据加入训练集4) 重复。这能极大提高数据利用效率。实操心得在项目初期不要追求最高精度的模拟。先用一个中等精度的模拟器快速构建起整个机器学习流水线验证想法的可行性并进行超参数搜索。待整个流程跑通后再将计算资源投入到最终的高精度数据生成上。使用JAX的jit编译和向量化可以极大加速仿真和训练过程。4.4 分布外检测阈值的设定如何设定那个决定“是否报警”的阈值这没有黄金标准。方法基于验证集的ROC曲线。我们需要一个既包含分布内数据ID也包含分布外数据OOD的验证集。OOD数据可以是我们故意构造的非标准宇宙学模拟。然后我们计算验证集中每个样本的分歧度分数并绘制接收者操作特征曲线。通过调整阈值在误报率将ID数据误判为OOD和检出率正确识别OOD数据之间取得平衡。注意阈值的选取与科学目标强相关。如果你在进行盲分析寻找新物理的蛛丝马迹你可能愿意承受较高的误报率以换取更高的检出率宁错杀不放过。如果你只是用模型进行标准参数的常规推断那么可能会设定一个严格的阈值只对极端异常值报警。最终建议不要只依赖一个单一的阈值或指标。将分歧度分数、后验预测p值、以及模型输出的后验分布形态是否多峰、是否超出物理合理范围等多个指标结合起来形成一个综合的“异常评分”。同时永远保留人工检查的环节最终的判断权应交由领域专家。
