是否可以将 Tensorflow Graphics 的 Levenberg-Marquardt 优化器与 Tensorflow 2.0 模型集成？

Question

我有一个 Tensorflow 2.0 tf.keras.Sequential 模型。现在，我的技术规范规定使用 Levenberg-Marquardt 优化器来拟合模型。 Tensorflow 2.0 没有提供它作为开箱即用的优化器，但它在 Tensorflow Graphics 模块中可用。

tfg.math.optimizer.levenberg_marquardt.minimize 函数接受 residuals （残差是 Python callable return 张量）和 variables （列表张量对应于我的模型权重）作为参数。

将我的模型转换为残差和变量的最佳方法是什么？

如果我正确理解 minimize 函数的工作原理，我必须提供两个残差。第一个残差必须为每个学习案例调用我的模型并将所有结果聚合到一个张量中。第二个残差必须 return 所有标签作为一个常数张量。问题是 tf.keras.Sequential.predict 函数 return 是一个 numpy 数组而不是张量。我相信如果我将它转换为张量，最小化器将无法计算关于变量的雅可比矩阵。

同样的问题是变量。似乎没有办法将模型中的所有权重提取到张量列表中。

Answer 1

从 implementation/API 的角度来看，tfg.math.optimizer.levenberg_marquardt.minimize 和 Keras 优化器之间存在重大差异。

Keras 优化器，例如 tf.keras.optimizers.Adam 使用梯度作为输入并更新 tf.Variables。

相比之下，tfg.math.optimizer.levenberg_marquardt.minimize 本质上是在图形模式下展开优化循环（使用 tf.while_loop 构造）。它采用初始参数值并生成更新的参数值，这与 Adam & co 不同，Adam & co 仅应用一次迭代并实际通过 assign_add.

更改 tf.Variables 的值

稍微回到理论上的大局，Levenberg-Marquardt 不是适用于任何非线性优化问题（例如 Adam）的一般梯度下降类求解器。它专门针对 非线性最小二乘 优化，因此它不是像 Adam 这样的优化器的直接替代品。在梯度下降中，我们计算损失相对于参数的梯度。在 Levenberg-Marquardt 中，我们计算关于参数的残差的 Jacobian。具体来说，它使用 tf.linalg.lstsq（内部对从 Jacobian 计算的 Gram 矩阵使用 Cholesky 分解）重复求解 delta_params 的线性化问题 Jacobian @ delta_params = residuals，并将 delta_params 应用为更新.

请注意，此 lstsq 操作在参数数量方面具有三次复杂性，因此在神经网络的情况下，它只能应用于相当小的网络。

另请注意，Levenberg-Marquardt 通常用作 批处理算法 ，而不是像 SGD 这样的小批量算法，尽管没有什么可以阻止您在不同的小批处理中应用 LM 迭代每次迭代。

我认为您可能只能通过类似

的方式从 tfg 的 LM 算法中获得一次迭代

from tensorflow_graphics.math.optimizer.levenberg_marquardt import minimize as lm_minimize

for input_batch, target_batch in dataset:

    def residual_fn(trainable_params):
        # do not use trainable params, it will still be at its initial value, since we only do one iteration of Levenberg Marquardt each time.
        return model(input_batch) - target_batch

    new_objective_value, new_params = lm_minimize(residual_fn, model.trainable_variables, max_iter=1)
    for var, new_param in zip(model.trainable_variables, new_params):
        var.assign(new_param)

相比之下，我认为以下朴素的方法将不起作用，我们在计算残差之前分配模型参数：

from tensorflow_graphics.math.optimizer.levenberg_marquardt import minimize as lm_minimize

dataset_iterator = ...

def residual_fn(params):
    input_batch, target_batch = next(dataset_iterator)
    for var, param in zip(model.trainable_variables, params):
        var.assign(param)
    return model(input_batch) - target_batch

final_objective, final_params = lm_minimize(residual_fn, model.trainable_variables, max_iter=10000)
for var, final_param in zip(model.trainable_variables, final_params):
    var.assign(final_param)

主要的概念性问题是 residual_fn 的输出没有关于其输入 params 的梯度，因为这种依赖性经历了 tf.assign。但它甚至可能在此之前失败，因为使用了图形模式中不允许的构造。

总的来说，我认为最好编写自己的适用于 tf.Variables 的 LM 优化器，因为 tfg.math.optimizer.levenberg_marquardt.minimize 有一个非常不同的 API，它并不真正适合优化 Keras 模型参数因为如果没有 tf.assign.

就无法直接计算 model(input, parameters) - target_value

是否可以将 Tensorflow Graphics 的 Levenberg-Marquardt 优化器与 Tensorflow 2.0 模型集成？

Is it possible to integrate Levenberg-Marquardt optimizer from Tensorflow Graphics with a Tensorflow 2.0 model?

tf.keras

tensorflow2.0