Clarification on Variable Bitrate Compression Implementation

[teaandscones]

Hello,

Thank you for the excellent work!

I’m relatively new to the library and I was wondering if you could provide some guidance on how to adapt the library to perform what is described in the section “Algorithms for variable bitrate compression with a single set of NN weights” of your paper: https://arxiv.org/pdf/2402.18930.

From equations (3) and (4), it seems that scaling the latent tensor before and after the rounding operations might be sufficient.

However, this is not entirely clear to me. Shouldn’t the entropy model also be adjusted accordingly?

Looking at where the rounding operation occurs in the library,

CompressAI/compressai/entropy_models/entropy_models.py

Line 169 in 4fbc02f

outputs = torch.round(outputs)

I noticed that the means are subtracted beforehand.

CompressAI/compressai/entropy_models/entropy_models.py

Line 167 in 4fbc02f

outputs -= means

Therefore, scaling the latent tensor before rounding would also scale the means.

But these means are also passed to the entropy model when computing the bitstreams. For example, they are first retrieved via _get_medians() and then passed to the compress method of the entropy model.

CompressAI/compressai/entropy_models/entropy_models.py

Lines 556 to 562 in 4fbc02f

	def compress(self, x):
	indexes = self._build_indexes(x.size())
	medians = self._get_medians().detach()
	spatial_dims = len(x.size()) - 2
	medians = self._extend_ndims(medians, spatial_dims)
	medians = medians.expand(x.size(0), *([-1] * (spatial_dims + 1)))
	return super().compress(x, indexes, medians)

Could you please clarify whether the entropy model needs to be modified as well, or if scaling alone is sufficient?

Thank you in advance!