Quanstions about entropy model and function update()

[AlbertoPresta]

Hi,

sorry if I bother you with another question related to the entropy model calculated in the function update().

1. Why the quantized CDF is larger by two with respect to the cmf? is this because the first and the last values are not part of the CDF and are equal to 0 and 2^precision -1
2. When I update a pre-trained network (and also a network that I trained in my own dataset) the value of the tail_mass calculated in the function update() is equal to one, is it right? it seems quite weird to me.
3. the way the samples are calculated (samples = samples[None, :] + pmf_length[:, None, None]) is to move around the most common value of the latent specific dimension?

thanks in advance for the answers,
Alberto

[YodaEmbedding]

1. Two reasons:
   • The cdf is indexed at "bin" boundaries, whereas each element of pmf represents the bins themselves.
     The probability of a symbol x is recovered via pmf[x] = cdf[x + 1] - cdf[x].
     

     CompressAI/compressai/entropy_models/entropy_models.py

     Lines 204 to 212 in a29c3d5

     	def _pmf_to_cdf(self, pmf, tail_mass, pmf_length, max_length):
     	cdf = torch.zeros(
     	(len(pmf_length), max_length + 2), dtype=torch.int32, device=pmf.device
     	)
     	for i, p in enumerate(pmf):
     	prob = torch.cat((p[: pmf_length[i]], tail_mass[i]), dim=0)
     	_cdf = pmf_to_quantized_cdf(prob, self.entropy_coder_precision)
     	cdf[i, : _cdf.size(0)] = _cdf
     	return cdf

     
     

     CompressAI/compressai/cpp_exts/ops/ops.cpp

     Line 71 in a29c3d5

     std::partial_sum(cdf.begin(), cdf.end(), cdf.begin());

   • One value needs to be allocated to represent all the low probability symbols outside the pmf range, where the "bypass" coding engine kicks in to try and encode those out-of-range symbols. The probability of this bypass symbol is roughly equal to the tail_mass (but not exactly equal, due to rounding precision).
     

     CompressAI/compressai/cpp_exts/rans/rans_interface.cpp

     Lines 140 to 147 in a29c3d5

     	assert(value < cdfs_sizes[cdf_idx] - 1);
     	_syms.push_back({static_cast<uint16_t>(cdf[value]),
     	static_cast<uint16_t>(cdf[value + 1] - cdf[value]),
     	false});
     	/* Bypass coding mode (value == max_value -> sentinel flag) */
     	if (value == max_value) {

2. This initializes it to 1e-9:
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Line 346 in a29c3d5

   tail_mass: float = 1e-9,

   
   For the EntropyBottleneck, it is the sum of the masses of the outer left/right regions of the logistic distribution:
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Line 424 in a29c3d5

   tail_mass = torch.sigmoid(lower[:, 0, :1]) + torch.sigmoid(-upper[:, 0, -1:])

   
   For the GaussianConditional, it is the sum of the masses of the outer left/right regions of the normal distribution:
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Line 642 in a29c3d5

   tail_mass = 2 * lower[:, :1]

3. Samples calculation.
   For the EntropyBottleneck, it looks like the additional dimension is needed due to the K-fold "matrix multiplication" for the small entropy modelling network described in the appendix of Balle2018:
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Lines 410 to 423 in a29c3d5

   	max_length = pmf_length.max().item()
   	device = pmf_start.device
   	samples = torch.arange(max_length, device=device)
   	samples = samples[None, :] + pmf_start[:, None, None]
   	half = float(0.5)
   	lower = self._logits_cumulative(samples - half, stop_gradient=True)
   	upper = self._logits_cumulative(samples + half, stop_gradient=True)
   	sign = -torch.sign(lower + upper)
   	pmf = torch.abs(torch.sigmoid(sign * upper) - torch.sigmoid(sign * lower))
   	pmf = pmf[:, 0, :]

   
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Lines 436 to 455 in a29c3d5

   	def _logits_cumulative(self, inputs: Tensor, stop_gradient: bool) -> Tensor:
   	# TorchScript not yet working (nn.Mmodule indexing not supported)
   	logits = inputs
   	for i in range(len(self.filters) + 1):
   	matrix = getattr(self, f"_matrix{i:d}")
   	if stop_gradient:
   	matrix = matrix.detach()
   	logits = torch.matmul(F.softplus(matrix), logits)
   	bias = getattr(self, f"_bias{i:d}")
   	if stop_gradient:
   	bias = bias.detach()
   	logits += bias
   	if i < len(self.filters):
   	factor = getattr(self, f"_factor{i:d}")
   	if stop_gradient:
   	factor = factor.detach()
   	logits += torch.tanh(factor) * torch.tanh(logits)
   	return logits

   
   For the GaussianConditional, no extra dimension is needed since there's no extra matrix-based network needed and one can model the normal distribution directly:
   

   CompressAI/compressai/entropy_models/entropy_models.py

   Lines 632 to 640 in a29c3d5

   	samples = torch.abs(
   	torch.arange(max_length, device=device).int() - pmf_center[:, None]
   	)
   	samples_scale = self.scale_table.unsqueeze(1)
   	samples = samples.float()
   	samples_scale = samples_scale.float()
   	upper = self._standardized_cumulative((0.5 - samples) / samples_scale)
   	lower = self._standardized_cumulative((-0.5 - samples) / samples_scale)
   	pmf = upper - lower

[AlbertoPresta]

first of all thanks for the precious and precise answers.
I got all the points you explained...but in practice, something weird happens ( or at least something that I can't understand fully):
I downloaded a pretrained model with the following command:

net = bmshj2018_factorized(quality=8, pretrained=True).eval().to(torch.device("cpu"))
net.update(force=True)

The I try to print out the quantized cdf and I have something like this (I show only the first dimension):

quantized_cdf = Tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
10, 11, 12, 13, 14, 15, 65536, 0, ...,0])
tail_mass = [1.]

Moreover, all dimensions have tail_mass equal to 1...what does it mean? should it equal to (or similar to) 1e-9? what is the domain of the quantized cdf if I want to plot it?

thanks in advance for the answers.
Alberto

[YodaEmbedding]

It might be that the model checkpoint wasn't loaded correctly. Could you please post a MCVE to help us verify this behavior?

[AlbertoPresta]

Yes, of course, I'll try to be as clear as possible:

from compressai.zoo import bmshj2018_factorized
import torch 

def upd(net):
    """
    this function simply reproduces what happens in the update function. 
    I wrote this function just to compute again the tail mass, the pmf and the quantized cdf
    """"

    medians = net.entropy_bottleneck.quantiles[:, 0, 1]

    minima = medians - net.entropy_bottleneck.quantiles[:, 0, 0]
    minima = torch.ceil(minima).int()
    minima = torch.clamp(minima, min=0)

    maxima = net.entropy_bottleneck.quantiles[:, 0, 2] - medians
    maxima = torch.ceil(maxima).int()
    maxima = torch.clamp(maxima, min=0)

    pmf_start = medians - minima
    pmf_length = maxima + minima + 1
        

    max_length = pmf_length.max().item()
    device = pmf_length.device
    samples = torch.arange(max_length, device=device)

    samples = samples[None, :] + pmf_length[:, None, None]
        
    half = float(0.5)
        
    lower = net.entropy_bottleneck._logits_cumulative(samples - half, stop_gradient=True)
    upper = net.entropy_bottleneck._logits_cumulative(samples + half, stop_gradient=True)
    sign = -torch.sign(lower + upper)
    pmf = torch.abs(torch.sigmoid(sign * upper) - torch.sigmoid(sign * lower))

    pmf = pmf[:, 0, :]
    tail_mass = torch.sigmoid(lower[:, 0, :1]) + torch.sigmoid(-upper[:, 0, -1:])
    quantized_cdf = net.entropy_bottleneck._pmf_to_cdf(pmf, tail_mass, pmf_length, max_length)
    
    return quantized_cdf, tail_mass




ml = torch.device("cpu")
net = bmshj2018_factorized(quality=1, pretrained=True).eval().to(ml)
net.update(force = True)

quantized_cdf, tail_mass = upd(net) # the outputs are like the ones I wrote in the previous answer

[YodaEmbedding]

Note that there are two lines that are different between upd() and the original EntropyBottleneck.update():

-    device = pmf_length.device
+    device = pmf_start.device
     samples = torch.arange(max_length, device=device)
-    samples = samples[None, :] + pmf_length[:, None, None]
+    samples = samples[None, :] + pmf_start[:, None, None]

With these fixes, 0 <= tail_mass <= 6.2472e-10.