天文统计学 - 光头博士

Simulation-based Inference、emcee、dynesty、HMC/LMC, JAX

2020-5-14
今天早上还听到普林的天体物理和机器学习双料大牛Miles Cranmer同学讲,他觉得天体物理里面的机器学习应用,既太多,又太少。太多是指太多人在没什么必要用机器学习的地方耍花活。非要在低维数据空间里面用什么深度学习,以便跟人说我在研究中用上深度学习了,这种就没什么意思了。太少指的是,还有很多我们当前认为很困难的问题,甚至是不敢想的问题,其实都是机器学习模型大有可为的地方。

比如就是上面提到的这类“仿真器”的应用。目前依然是宇宙学家走在了前面,少数做星系和恒星光谱的人也开始留意这类应用。狭义的“仿真器” (Emulator)在我看来更像是一个磕了兴奋剂的差值方法(Interpolator). 我有一组覆盖了非常有限参数空间的观测或者模拟,你能不能学习出一个Emulator,把剩下的参数空间给我“差值”出来?在二维空间好像可以用LOESS算法这样的东西去带入理解,但是到了高维空间,就要加入很多很多我还没搞懂的技术了,比如高斯过程,比如神经网络等等。更广义地看,这类“仿真器”应该是“Simulation-based inference”领域的一部分。

2022-04-8】 HMC/LMC, emcee, dynesty, JAX
中午Lunch Talk,听到了一个挺好玩的东西。是关于现代Bayesian统计推断中的一种Monte Carlo采样算法,叫Langevin Monte Carlo (郎之万蒙特卡洛;LMC)。和汉密尔顿蒙特卡洛 (Hamiltonian Monte Carlo;HMC) 类似,都是把理论物理中的一些非常基础的概念抽象出来,应用到统计采样中。HMC自然是来自理论力学里最基础的汉密尔顿动力学;LMC其实(大概)可以看成是HMC的一种改进或者说“变种”,借鉴了统计物理里面描述布朗运动(如分子动力学)的郎之万动力学,或者说郎之万扩散 (Langevin Diffusion) 概念。据说这类算法可以更好地解决非常高维空间以及参数空间相关性很高的情况下的采样问题,可以和深度学习之类的机器学习建立联系。

我是今天才听说的,远远还没有搞懂。但是不得不承认,统计推断和统计学习这个领域的进展真的很快,处于应用下游的天体物理经常是还没开始就落后了。其实天体物理学家开始在宇宙学之外大规模地使用Bayesian学习方法就是近10年的事情。其中一个叫emcee的Ensemble MCMC算法和软件的开发起了极大的推动作用。到现在为止,在使用了MCMC的天文文章里,emcee依然占绝对主导。但其实其中的很多应用都或多或少有问题,有的是把在高维(>10)完全搞不定的emcee用到了几十个参数的框架里,有的是从来不考虑MC的收敛,等等。

宇宙学家一直对另一种叫做嵌套采样(Nested Sampling)的算法情有独钟。这个思路的现代算法MultiNest应该说就是宇宙学家开发的。它和只考虑Posterior的MCMC不同,而是积分参数空间,尝试去计算Bayesian evidence。其中一个原因是这个方法的高维表现更好,再有一个是很多人一直迷信Bayesian evidence可以用来比较模型 (不是很好,因为对Pirors的依赖太大)。。。再有就是早期的宇宙学应用真的不在乎算力,扔到计算平台上跑1个月也没事。我在日本的时候,看着身边的小天才沈佳士(Josh Speagle)同学开发了Dynesty (Dynamic Nested Sampling;面向Posterior的Nested sampling,和MCMC比依然有优势),就误以为这已经是很领先的算法了。。。谁知道统计界早就把目光投向了依赖于Likelihood梯度的HMC类算法。

其实,Josh很早就给我安利过HMC。但是天体物理的很多数据应用都没有办法很容易的得到这个宝贵的“梯度”。比如给我一个描述星系面亮度分布轮廓的函数,我能知道这个函数对各个参数的梯度。。。但实际数据上需要卷积上一个PSF,一下子就找不到解析的梯度了。。。除非你用N个高斯去近似这个函数。等等。不过这几年随着统计采样算法进步的,还有基于机器学习爆发推广起来的"仿真器" (Emulator;可以基于简单的PCA+高斯过程;也可以基于某种神经网络架构)和日益强大的Autodiff算法让一切都变得有可能了。。。我已经被安利了好久的Jax大法了。

仅在观测宇宙学领域内,就至少有三个活跃的团队在系统性地把“可微分”(Differentiable) 算法引用到各种应用里,而且大部分都是基于Jax的。感兴趣地可以到 Github 上搜一下 The CosmoStat Laboratory 的 Jax-lensing 工具;还有一个叫“可微分宇宙企划” (The Differentiable Universe Initiative) 这个开放开发团队的一系列工作;在美国的Argonne国家实验室,我的合作者Andrew Hearin正在以惊人的速度用Jax“改写”肉眼可见的各种工具 (diffstar; dsps; diffmah…总之万物皆可diff;其中有些应用甚至都不是为了最终可以用高级的统计推断工具,单纯地就是借助GPU超算和Jax这个“黑魔法”来提升性能)。

2022-02-06Transformers Can Do Bayesian Inference
利用Transformers深度网络,把估计Posterior分布变成一个监督学习问题。底层上和Simulation-Based Inference (SBI) 神似,都是基于能够生成仿真数据的能力。具体上似乎是:我有一个需要进行推断的模型。基于这个模型我有一个正向建模的任务,能够通过一组参数y,生成出一组仿真数据x。然后对于y我是有先验分布的考虑了。那么我就可以按照这个Prior去通过这个建模任务不断生成数据。然后要做的事情就是引入一个牛逼的深度学习模型,通过数据去学习如何预测参数y。训练完成后,引入我实际需要进行推断的数据,这个时候模型能够给出一个与这组数据对应的参数y的近似Posterior分布。