residual_vq

residual_vq

Residual vector quantization with cascading codebook layers.

Each layer quantizes the residual left by previous layers, producing a hierarchical discrete representation. Coarse layers capture the dominant structure, fine layers capture residual detail.

ResidualVQ

ResidualVQ(input_dimension, code_dim, num_codes, num_layers=1, ema_decay=0.99, dead_code_threshold=1.0, kmeans_init=True)

Bases: Module

Multi-layer residual vector quantizer.

Cascades N VectorQuantize layers. Each layer quantizes the residual from the previous layer. The final quantized output is the sum of all layers' contributions.

Parameters:

Name	Type	Description	Default
input_dimension	int	Dimension of the input vectors.	required
code_dim	int	Dimension of each codebook vector per layer.	required
num_codes	int	Number of codebook entries per layer (K).	required
num_layers	int	Number of residual VQ layers.	1
ema_decay	float	EMA decay for codebook updates in each layer.	0.99
dead_code_threshold	float	Dead code replacement threshold per layer.	1.0
kmeans_init	bool	Initialize each layer's codebook from data.	True

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

def __init__(
    self,
    input_dimension: int,
    code_dim: int,
    num_codes: int,
    num_layers: int = 1,
    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}.")
    if num_layers <= 0:
        raise ValueError(f"num_layers must be positive, got {num_layers}.")
    self.input_dimension = input_dimension
    self.code_dim = code_dim
    self.num_codes = num_codes
    self.num_layers = num_layers

    self.layers = nn.ModuleList(
        [
            VectorQuantize(
                input_dimension=input_dimension,
                code_dim=code_dim,
                num_codes=num_codes,
                ema_decay=ema_decay,
                dead_code_threshold=dead_code_threshold,
                kmeans_init=kmeans_init,
            )
            for _ in range(num_layers)
        ]
    )

forward

forward(z_e)

Quantize input through cascading residual layers.

Parameters:

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

Returns:

Type

Description

tuple[Tensor, list[Tensor], Tensor, Tensor]

Tuple of: z_q: Sum of all layers' quantized outputs, shape (B, input_dimension). all_indices: List of per-layer codebook indices, each shape (B,). z_e_per_layer: Per-layer pre-quantization encoder outputs in code space, stacked along dim 0, shape (L, B, code_dim). Carries gradient; used as commitment-loss target. z_q_per_layer: Per-layer hard-quantized codebook vectors in code space, stacked along dim 0, shape (L, B, code_dim). Detached; paired with z_e_per_layer for per-layer commitment loss.

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

def forward(
    self, z_e: torch.Tensor
) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor, torch.Tensor]:
    """Quantize input through cascading residual layers.

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

    Returns:
        Tuple of:
            z_q: Sum of all layers' quantized outputs, shape (B, input_dimension).
            all_indices: List of per-layer codebook indices, each shape (B,).
            z_e_per_layer: Per-layer pre-quantization encoder outputs in
                code space, stacked along dim 0, shape (L, B, code_dim).
                Carries gradient; used as commitment-loss target.
            z_q_per_layer: Per-layer hard-quantized codebook vectors in
                code space, stacked along dim 0, shape (L, B, code_dim).
                Detached; paired with z_e_per_layer for per-layer
                commitment loss.
    """
    residual = z_e  # (B, input_dimension)
    z_q_hard_total = torch.zeros_like(z_e)  # (B, input_dimension)
    z_q_straight_through_total = torch.zeros_like(z_e)  # (B, input_dimension)
    all_indices = []
    all_z_e_projected = []
    all_z_q_code = []

    for layer in self.layers:
        z_q_layer, indices, z_e_projected, z_q_code = layer(
            residual
        )  # (B, input_dimension), (B,), (B, code_dim), (B, code_dim)
        residual = (
            residual - z_q_layer.detach()
        )  # (B, input_dimension) — detach to stop gradient across layers
        z_q_hard_total = z_q_hard_total + z_q_layer.detach()  # (B, input_dimension)
        z_q_straight_through_total = (
            z_q_straight_through_total + z_q_layer
        )  # (B, input_dimension)
        all_indices.append(indices)
        all_z_e_projected.append(z_e_projected)
        all_z_q_code.append(z_q_code)

    z_e_per_layer = torch.stack(all_z_e_projected, dim=0)  # (L, B, code_dim)
    z_q_per_layer = torch.stack(all_z_q_code, dim=0)  # (L, B, code_dim)
    # Keep the hard RVQ sum in the forward pass, but average the per-layer
    # straight-through paths so an L-layer RVQ does not multiply encoder
    # gradients by L when the projections are identities.
    z_q_total = (
        z_q_hard_total
        + (z_q_straight_through_total - z_q_straight_through_total.detach())
        / self.num_layers
    )

    return z_q_total, all_indices, z_e_per_layer, z_q_per_layer

decode_from_indices

decode_from_indices(all_indices)

Reconstruct quantized output from codebook indices.

Used by the prior at inference to convert predicted indices into the quantized embedding that the decoder expects.

Parameters:

Name	Type	Description	Default
all_indices	list[Tensor]	List of per-layer codebook indices, each shape (B,).	required

Returns:

Type	Description
Tensor	Reconstructed quantized output, shape (B, input_dimension).

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

def decode_from_indices(self, all_indices: list[torch.Tensor]) -> torch.Tensor:
    """Reconstruct quantized output from codebook indices.

    Used by the prior at inference to convert predicted indices into
    the quantized embedding that the decoder expects.

    Args:
        all_indices: List of per-layer codebook indices, each shape (B,).

    Returns:
        Reconstructed quantized output, shape (B, input_dimension).
    """
    z_q_total = torch.zeros(
        all_indices[0].shape[0],
        self.input_dimension,
        device=all_indices[0].device,
    )  # (B, input_dimension)

    for layer, indices in zip(self.layers, all_indices, strict=True):
        codebook_vectors = layer.codebook.embed[indices]  # (B, code_dim)
        z_q_total = z_q_total + layer.project_out(
            codebook_vectors
        )  # (B, input_dimension)

    return z_q_total