LazyAdamOptimizerを見てみた

機械学習 Optimizer TensorFlow

見てみた

概要

Adam Optimizerの sparse更新をより効率化したもの。

既存のAdamアルゴリズムは2つの移動平均の各train variableに対して行います。 計算はすべてのStepで更新されます。 このクラスは sparse variableに対して勾配更新の緩やかな処理を行います。 すべてのインデックスのアキュムレータを更新するのではなく、 現在のバッチに含まれるスパース変数インデックスの移動平均アキュムレータのみを更新します。

オリジナルのAdamオプティマイザと比較すると、一部のアプリケーションではモデルトレーニングのスループットが大幅に向上しています。
ただし、元のAdamアルゴリズムとはわずかに異なるセマンティクスが提供されているため、異なる経験的結果が生じる可能性があります。

詳細

  • lazy_adam_optimizer.py
    では、Adam Optimizerから2つの関数のみoverrideしている。

  • _apple_sparse
    ※こちらを例に以後違いを説明します。

  • _resource_apply_sparse

_apple_sparse

Adam Optimizerでは以下のように関数を呼んでいる。

LazyAdam Optimizerでは*\_shared関数は呼んでいない

  def _apply_sparse(self, grad, var):
    beta1_power, beta2_power = self._get_beta_accumulators()
    beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype)
    beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype)
    lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype)
    beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype)
    beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype)
    epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype)
    lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power))

    # \\(m := beta1 * m + (1 - beta1) * g_t\\)
    m = self.get_slot(var, "m")
    m_t = state_ops.scatter_update(m, grad.indices,
                                   beta1_t * array_ops.gather(m, grad.indices) +
                                   (1 - beta1_t) * grad.values,
                                   use_locking=self._use_locking)

    # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\)
    v = self.get_slot(var, "v")
    v_t = state_ops.scatter_update(v, grad.indices,
                                   beta2_t * array_ops.gather(v, grad.indices) +
                                   (1 - beta2_t) * math_ops.square(grad.values),
                                   use_locking=self._use_locking)

    # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\)
    m_t_slice = array_ops.gather(m_t, grad.indices)
    v_t_slice = array_ops.gather(v_t, grad.indices)
    denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t
    var_update = state_ops.scatter_sub(var, grad.indices,
                                       lr * m_t_slice / denominator_slice,
                                       use_locking=self._use_locking)
    return control_flow_ops.group(var_update, m_t, v_t)

m_t := beta_1 * m_ {t-1} *+ (1 - beta_1) * g

の箇所だけを比較すると

  • Adam Optimizer
    # m_t = beta1 * m + (1 - beta1) * g_t
    m = self.get_slot(var, "m")
    m_scaled_g_values = grad * (1 - beta1_t)
    m_t = state_ops.assign(m, m * beta1_t,
                           use_locking=self._use_locking)
    with ops.control_dependencies([m_t]):
      m_t = scatter_add(m, indices, m_scaled_g_values)
  • LazyAdam Optimizer
    # \\(m := beta1 * m + (1 - beta1) * g_t\\)
    m = self.get_slot(var, "m")
    m_t = state_ops.scatter_update(m, grad.indices,
                                   beta1_t * array_ops.gather(m, grad.indices) +
                                   (1 - beta1_t) * grad.values,
                                   use_locking=self._use_locking)

となっている。

LazyAdamで使っている scatter_update はindicesの部分(sparse)のみ更新する関数

そのあと v_tも同じように計算し

    m_t_slice = array_ops.gather(m_t, grad.indices)
    v_t_slice = array_ops.gather(v_t, grad.indices)

切り取って

    denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t
    var_update = state_ops.scatter_sub(var, grad.indices,
                                       lr * m_t_slice / denominator_slice,
                                       use_locking=self._use_locking)

更新値を算出。

まとめ

Moving-Averageの計算において、部分な値にて計算しているだけみたい。
(一部 addから updateの変更も合っているけど)

参考