Ray Tracing抽样器：为所有人提供神经网络的贝叶斯抽样

我们推导了基于在折射率$n(x)$为所需似然$\mathcal{L}(x)$函数的介质中沿射线路径进行的Markov Chain Monte Carlo抽样器。该抽样方法在参数空间中以恒定速度传播射线，相对于Hamiltonian Monte Carlo（HMC）在随机梯度上拥有数倍甚至数十倍的抗热能力，能够跨越任何似然障碍，包括参数空间中的孔洞。使用射线追踪方法，我们对各种不同体系结构的神经网络输出后验分布进行抽样，最高可达15亿参数的GPT-2（生成式预训练变换器2）体系结构，全部在单个消费级GPU上完成。我们还展示了包括传统HMC、微观HMC、Metropolis、Gibbs以及甚至Monte Carlo积分在内的各种prior抽样器，均可在广义射线追踪框架下的任意权重函数上进行抽样。C、JAX和PyTorch的公共代码和文档均可在https://bitbucket.org/pbehroozi/ray-tracing-sampler/src获取。