About Auxiliary loss optimization

[cucNaifuXue]

Hi,

I am confused about the 'Auxiliary loss'. It should be minimized but I don't know it's influence on the whole model's R_D performance.

Should I adjust lr for aux optimizer according to some metrics when training my own model?

In tutorial "Custom model", the example looks like the Factorized Prior model, bmshj2018-factorized. The Rate and Distortion are included in main loss, and I don't know how aux loss works.

Thanks for answers.

[YodaEmbedding]

aux_loss adjusts only the EntropyBottleneck.quantiles to ensure that the trained distribution fits a few target boundary conditions at runtime. Nothing further. It also does not take any image data as input. It has no effect on the training of the model since EntropyBottleneck.quantiles is not used for the model's forward pass during training.

CompressAI/compressai/optimizers/net_aux.py

Lines 46 to 57 in 216af74

	parameters = {
	"net": {
	name
	for name, param in net.named_parameters()
	if param.requires_grad and not name.endswith(".quantiles")
	},
	"aux": {
	name
	for name, param in net.named_parameters()
	if param.requires_grad and name.endswith(".quantiles")
	},
	}

CompressAI/compressai/entropy_models/entropy_models.py

Lines 431 to 434 in 216af74

	def loss(self) -> Tensor:
	logits = self._logits_cumulative(self.quantiles, stop_gradient=True)
	loss = torch.abs(logits - self.target).sum()
	return loss

CompressAI/compressai/entropy_models/entropy_models.py

Lines 378 to 383 in 216af74

	self.quantiles = nn.Parameter(torch.Tensor(channels, 1, 3))
	init = torch.Tensor([-self.init_scale, 0, self.init_scale])
	self.quantiles.data = init.repeat(self.quantiles.size(0), 1, 1)
	target = np.log(2 / self.tail_mass - 1)
	self.register_buffer("target", torch.Tensor([-target, 0, target]))

Thus, I would say that minimizing aux_loss as much as possible isn't critical since it has nearly no effect on the trained RD performance. All it does is make sure the support of the encoding distribution is finite, ensuring that we don't use too many symbols at runtime, and that symbols have some probability over a precision threshold. Potentially, it may also act as a small regularizer that keeps the distributions in check. The .quantiles are not used at all during training, so aux_loss is not relevant until after training finishes.

---

Related discussions:

• Understanding probability distribution for entropy coding #139 (comment)
• Quanstions about entropy model and function update() #146

[Freed-Wu]

it has nearly no effect on the trained RD performance.

It means user shouldn't adjust lr for aux optimizer according to some metrics when training model, really?

[YodaEmbedding]

The only time the .quantiles parameters are referenced during training is in _get_medians, which references the .quantiles[:, :, 1] (midpoints):

CompressAI/compressai/entropy_models/entropy_models.py

Lines 385 to 387 in 23e9c70

	def _get_medians(self) -> Tensor:
	medians = self.quantiles[:, :, 1:2]
	return medians

...which is only used to offset prior to adding noise:

CompressAI/compressai/entropy_models/entropy_models.py

Lines 496 to 498 in 23e9c70

	outputs = self.quantize(
	values, "noise" if training else "dequantize", self._get_medians()
	)

The model should not be too heavily influenced by a near-constant "midpoint" bias.

---

Funnily enough:

• One can probably minimize the aux loss at any point they desire. You can possibly even do it just before each training epoch (for the medians), and then once after training ends (for the tails). We do it continuously during training, under assumption that the "midpoints" remain reasonably stable, though I wonder how true this assumption is, particularly in the first 1-3 epochs... And I wonder if it would be a better usage of the GPU to reduce the amount of aux_optimizer.step()s.
• It is not necessary to formulate the minimization of the aux loss as a SGD problem. What it does is so simple that a mere binary search of the CDF could probably do the trick, and much, much quicker... All it needs to do is solve the following three equations for the learned CDF function c(y):
  • c(q_0) = tail_mass
  • c(q_1) = 0.5
  • c(q_2) = 1 - tail_mass