具有随机估计梯度的 jax 最小化

Question

我正在尝试使用tensorflow_probability.substrates.jax和jax.scipy.optimize.minimize中的 bfgs 优化器来最小化根据伪随机样本估计的 function f ，并以jax.random.PRNGKey作为参数。 要将此 function 与 jax/tfp bfgs 最小化器一起使用，我将 function 包装在 lambda function 中

seed = 100
key  = jax.random.PRNGKey(seed)
fun = lambda x: return f(x,key)
result = jax.scipy.optimize.minimize(fun = fun, ...)

当最小化例程调用要最小化的 function 以便我以可重现的方式使用不同的伪随机数时，更新密钥的最佳方法是什么？ 也许是全局关键变量？ 如果是，有没有我可以效仿的例子？

其次，有没有办法让优化在一定时间后停止，就像 scipy 中的回调一样？ 我可以直接使用 bfgs/l-bfgs-b/ 等的 scipy 实现，并使用 jax ony 估计 function 和 tis 梯度，这似乎有效。 scipy、jax.scipy 和 tfp.jax bfgs 实现之间有区别吗？

最后，有没有办法在 jax.scipy 或 tfp 中的 bfgs 优化期间打印fun的 arguments 的值，假设f是 jitted？

谢谢！

Answer 1

没有办法用jax.scipy.optimize.minimize做你要求的事情，因为最小化器不提供任何方法来跟踪 function 调用之间的变化 state，并且不提供优化器中的任何内置随机性。

如果您对 JAX 中的随机优化感兴趣，您可以尝试JAXOpt 中的随机优化，它提供了一组更加灵活的优化例程。

关于你的第二个问题，如果你想在 jit 编译优化或其他循环过程中打印值，你可以使用jax.debug.print 。