巴黎大学的一名研究人员使用Graphcore IPU来加速面向宇宙学应用的神经网络训练。
在这篇新发表的论文中，研究员Bastien Arcelin探讨了IPU处理器在两个宇宙学深度学习用例的适用性和性能：从经过训练的VAE潜在空间生成星系图像，以及使用确定性深度神经网络和贝叶斯神经网络（BNN）估计星系形状。
大数据
未来即将进行的天文勘测将产生空前数量的观测数据。其中一个项目是维拉·C·鲁宾天文台（Vera C. Rubin Observatory）的时空遗迹巡天（Legacy Survey of Space and Time，LSST），预计每晚生成20TB的数据，在其十年的运营期间总共将生成约60PB的数据。
宇宙学研究人员越来越多地使用神经网络来管理如此庞大而复杂的数据集。
然而，现实世界观测中涉及的不可知或不可控变量的数量非常巨大，使得针对此类数据的神经网络训练变得极具挑战性。
与许多人工智能应用程序一样，模拟数据（精确参数可知可控）最适合训练。
由于这些模拟数据很大也很复杂，因此相关的神经网络需要快速、准确，并且能够针对某些应用程序准确地描绘认知不确定性的特点。
这种计算需求引出了一个问题：专为人工智能设计的硬件能否为宇宙学中的深度学习提供卓越的性能？
Graphcore IPU旨在高效处理计算图（每个人工智能算法背后的基本结构），这使其成为测试该理论的理想示范处理器。
这项研究在一个第一代Graphcore芯片——GC2 IPU，和一个Nvidia V100 GPU之间进行了性能对比。
在用小批量训练神经网络时，发现IPU执行DNN的速度至少是GPU的两倍，执行BNN的速度至少是GPU的4倍。GC2 IPU实现这种性能的功耗仅为GPU的一半。
在所有实验中都使用了TensorFlow 2.1框架，因为TensorFlow 1和2在具有集成XLA后端的IPU上都得到完全支持。
星系图像生成
基于模拟宇宙现象（如星系）的数据传统上基于简单的分析轮廓（Analytic Profiles），如塞西克轮廓（Sérsic Profiles） 。这是一种缓慢的生成技术，增加了形成模型偏差的风险。
深度学习方法有可能加速数据模拟。生成性神经网络越来越多地被用于为宇宙学中的各种应用模拟星系（例如：Lanusse等人，2020年；Regier、McAuliffe和Prabhat，2015年；或Arcelin等人，2021年）。
图像是通过对已在孤立星系图像上训练过的变分自编码器（Variational Autoencoder，VAE）的潜在空间分布进行采样来生成的。
在这种情况下，可以看到IPU在生成小批量图像时性能优于GPU。
在Graphcore第二代IPU——GC200上运行这个工作负载，有可能在生成更大批量时显著提高性能，这是因为处理器的片上存储被极大地扩展了。
星系形状参数估计
下一代天文勘测将比以往任何时候都更深入地观察天空，因此混合（即重叠）物体的可能性也会有所增加。显然，当星系被部分覆盖时，测量星系形状变得困难得多。
现有的星系形状测量方法不能准确地测量这些重叠的物体，因此当勘测工作需要观察到天空的更深处时，就需要新的技术。
为测量孤立星系和混合星系的形状椭率参数，巴黎大学研究员Bastien Arcelin在开发一种具有深度神经网络和卷积层的新技术。
第一个实验使用确定性神经网络进行，第二个测试使用了BNN。一旦经过训练，确定性网络中的可训练参数就有一个固定值，而且即便网络被两次馈送相同的孤立星系图像，这个固定值也不会改变。这与BNN形成了反差。在BNN里，权重本身被赋予概率分布而不是单个值。这意味着用相同的图像向BNN网络馈送两次会导致概率分布的两次不同采样，从而产生两个略有不同的输出。通过对这些分布进行多次采样，可以估计出结果的认知不确定性。
对于确定性神经网络，与GPU相比，Graphcore IPU可实现更快的训练速度，尤其是在使用小批量时。这并不令人惊讶，因为IPU技术以其在计算机视觉应用中的卷积神经网络上的高效性能而闻名。
接下来，该技术在贝叶斯神经网络上进行了测试。对于宇宙学研究人员来说，建立对神经网络预测的置信水平至关重要。研究人员可以使用BNN来计算与预测相关的认知不确定性，从而帮助确定这些预测中的置信水平，因为认知不确定性与大的测量误差密切相关。
在BNN的情况下，IPU的性能再次优于GPU，训练速度至少快四倍。
显然，IPU能够显著减少用于星系形状参数估计的人工神经网络训练时间，并且在本研究中，使用一个IPU时在小批量下表现最佳。对于更大的批尺寸，网络可以拆分到多个IPU上以提高性能和效率。