Repository: google-deepmind/optax

Issue:
Implement log(cosh()) loss function
log-cosh is a doubly differentiable alternative to the huber loss. A naive implementation is prone to overflow (since cosh has an e^x term), so I think it'd be a useful addition to the library. Plus, it's implemented in other libraries, such as [TensorFlow](https://www.tensorflow.org/api_docs/python/tf/keras/losses/log_cosh).

If this sounds like a relevant addition, I'd be happy to contribute it!


Discussion:
This would be a great contribution! Looking forward to it!

Resolving patch:
diff --git a/optax/_src/loss.py b/optax/_src/loss.py
--- a/optax/_src/loss.py
+++ b/optax/_src/loss.py
@@ -227,3 +227,29 @@ def cosine_distance(
   chex.assert_type([targets, predictions], float)
   # cosine distance = 1 - cosine similarity.
   return 1. - cosine_similarity(predictions, targets, epsilon)
+
+
+def log_cosh(
+    predictions: chex.Array,
+    targets: Optional[chex.Array] = None,
+) -> chex.Array:
+  """Calculates the log-cosh loss for a set of predictions.
+
+  log(cosh(x)) is approximately `(x**2) / 2` for small x and `abs(x) - log(2)`
+  for large x.  It is a twice differentiable alternative to the Huber loss.
+
+  References:
+    [Chen et al, 2019](https://openreview.net/pdf?id=rkglvsC9Ym)
+
+  Args:
+    predictions: a vector of arbitrary shape.
+    targets: a vector of shape compatible with predictions; if not provided
+      then it is assumed to be zero.
+
+  Returns:
+    the log-cosh loss.
+  """
+  chex.assert_type([predictions], float)
+  errors = (predictions - targets) if (targets is not None) else predictions
+  # log(cosh(x)) = log((exp(x) + exp(-x))/2) = log(exp(x) + exp(-x)) - log(2)
+  return jnp.logaddexp(errors, -errors) - jnp.log(2.0).astype(errors.dtype)


Test patch:
diff --git a/optax/_src/loss_test.py b/optax/_src/loss_test.py
--- a/optax/_src/loss_test.py
+++ b/optax/_src/loss_test.py
@@ -204,5 +204,38 @@ def test_batched_similarity(self):
         1. - self.exp, atol=1e-4)
 
 
+class LogCoshTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Test large values for overflow
+    self.ys = jnp.array([500, -2., -1., 0.5, 1.])
+    self.ts = jnp.array([-200, -1.5, 0., -1, 1.])
+    # computed using tensorflow.keras.losses.log_cosh v2.4.1
+    self.exp = jnp.array([699.3068, 0.12011445, 0.4337809, 0.85544014, 0.])
+    self.exp_ys_only = jnp.array(
+        [499.30685, 1.3250027, 0.4337809, 0.12011451, 0.43378082])
+
+  @chex.all_variants
+  def test_scalar(self):
+    out = self.variant(loss.log_cosh)(self.ys[0], self.ts[0])
+    np.testing.assert_allclose(out, self.exp[0], rtol=1e-6, atol=1e-6)
+
+  @chex.all_variants
+  def test_batched(self):
+    out = self.variant(loss.log_cosh)(self.ys, self.ts)
+    np.testing.assert_allclose(out, self.exp, rtol=1e-6, atol=1e-6)
+
+  @chex.all_variants
+  def test_scalar_predictions_only(self):
+    out = self.variant(loss.log_cosh)(self.ys[0])
+    np.testing.assert_allclose(out, self.exp_ys_only[0], rtol=1e-6, atol=1e-6)
+
+  @chex.all_variants
+  def test_batched_predictions_only(self):
+    out = self.variant(loss.log_cosh)(self.ys)
+    np.testing.assert_allclose(out, self.exp_ys_only, rtol=1e-6, atol=1e-6)
+
+
 if __name__ == '__main__':
   absltest.main()