feat: change Lipschitz cst estimation to Jacobian method

To estimate the Lipschitz constant of a model, we now use the Jacobian
method that is more robust than random perturbations (subject to numerical
errors).

The Jacobian method computes Jacobians for a batch of inputs and gets
the max singular values of each Jacobian matrix, corresponding to the local
Lipschitz constant.
The function then returns the greatest Lipschitz constant estimation in the
batch.

Author: cofri

deel/lip/utils.py

@@ -13,64 +13,60 @@
 
 
 def evaluate_lip_const_gen(
-    model: Model,
-    generator: Generator[Tuple[np.ndarray, np.ndarray], Any, None],
-    eps=1e-4,
-    seed=None,
+    model: Model, generator: Generator[Tuple[np.ndarray, np.ndarray], Any, None]
 ):
     """
-    Evaluate the Lipschitz constant of a model, with the naive method.
-    Please note that the estimation of the lipschitz constant is done locally around
-    input sample. This may not correctly estimate the behaviour in the whole domain.
+    Evaluate the Lipschitz constant of a model, using the Jacobian of the model.
+    Please note that the estimation of the Lipschitz constant is done locally around
+    input samples. This may not correctly estimate the behaviour in the whole domain.
     The computation might also be inaccurate in high dimensional space.
 
     This is the generator version of evaluate_lip_const.
 
     Args:
         model: built keras model used to make predictions
         generator: used to select datapoints where to compute the lipschitz constant
-        eps (float): magnitude of noise to add to input in order to compute the constant
-        seed (int): seed used when generating the noise ( can be set to None )
 
     Returns:
         float: the empirically evaluated lipschitz constant.
 
     """
     x, _ = generator.send(None)
-    return evaluate_lip_const(model, x, eps, seed=seed)
+    return evaluate_lip_const(model, x)
 
 
-def evaluate_lip_const(model: Model, x, eps=1e-4, seed=None):
+def evaluate_lip_const(model: Model, x):
     """
-    Evaluate the Lipschitz constant of a model, with the naive method.
+    Evaluate the Lipschitz constant of a model, using the Jacobian of the model.
     Please note that the estimation of the lipschitz constant is done locally around
-    input sample. This may not correctly estimate the behaviour in the whole domain.
+    input samples. This may not correctly estimate the behaviour in the whole domain.
 
     Args:
         model: built keras model used to make predictions
         x: inputs used to compute the lipschitz constant
-        eps (float): magnitude of noise to add to input in order to compute the constant
-        seed (int): seed used when generating the noise ( can be set to None )
 
     Returns:
-        float: the empirically evaluated lipschitz constant. The computation might also
+        float: the empirically evaluated Lipschitz constant. The computation might also
             be inaccurate in high dimensional space.
 
     """
-    y_pred = model.predict(x)
-    # x = np.repeat(x, 100, 0)
-    # y_pred = np.repeat(y_pred, 100, 0)
-    x_var = x + K.random_uniform(
-        shape=x.shape, minval=eps * 0.25, maxval=eps, seed=seed
-    )
-    y_pred_var = model.predict(x_var)
-    dx = x - x_var
-    dfx = y_pred - y_pred_var
-    ndx = K.sqrt(K.sum(K.square(dx), axis=range(1, len(x.shape))))
-    ndfx = K.sqrt(K.sum(K.square(dfx), axis=range(1, len(y_pred.shape))))
-    lip_cst = K.max(ndfx / ndx)
-    print(f"lip cst: {lip_cst:.3f}")
-    return lip_cst
+    batch_size = x.shape[0]
+    x = tf.constant(x, dtype=model.input.dtype)
+
+    # Get the jacobians of the model w.r.t. the inputs
+    with tf.GradientTape() as tape:
+        tape.watch(x)
+        y_pred = model(x, training=False)
+    batch_jacobian = tape.batch_jacobian(y_pred, x)
+
+    # Reshape the jacobians (in case of multi-dimensional input/output like in conv)
+    xdim = tf.reduce_prod(x.shape[1:])
+    ydim = tf.reduce_prod(y_pred.shape[1:])
+    batch_jacobian = tf.reshape(batch_jacobian, (batch_size, ydim, xdim))
+
+    # Compute the spectral norm of the jacobians and return the maximum
+    b = tf.norm(batch_jacobian, ord=2, axis=[-2, -1]).numpy()
+    return tf.reduce_max(b)
 
 
 def _padding_circular(x, circular_paddings):

tests/test_layers.py

@@ -224,7 +224,7 @@ def train_k_lip_model(
         linear_generator(batch_size, input_shape, kernel),
         steps=10,
     )
-    empirical_lip_const = evaluate_lip_const(model=model, x=x, seed=42)
+    empirical_lip_const = evaluate_lip_const(model=model, x=x)
     # save the model
     model_checkpoint_path = os.path.join(logdir, "model.keras")
     model.save(model_checkpoint_path, overwrite=True)
@@ -237,7 +237,7 @@ def train_k_lip_model(
         linear_generator(batch_size, input_shape, kernel),
         steps=10,
     )
-    from_empirical_lip_const = evaluate_lip_const(model=model, x=x, seed=42)
+    from_empirical_lip_const = evaluate_lip_const(model=model, x=x)
     # log metrics
     file_writer = tf.summary.create_file_writer(os.path.join(logdir, "metrics"))
     file_writer.set_as_default()