Method forward

axlearn/common/quantizer.py:457–555

Quantization with mse loss. Gradients for the codebook are computed only from MSE loss, while the gradients for the inputs are computed from both MSE and decoder losses approximated using the straight-through estimator. Args: inputs: Tensor of shape [bat

(self, inputs: Tensor, *, paddings: Tensor)

            return None

    def forward(self, inputs: Tensor, *, paddings: Tensor) -> BaseQuantizer.Output:
        """Quantization with mse loss.

        Gradients for the codebook are computed only from MSE loss, while the gradients for
        the inputs are computed from both MSE and decoder losses approximated using
        the straight-through estimator.

        Args:
            inputs: Tensor of shape [batch_size, seq_len, input_dim]. input_dim must equal to
                cfg.num_codebooks * cfg.codebook_dim.
            paddings: 0/1 Tensor of shape [batch_size, seq_len].

        Returns:
            BaseQuantizer.Output.
            module_outputs contains kmeans_loss and commitment_loss.

        Raises:
            ValueError: if inputs' last dimension does not match with the codebook dimensions.
        """
        cfg = self.config
        input_dim = cfg.num_codebooks * cfg.codebook_dim
        batch_size, seq_len = inputs.shape[:2]
        if inputs.shape[-1] != input_dim:
            raise ValueError(
                "inputs feature dimension should match with dims from all codebooks."
                f"{inputs.shape[-1]} != {cfg.num_codebooks} x {cfg.codebook_dim}."
            )
        inputs_by_group = jnp.reshape(
            inputs, [batch_size, seq_len, cfg.num_codebooks, cfg.codebook_dim]
        )
        quantized_inputs = quantize_by_nearest_neighbor(
            inputs=inputs_by_group,
            codebook=self.parameters["codebook"],
            metric=(
                SimilarityMetric.DOT_PRODUCT
                if cfg.normalize_codebook
                else SimilarityMetric.L2_DISTANCE
            ),
        )
        if cfg.normalize_inputs:
            inputs_by_group = l2_normalize(inputs_by_group, axis=-1, eps=1e-12)
        if cfg.normalize_codebook:
            self.parameters["codebook"] = l2_normalize(
                self.parameters["codebook"], axis=-1, eps=1e-12
            )
        quantized_inputs = _apply_paddings(outputs=quantized_inputs, paddings=paddings)

        # [batch_size, seq_len, input_dim].
        q_vecs = jnp.reshape(quantized_inputs.quantized_vectors, [batch_size, seq_len, input_dim])

        # Compute mean squared errors between q_vecs and inputs on non-padded frames.
        # Number of valid frames * input_dim.
        num_frames = jnp.sum(safe_not(paddings))
        denominator = jnp.maximum(num_frames * input_dim, 1)
        # Eq.3 of VQ-VAE paper https://arxiv.org/pdf/1711.00937.pdf.
        # The codebook is optimized by kmeans_loss only.
        inputs_to_loss = (
            jnp.reshape(inputs_by_group, [batch_size, seq_len, -1])