Merge pull request #780 from Trusted-AI/development_issue_776

Increase precision for exponentiation in Wasserstein

Author: beat-buesser

art/attacks/evasion/wasserstein.py

@@ -40,6 +40,8 @@
 
 logger = logging.getLogger(__name__)
 
+EPS_LOG = 10 ** -10
+
 
 class Wasserstein(EvasionAttack):
     """
@@ -377,6 +379,7 @@ def _conjugate_sinkhorn(self, x: np.ndarray, grad: np.ndarray, cost_matrix: np.n
         exp_alpha = np.exp(-alpha)
 
         beta = -self.regularization * grad
+        beta = beta.astype(np.float64)
         exp_beta = np.exp(-beta)
 
         # Check for overflow
@@ -399,6 +402,7 @@ def _conjugate_sinkhorn(self, x: np.ndarray, grad: np.ndarray, cost_matrix: np.n
 
         for _ in range(self.conjugate_sinkhorn_max_iter):
             # Block coordinate descent iterates
+            x[x == 0.0] = EPS_LOG  # Prevent divide by zero in np.log
             alpha[I_nonzero_] = (np.log(self._local_transport(K, exp_beta, self.kernel_size)) - np.log(x))[I_nonzero_]
             exp_alpha = np.exp(-alpha)
 
@@ -474,6 +478,7 @@ def _projected_sinkhorn(
 
         for _ in range(self.projected_sinkhorn_max_iter):
             # Block coordinate descent iterates
+            x_init[x_init == 0.0] = EPS_LOG  # Prevent divide by zero in np.log
             alpha = np.log(self._local_transport(K, exp_beta, self.kernel_size)) - np.log(x_init)
             exp_alpha = np.exp(-alpha)
 
@@ -733,8 +738,10 @@ def _check_params(self) -> None:
         if self.eps_step <= 0:
             raise ValueError("The perturbation step-size `eps_step` has to be positive.")
 
-        if self.eps_step > self.eps:
-            raise ValueError("The iteration step `eps_step` has to be smaller than the total attack `eps`.")
+        if self.norm == "inf" and self.eps_step > self.eps:
+            raise ValueError(
+                "The iteration step `eps_step` has to be smaller than or equal to the total attack budget `eps`."
+            )
 
         if self.eps_iter <= 0:
             raise ValueError("The number of epsilon iterations `eps_iter` has to be a positive integer.")