vector_quantize

vector_quantize

Single-layer vector quantization with straight-through gradient estimator.

Wraps an EuclideanCodebook with optional input/output projections. Gradients flow through the quantization step via the straight-through estimator: forward uses the discrete codebook lookup, backward pretends it was an identity. Loss computation is handled externally by the metrics module (VQCommitmentLoss).

VectorQuantize

VectorQuantize(input_dimension, code_dim, num_codes, ema_decay=0.99, dead_code_threshold=1.0, kmeans_init=True)

Bases: Module

Single-layer vector quantizer with straight-through gradients.

Parameters:

Name	Type	Description	Default
input_dimension	int	Dimension of the input vectors from the encoder.	required
code_dim	int	Dimension of each codebook vector. If different from input_dimension, linear projections are added.	required
num_codes	int	Number of codebook entries (K).	required
ema_decay	float	EMA decay for codebook updates.	0.99
dead_code_threshold	float	Cluster size below which a code is replaced.	1.0
kmeans_init	bool	Initialize codebook from the first batch.	True

Source code in src/versatil/models/decoding/latent/vq/vector_quantize.py

def __init__(
    self,
    input_dimension: int,
    code_dim: int,
    num_codes: int,
    ema_decay: float = 0.99,
    dead_code_threshold: float = 1.0,
    kmeans_init: bool = True,
):
    super().__init__()
    if input_dimension <= 0:
        raise ValueError(
            f"input_dimension must be positive, got {input_dimension}."
        )
    if code_dim <= 0:
        raise ValueError(f"code_dim must be positive, got {code_dim}.")
    if num_codes <= 0:
        raise ValueError(f"num_codes must be positive, got {num_codes}.")
    self.input_dimension = input_dimension
    self.code_dim = code_dim
    self.num_codes = num_codes

    needs_projection = input_dimension != code_dim
    self.project_in = (
        nn.Linear(input_dimension, code_dim, bias=False)
        if needs_projection
        else nn.Identity()
    )
    self.project_out = (
        nn.Linear(code_dim, input_dimension, bias=False)
        if needs_projection
        else nn.Identity()
    )

    self.codebook = EuclideanCodebook(
        num_codes=num_codes,
        code_dim=code_dim,
        ema_decay=ema_decay,
        dead_code_threshold=dead_code_threshold,
        kmeans_init=kmeans_init,
    )

forward

forward(z_e)

Quantize input via nearest codebook lookup with straight-through gradient.

Parameters:

Name	Type	Description	Default
z_e	Tensor	Encoder output, shape (B, input_dimension).	required

Returns:

Type

Description

tuple[Tensor, Tensor, Tensor, Tensor]

Tuple of: z_q: Quantized output with straight-through gradient, shape (B, input_dimension). indices: Codebook indices, shape (B,). z_e_projected: Pre-quantization encoder output in code space, shape (B, code_dim). Carries gradient; used as the commitment-loss target for the encoder. z_q_code: Hard-quantized codebook vector in code space, shape (B, code_dim), detached from the computation graph. Paired with z_e_projected for per-layer commitment loss.

Source code in src/versatil/models/decoding/latent/vq/vector_quantize.py

def forward(
    self, z_e: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Quantize input via nearest codebook lookup with straight-through gradient.

    Args:
        z_e: Encoder output, shape (B, input_dimension).

    Returns:
        Tuple of:
            z_q: Quantized output with straight-through gradient,
                shape (B, input_dimension).
            indices: Codebook indices, shape (B,).
            z_e_projected: Pre-quantization encoder output in code space,
                shape (B, code_dim). Carries gradient; used as the
                commitment-loss target for the encoder.
            z_q_code: Hard-quantized codebook vector in code space,
                shape (B, code_dim), detached from the computation
                graph. Paired with z_e_projected for per-layer
                commitment loss.
    """
    z_e_projected = self.project_in(z_e)  # (B, code_dim)

    quantized, indices = self.codebook(z_e_projected)  # (B, code_dim), (B,)

    # Straight-through estimator: forward uses quantized, backward uses z_e_projected
    z_q = z_e_projected + (quantized - z_e_projected).detach()  # (B, code_dim)

    z_q = self.project_out(z_q)  # (B, input_dimension)

    return z_q, indices, z_e_projected, quantized.detach()