Improve with providing Nuclear Norm Constraint (#1107)

* Improve with providing Nuclear Norm Constraint

* Update based on feedback

* Update Based on Feedback

* Solve Conflict

* Solve Conflict

* change format and fix conflict

Co-authored-by: gabrieldemarmiesse <[email protected]>

Author: pkan2

tensorflow_addons/optimizers/conditional_gradient.py

@@ -27,7 +27,8 @@ class ConditionalGradient(tf.keras.optimizers.Optimizer):
 
     This optimizer helps handle constraints well.
 
-    Currently only supports frobenius norm constraint.
+    Currently only supports frobenius norm constraint or nuclear norm
+    constraint.
     See https://arxiv.org/pdf/1803.06453.pdf
 
     ```
@@ -42,6 +43,13 @@ class ConditionalGradient(tf.keras.optimizers.Optimizer):
     gradient is 0.
 
     In this implementation, `epsilon` defaults to $10^{-7}$.
+
+    For nucler norm constraint, the formula is as following:
+
+    ```
+    variable -= (1-learning_rate) * (variable
+        + lambda_ * top_singular_vector(gradient))
+    ```
     """
 
     @typechecked
@@ -50,6 +58,7 @@ def __init__(
         learning_rate: Union[FloatTensorLike, Callable],
         lambda_: Union[FloatTensorLike, Callable] = 0.01,
         epsilon: FloatTensorLike = 1e-7,
+        ord: str = "fro",
         use_locking: bool = False,
         name: str = "ConditionalGradient",
         **kwargs
@@ -64,6 +73,8 @@ def __init__(
             epsilon: A `Tensor` or a floating point value. A small constant
                 for numerical stability when handling the case of norm of
                 gradient to be zero.
+            ord: Order of the norm. Supported values are `'fro'`
+                and `'nuclear'`. Default is `'fro'`, which is frobenius norm.
             use_locking: If `True`, use locks for update operations.
             name: Optional name prefix for the operations created when
                 applying gradients. Defaults to 'ConditionalGradient'.
@@ -78,13 +89,21 @@ def __init__(
         self._set_hyper("learning_rate", kwargs.get("lr", learning_rate))
         self._set_hyper("lambda_", lambda_)
         self.epsilon = epsilon or tf.keras.backend.epsilon()
+        supported_norms = ["fro", "nuclear"]
+        if ord not in supported_norms:
+            raise ValueError(
+                "'ord' must be a supported matrix norm in %s, got '%s' instead"
+                % (supported_norms, ord)
+            )
+        self.ord = ord
         self._set_hyper("use_locking", use_locking)
 
     def get_config(self):
         config = {
             "learning_rate": self._serialize_hyperparameter("learning_rate"),
             "lambda_": self._serialize_hyperparameter("lambda_"),
             "epsilon": self.epsilon,
+            "ord": self.ord,
             "use_locking": self._serialize_hyperparameter("use_locking"),
         }
         base_config = super().get_config()
@@ -106,23 +125,64 @@ def _prepare_local(self, var_device, var_dtype, apply_state):
             self.epsilon, var_dtype
         )
 
-    def _resource_apply_dense(self, grad, var, apply_state=None):
-        def frobenius_norm(m):
-            return tf.math.reduce_sum(m ** 2) ** 0.5
+    @staticmethod
+    def _frobenius_norm(m):
+        return tf.reduce_sum(m ** 2) ** 0.5
+
+    @staticmethod
+    def _top_singular_vector(m):
+        # handle the case where m is a tensor of rank 0 or rank 1.
+        # Example:
+        #   scalar (rank 0) a, shape []=> [[a]], shape [1,1]
+        #   vector (rank 1) [a,b], shape [2] => [[a,b]], shape [1,2]
+        original_rank = tf.rank(m)
+        shape = tf.shape(m)
+        first_pad = tf.cast(tf.less(original_rank, 2), dtype=tf.int32)
+        second_pad = tf.cast(tf.equal(original_rank, 0), dtype=tf.int32)
+        new_shape = tf.concat(
+            [
+                tf.ones(shape=first_pad, dtype=tf.int32),
+                tf.ones(shape=second_pad, dtype=tf.int32),
+                shape,
+            ],
+            axis=0,
+        )
+        n = tf.reshape(m, new_shape)
+        st, ut, vt = tf.linalg.svd(n, full_matrices=False)
+        n_size = tf.shape(n)
+        ut = tf.reshape(ut[:, 0], [n_size[0], 1])
+        vt = tf.reshape(vt[:, 0], [n_size[1], 1])
+        st = tf.matmul(ut, tf.transpose(vt))
+        # when we return the top singular vector, we have to remove the
+        # dimension we have added on
+        st_shape = tf.shape(st)
+        begin = tf.cast(tf.less(original_rank, 2), dtype=tf.int32)
+        end = 2 - tf.cast(tf.equal(original_rank, 0), dtype=tf.int32)
+        new_shape = st_shape[begin:end]
+        return tf.reshape(st, new_shape)
 
+    def _resource_apply_dense(self, grad, var, apply_state=None):
         var_device, var_dtype = var.device, var.dtype.base_dtype
         coefficients = (apply_state or {}).get(
             (var_device, var_dtype)
         ) or self._fallback_apply_state(var_device, var_dtype)
-        norm = tf.convert_to_tensor(
-            frobenius_norm(grad), name="norm", dtype=var.dtype.base_dtype
-        )
         lr = coefficients["learning_rate"]
         lambda_ = coefficients["lambda_"]
         epsilon = coefficients["epsilon"]
-        var_update_tensor = tf.math.multiply(var, lr) - (1 - lr) * lambda_ * grad / (
-            norm + epsilon
-        )
+        if self.ord == "fro":
+            norm = tf.convert_to_tensor(
+                self._frobenius_norm(grad), name="norm", dtype=var.dtype.base_dtype
+            )
+            s = grad / (norm + epsilon)
+        else:
+            top_singular_vector = tf.convert_to_tensor(
+                self._top_singular_vector(grad),
+                name="top_singular_vector",
+                dtype=var.dtype.base_dtype,
+            )
+            s = top_singular_vector
+
+        var_update_tensor = tf.math.multiply(var, lr) - (1 - lr) * lambda_ * s
         var_update_kwargs = {
             "resource": var.handle,
             "value": var_update_tensor,
@@ -131,23 +191,28 @@ def frobenius_norm(m):
         return tf.group(var_update_op)
 
     def _resource_apply_sparse(self, grad, var, indices, apply_state=None):
-        def frobenius_norm(m):
-            return tf.reduce_sum(m ** 2) ** 0.5
-
         var_device, var_dtype = var.device, var.dtype.base_dtype
         coefficients = (apply_state or {}).get(
             (var_device, var_dtype)
         ) or self._fallback_apply_state(var_device, var_dtype)
-        norm = tf.convert_to_tensor(
-            frobenius_norm(grad), name="norm", dtype=var.dtype.base_dtype
-        )
         lr = coefficients["learning_rate"]
         lambda_ = coefficients["lambda_"]
         epsilon = coefficients["epsilon"]
         var_slice = tf.gather(var, indices)
-        var_update_value = tf.math.multiply(var_slice, lr) - (
-            1 - lr
-        ) * lambda_ * grad / (norm + epsilon)
+        if self.ord == "fro":
+            norm = tf.convert_to_tensor(
+                self._frobenius_norm(grad), name="norm", dtype=var.dtype.base_dtype
+            )
+            s = grad / (norm + epsilon)
+        else:
+            top_singular_vector = tf.convert_to_tensor(
+                self._top_singular_vector(grad),
+                name="top_singular_vector",
+                dtype=var.dtype.base_dtype,
+            )
+            s = top_singular_vector
+
+        var_update_value = tf.math.multiply(var_slice, lr) - (1 - lr) * lambda_ * s
         var_update_kwargs = {
             "resource": var.handle,
             "indices": indices,