Method forward

axlearn/common/quantizer.py:610–676

Quantization using Gumbel softmax trick. The code is selected based on the largest index of inputs. Inputs Gradients is computed with Gumbel softmax straight-through estimator. Args: inputs: tensor of shape [batch_size, seq_len, input_dim]. paddings:

(  # pytype: disable=signature-mismatch
        self, inputs: Tensor, *, paddings: Tensor
    )

        return params

    def forward(  # pytype: disable=signature-mismatch
        self, inputs: Tensor, *, paddings: Tensor
    ) -> tuple[BaseQuantizer.Output, dict[str, Tensor]]:
        """Quantization using Gumbel softmax trick.

        The code is selected based on the largest index of inputs. Inputs Gradients is computed
        with Gumbel softmax straight-through estimator.

        Args:
            inputs: tensor of shape [batch_size, seq_len, input_dim].
            paddings: 0/1 Tensor of shape [batch_size, seq_len].

        Returns:
            BaseQuantizer.Output.
            module_outputs contains temperature `tau` and prediction probability `probs`.
        """
        cfg = self.config
        # [batch_size, seq_len, num_codebooks, vocab_size].
        logits = self.input_proj(inputs=inputs)

        if self.is_training:
            tau = self.temperature_schedule(self.parameters["step"])
            logits = (
                logits + jax.random.gumbel(self.prng_key, shape=logits.shape, dtype=logits.dtype)
            ) / tau
            self.add_state_update("step", self.parameters["step"] + 1)

        # [batch_size, seq_len, num_codebooks].
        ids = jnp.argmax(logits, axis=-1)

        if not self.is_training:
            outputs = self.lookup(ids=ids)
            outputs = _apply_paddings(outputs=outputs, paddings=paddings)
        else:
            # [batch_size, seq_len, 1].
            mask = safe_not(paddings)[:, :, None]
            ids = ids * mask + (-1) * safe_not(mask)
            # TODO(dhwang2): optimize memory by scan for long context training.
            # [batch_size, seq_len, num_codebooks, vocab_size].
            onehots = _ids_to_onehots(ids, codebook_size=cfg.codebook_size, dtype=inputs.dtype)
            # We need this to stop gradients on the padded frames.
            onehots = onehots * mask[:, :, :, None]
            # [batch_size, seq_len, num_codebooks, vocab_size].
            y_soft = jax.nn.softmax(logits, axis=-1)
            y_soft = y_soft * mask[:, :, :, None]

            # Straight-through estimator such that dL/y_soft = dL/onehots.
            onehots = y_soft + jax.lax.stop_gradient(onehots - y_soft)
            batch_dims = _einsum_dims[: onehots.ndim - 2]
            quantized_vectors = jnp.einsum(
                f"{batch_dims}gv,vgh->{batch_dims}gh", onehots, self.parameters["codebook"]
            )
            quantized_vectors = quantized_vectors * mask[:, :, :, None]
            outputs = self.Output(
                # [batch_size, seq_len, num_codebooks].
                ids=ids,
                # [batch_size, seq_len, num_codebooks, codebook_dim].
                quantized_vectors=quantized_vectors,