Fix RAdam stepSize (#700)

`RAdam` step size was missing a factor of `learningRate / (1 - beta1Power)`.
Update `RAdam` numerical correctness test.

Co-authored-by: t.ae <t-ae@users.noreply.github.com>
Co-authored-by: Dan Zheng <danielzheng@google.com>

Author: 3 people

Sources/TensorFlow/Optimizers/MomentumBased.swift

@@ -592,18 +592,18 @@ where
     let N_sma_inf = 2 / (1 - beta2) - 1
     let N_sma_t = N_sma_inf - 2 * step * beta2Power / (1 - beta2Power)
 
-    if N_sma_t > 5 {
+    if N_sma_t >= 5 {
       // Compute bias-corrected second moments, rectification and adapted momentum.
       let secondMoments_h = Model.TangentVector.sqrt(secondMoments).adding(epsilon)
-      let stepSize = sqrtf(
-        (N_sma_t - 4) * (N_sma_t - 2) * N_sma_inf
-          / ((N_sma_inf - 4) * (N_sma_inf - 2) * (N_sma_t)))
+      let stepSize =
+        sqrtf(
+          (N_sma_t - 4) * (N_sma_t - 2) * N_sma_inf
+            / ((N_sma_inf - 4) * (N_sma_inf - 2) * (N_sma_t))) * learningRate / (1 - beta1Power)
       model.move(
-        along: (firstMoments ./ secondMoments_h).scaled(
-          by: -stepSize * sqrtf(1 - beta2Power)))
+        along: (firstMoments ./ secondMoments_h).scaled(by: -stepSize * sqrtf(1 - beta2Power)))
     } else {
       // Update with un-adapted momentum.
-      let stepSize = self.learningRate * step / (1 - beta1Power)
+      let stepSize = learningRate / (1 - beta1Power)
       model.move(along: firstMoments.scaled(by: -stepSize))
     }
   }

Tests/TensorFlowTests/OptimizerTests.swift

@@ -19,18 +19,42 @@ class OptimizerTests: XCTestCase {
   /// A dense layer for testing optimizer convergence.
   // TODO: Consider replacing users with `Dense`.
   struct Model: Layer {
-    var dense1 = Dense<Float>(weight: [[0.8]], bias: [0.8], activation: identity)
+    var dense = Dense<Float>(weight: [[0.8]], bias: [0.8], activation: identity)
 
     @differentiable
     func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
-      dense1(input)
+      dense(input)
     }
   }
 
+  /// Check expected weight and bias after updating `model` with `optimizer` `stepCount` times.
+  ///
+  /// - Note: optimizer correctness reference implementations exist at
+  ///   `Utilities/ReferenceImplementations/optimizers.py`.
+  func testCorrectness<Opt: Optimizer>(
+    optimizer: Opt,
+    model: Model,
+    expectedWeight: Tensor<Float>,
+    expectedBias: Tensor<Float>,
+    stepCount: Int = 1000,
+    file: StaticString = #file,
+    line: UInt = #line
+  ) where Opt.Model == Model {
+    var optimizer = optimizer
+    var model = model
+    let grad = Model.TangentVector(dense: .init(weight: [[0.1]], bias: [0.2]))
+    for _ in 0..<stepCount {
+      optimizer.update(&model, along: grad)
+    }
+    XCTAssertEqual(model.dense.weight, expectedWeight, file: file, line: line)
+    XCTAssertEqual(model.dense.bias, expectedBias, file: file, line: line)
+  }
+
   /// Check that `model` converges after updating it with `optimizer` `stepCount` times.
   func testConvergence<Opt: Optimizer>(
     optimizer: Opt,
     model: Model,
+    stepCount: Int = 1000,
     file: StaticString = #file,
     line: UInt = #line
   ) where Opt.Model == Model {
@@ -40,7 +64,7 @@ class OptimizerTests: XCTestCase {
       .reshaped(to: [-1, 1])
     let y: Tensor<Float> = x + 1
 
-    for _ in 0..<1000 {
+    for _ in 0..<stepCount {
       let grad = gradient(at: model) { model -> Tensor<Float> in
         let yy = model(x)
         return meanSquaredError(predicted: yy, expected: y)
@@ -102,7 +126,7 @@ class OptimizerTests: XCTestCase {
   func testRAdam() {
     let model = Model()
     let optimizer = RAdam(for: model)
-    testConvergence(optimizer: optimizer, model: model)
+    testConvergence(optimizer: optimizer, model: model, stepCount: 1400)
   }
 
   /// A `Tensor<Float>` wrapper for testing optimizer numerical correctness.
@@ -207,10 +231,9 @@ class OptimizerTests: XCTestCase {
     let optimizer = RAdam(for: values, learningRate: 1e-3, epsilon: 1e-7)
     // FIXME(TF-759): Investigate large differences with Python reference implementation results:
     // `[ 0.46914074, -0.44463935, -0.44513944]`.
-    // Pending fix: https://github.com/tensorflow/swift-apis/pull/700
     testNumericalCorrectness(
       optimizer: optimizer, startingValues: values,
-      expectedValues: [ 443.81192, -443.80478, -443.85016])
+      expectedValues: [ 0.44664007, -0.44463903, -0.45914108])
   }
 
   static var allTests = [