Removes likelihood_bound argument.

With the current set of density models, it seems the numerical fix isn't
necessary any more. Moreover, this can cause gradients to stagnate in low
probability regions.

PiperOrigin-RevId: 340650635
Change-Id: Ib48b50e9a88526d67574705fa2e9e3a035a3ebb5

Author: Johannes Ballé

tensorflow_compression/python/entropy_models/BUILD

@@ -32,7 +32,6 @@ py_library(
     deps = [
         ":continuous_base",
         "//tensorflow_compression/python/distributions:helpers",
-        "//tensorflow_compression/python/ops:math_ops",
         "//tensorflow_compression/python/ops:range_coding_ops",
     ],
 )

tensorflow_compression/python/entropy_models/continuous_base.py

@@ -38,8 +38,7 @@ class ContinuousEntropyModelBase(tf.Module, metaclass=abc.ABCMeta):
 
   @abc.abstractmethod
   def __init__(self, prior, coding_rank, compression=False,
-               likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12, no_variables=False):
+               tail_mass=2**-8, range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -55,8 +54,6 @@ def __init__(self, prior, coding_rank, compression=False,
       compression: Boolean. If set to `True`, the range coding tables used by
         `compress()` and `decompress()` will be built on instantiation. If set
         to `False`, these two methods will not be accessible.
-      likelihood_bound: Float. Lower bound for likelihood values, to prevent
-        training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
@@ -76,7 +73,6 @@ def __init__(self, prior, coding_rank, compression=False,
       self._prior_shape = tuple(int(s) for s in prior.batch_shape)
       self._coding_rank = int(coding_rank)
       self._compression = bool(compression)
-      self._likelihood_bound = float(likelihood_bound)
       self._tail_mass = float(tail_mass)
       self._range_coder_precision = int(range_coder_precision)
       self._no_variables = bool(no_variables)
@@ -143,11 +139,6 @@ def compression(self):
     """Whether this entropy model is prepared for compression."""
     return self._compression
 
-  @property
-  def likelihood_bound(self):
-    """Lower bound for likelihood values."""
-    return self._likelihood_bound
-
   @property
   def tail_mass(self):
     """Approximate probability mass which is range encoded with overflow."""
@@ -285,7 +276,6 @@ def get_config(self):
         dtype=self._dtype.name,
         prior_shape=self._prior_shape,
         coding_rank=self._coding_rank,
-        likelihood_bound=self._likelihood_bound,
         tail_mass=self._tail_mass,
         range_coder_precision=self._range_coder_precision,
         cdf_width=self._cdf.shape.as_list()[1],
@@ -314,7 +304,6 @@ def from_config(cls, config):
       self._prior_shape = tuple(int(s) for s in config["prior_shape"])
       self._coding_rank = int(config["coding_rank"])
       self._compression = True
-      self._likelihood_bound = float(config["likelihood_bound"])
       self._tail_mass = float(config["tail_mass"])
       self._range_coder_precision = int(config["range_coder_precision"])
       self._no_variables = False

tensorflow_compression/python/entropy_models/continuous_batched.py

@@ -20,7 +20,6 @@
 
 from tensorflow_compression.python.distributions import helpers
 from tensorflow_compression.python.entropy_models import continuous_base
-from tensorflow_compression.python.ops import math_ops
 from tensorflow_compression.python.ops import range_coding_ops
 
 
@@ -74,8 +73,7 @@ class ContinuousBatchedEntropyModel(continuous_base.ContinuousEntropyModelBase):
   """
 
   def __init__(self, prior, coding_rank, compression=False,
-               likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12, no_variables=False):
+               tail_mass=2**-8, range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -93,8 +91,6 @@ def __init__(self, prior, coding_rank, compression=False,
       compression: Boolean. If set to `True`, the range coding tables used by
         `compress()` and `decompress()` will be built on instantiation. If set
         to `False`, these two methods will not be accessible.
-      likelihood_bound: Float. Lower bound for likelihood values, to prevent
-        training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
@@ -112,7 +108,6 @@ def __init__(self, prior, coding_rank, compression=False,
         prior=prior,
         coding_rank=coding_rank,
         compression=compression,
-        likelihood_bound=likelihood_bound,
         tail_mass=tail_mass,
         range_coder_precision=range_coder_precision,
         no_variables=no_variables,
@@ -173,11 +168,10 @@ def bits(self, bottleneck, training=True):
           tf.shape(bottleneck), minval=-.5, maxval=.5, dtype=bottleneck.dtype)
     else:
       quantized = self.quantize(bottleneck)
-    probs = self.prior.prob(quantized)
-    probs = math_ops.lower_bound(probs, self.likelihood_bound)
+    log_probs = self.prior.log_prob(quantized)
     axes = tuple(range(-self.coding_rank, 0))
-    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / (
-        -tf.math.log(tf.constant(2., dtype=probs.dtype)))
+    bits = tf.reduce_sum(log_probs, axis=axes) / (
+        -tf.math.log(tf.constant(2, dtype=log_probs.dtype)))
     return bits
 
   @tf.Module.with_name_scope

tensorflow_compression/python/entropy_models/continuous_batched_test.py

@@ -28,7 +28,6 @@ def test_can_instantiate(self):
     em = ContinuousBatchedEntropyModel(noisy, 1)
     self.assertIs(em.prior, noisy)
     self.assertEqual(em.coding_rank, 1)
-    self.assertEqual(em.likelihood_bound, 1e-9)
     self.assertEqual(em.tail_mass, 2**-8)
     self.assertEqual(em.range_coder_precision, 12)
     self.assertEqual(em.dtype, noisy.dtype)

tensorflow_compression/python/entropy_models/continuous_indexed.py

@@ -128,8 +128,7 @@ class ContinuousIndexedEntropyModel(continuous_base.ContinuousEntropyModelBase):
 
   def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
                compression=False, channel_axis=-1, dtype=tf.float32,
-               likelihood_bound=1e-9, tail_mass=2**-8,
-               range_coder_precision=12, no_variables=False):
+               tail_mass=2**-8, range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -165,8 +164,6 @@ def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
         dimension.
       dtype: `tf.dtypes.DType`. The data type of all floating-point
         computations carried out in this class.
-      likelihood_bound: Float. Lower bound for likelihood values, to prevent
-        training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
@@ -209,7 +206,6 @@ def __init__(self, prior_fn, index_ranges, parameter_fns, coding_rank,
         prior=prior,
         coding_rank=coding_rank,
         compression=compression,
-        likelihood_bound=likelihood_bound,
         tail_mass=tail_mass,
         range_coder_precision=range_coder_precision,
         no_variables=no_variables,
@@ -285,11 +281,10 @@ def bits(self, bottleneck, indexes, training=True):
     else:
       offset = helpers.quantization_offset(prior)
       quantized = self._quantize(bottleneck, offset)
-    probs = prior.prob(quantized)
-    probs = math_ops.lower_bound(probs, self.likelihood_bound)
+    log_probs = prior.log_prob(quantized)
     axes = tuple(range(-self.coding_rank, 0))
-    bits = tf.reduce_sum(tf.math.log(probs), axis=axes) / (
-        -tf.math.log(tf.constant(2., dtype=probs.dtype)))
+    bits = tf.reduce_sum(log_probs, axis=axes) / (
+        -tf.math.log(tf.constant(2, dtype=log_probs.dtype)))
     return bits
 
   @tf.Module.with_name_scope
@@ -439,8 +434,8 @@ class LocationScaleIndexedEntropyModel(ContinuousIndexedEntropyModel):
   """
 
   def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
-               compression=False, dtype=tf.float32, likelihood_bound=1e-9,
-               tail_mass=2**-8, range_coder_precision=12, no_variables=False):
+               compression=False, dtype=tf.float32, tail_mass=2**-8,
+               range_coder_precision=12, no_variables=False):
     """Initializer.
 
     Arguments:
@@ -466,8 +461,6 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
         be accessible.
       dtype: `tf.dtypes.DType`. The data type of all floating-point
         computations carried out in this class.
-      likelihood_bound: Float. Lower bound for likelihood values, to prevent
-        training instabilities.
       tail_mass: Float. Approximate probability mass which is range encoded with
         less precision, by using a Golomb-like code.
       range_coder_precision: Integer. Precision passed to the range coding op.
@@ -485,7 +478,6 @@ def __init__(self, prior_fn, num_scales, scale_fn, coding_rank,
         coding_rank=coding_rank,
         compression=compression,
         dtype=dtype,
-        likelihood_bound=likelihood_bound,
         tail_mass=tail_mass,
         range_coder_precision=range_coder_precision,
         no_variables=no_variables,

tensorflow_compression/python/entropy_models/continuous_indexed_test.py

@@ -31,7 +31,6 @@ def test_can_instantiate_one_dimensional(self):
         dict(loc=lambda _: 0, scale=lambda i: tf.exp(i / 8 - 5)), 1)
     self.assertIsInstance(em.prior, uniform_noise.NoisyNormal)
     self.assertEqual(em.coding_rank, 1)
-    self.assertEqual(em.likelihood_bound, 1e-9)
     self.assertEqual(em.tail_mass, 2**-8)
     self.assertEqual(em.range_coder_precision, 12)
     self.assertEqual(em.dtype, tf.float32)
@@ -50,7 +49,6 @@ def test_can_instantiate_n_dimensional(self):
     self.assertIsInstance(em.prior, uniform_noise.NoisyLogisticMixture)
     self.assertEqual(em.coding_rank, 1)
     self.assertEqual(em.channel_axis, -1)
-    self.assertEqual(em.likelihood_bound, 1e-9)
     self.assertEqual(em.tail_mass, 2**-8)
     self.assertEqual(em.range_coder_precision, 12)
     self.assertEqual(em.dtype, tf.float32)
@@ -63,7 +61,6 @@ def test_can_instantiate(self):
         uniform_noise.NoisyNormal, 64, lambda i: tf.exp(i / 8 - 5), 1)
     self.assertIsInstance(em.prior, uniform_noise.NoisyNormal)
     self.assertEqual(em.coding_rank, 1)
-    self.assertEqual(em.likelihood_bound, 1e-9)
     self.assertEqual(em.tail_mass, 2**-8)
     self.assertEqual(em.range_coder_precision, 12)
     self.assertEqual(em.dtype, tf.float32)