euclidean_codebook

euclidean_codebook

Euclidean codebook with EMA updates for vector quantization.

Manages a set of learnable embedding vectors updated via exponential moving averages of encoder outputs, following van den Oord et al. (2017) and Tolstikhin et al. (2018). Supports automatic initialization on the first batch and dead code replacement.

EuclideanCodebook

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

Bases: Module

Codebook of embedding vectors with nearest-neighbor lookup and EMA updates.

Parameters:

Name	Type	Description	Default
num_codes	int	Number of codebook entries (K).	required
code_dim	int	Dimension of each codebook vector.	required
ema_decay	float	Exponential moving average decay for codebook updates. Higher values (e.g. 0.99) produce slower, more stable updates.	0.99
dead_code_threshold	float	Minimum average cluster size below which a code is considered dead and replaced with a random encoder output.	1.0
kmeans_init	bool	Initialize codebook vectors from the first batch (True) or from N(0, 1) (False).	True

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

def __init__(
    self,
    num_codes: int,
    code_dim: int,
    ema_decay: float = 0.99,
    dead_code_threshold: float = 1.0,
    kmeans_init: bool = True,
):
    super().__init__()
    if num_codes <= 0:
        raise ValueError(f"num_codes must be positive, got {num_codes}.")
    if code_dim <= 0:
        raise ValueError(f"code_dim must be positive, got {code_dim}.")
    if not 0.0 <= ema_decay < 1.0:
        raise ValueError(
            f"ema_decay must be in the interval [0.0, 1.0), got {ema_decay}."
        )
    if dead_code_threshold < 0.0:
        raise ValueError(
            f"dead_code_threshold must be non-negative, got {dead_code_threshold}."
        )
    self.num_codes = num_codes
    self.code_dim = code_dim
    self.ema_decay = ema_decay
    self.dead_code_threshold = dead_code_threshold
    self.kmeans_init = kmeans_init

    if kmeans_init:
        embed = torch.zeros(
            num_codes, code_dim
        )  # (K, D) — placeholder until first batch
    else:
        embed = torch.randn(num_codes, code_dim)  # (K, D) — random init

    self.register_buffer("embed", embed)  # (K, D)
    self.register_buffer("cluster_size", torch.zeros(num_codes))  # (K,)
    self.register_buffer("embed_avg", embed.clone())  # (K, D)
    self.register_buffer(
        "initialized", torch.tensor(not kmeans_init)
    )  # scalar bool

forward

forward(z_e)

Quantize encoder outputs to nearest codebook entries.

During training, updates codebook via EMA and replaces dead codes.

Parameters:

Name	Type	Description	Default
z_e	Tensor	Encoder outputs, shape (B, D).	required

Returns:

Type	Description
tuple[Tensor, Tensor]	Tuple of: quantized: Nearest codebook vectors, shape (B, D). indices: Codebook indices for each input, shape (B,).

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

def forward(self, z_e: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Quantize encoder outputs to nearest codebook entries.

    During training, updates codebook via EMA and replaces dead codes.

    Args:
        z_e: Encoder outputs, shape (B, D).

    Returns:
        Tuple of:
            quantized: Nearest codebook vectors, shape (B, D).
            indices: Codebook indices for each input, shape (B,).
    """
    if self.kmeans_init and not self.initialized:
        self._initialize_from_data(z_e.detach())
        self._sync_buffers_from_primary_rank()

    # z_e: (B, D), self.embed: (K, D) -> dist: (B, K)
    dist = torch.cdist(z_e, self.embed)
    indices = dist.argmin(dim=-1)  # (B,)
    quantized = self.embed[indices]  # (B, D)

    if self.training:
        # One-hot assignment matrix: (B, K)
        one_hot = torch.zeros(z_e.shape[0], self.num_codes, device=z_e.device)
        one_hot.scatter_(
            1, indices.unsqueeze(1), 1.0
        )  # indices: (B,) -> (B, 1) for scatter

        # Cluster statistics for EMA
        new_cluster_size = one_hot.sum(dim=0)  # (K,)
        new_embed_sum = one_hot.T @ z_e.detach()  # (K, B) @ (B, D) -> (K, D)
        if torch_dist.is_available() and torch_dist.is_initialized():
            torch_dist.all_reduce(new_cluster_size, op=torch_dist.ReduceOp.SUM)
            torch_dist.all_reduce(new_embed_sum, op=torch_dist.ReduceOp.SUM)

        # EMA update: running_avg = decay * old + (1 - decay) * new
        self.cluster_size.data.mul_(self.ema_decay).add_(
            new_cluster_size, alpha=1.0 - self.ema_decay
        )  # (K,)
        self.embed_avg.data.mul_(self.ema_decay).add_(
            new_embed_sum, alpha=1.0 - self.ema_decay
        )  # (K, D)

        # Normalize to get updated embeddings: embed = embed_avg / cluster_size
        smoothed_cluster_size = (
            self.cluster_size + 1e-5
        )  # (K,) — avoid division by zero
        self.embed.data.copy_(
            self.embed_avg
            / smoothed_cluster_size.unsqueeze(1)  # (K, 1) broadcast over D
        )  # (K, D)

        if self._replace_dead_codes(z_e.detach()):
            # In case of distributed training, keep nodes in sync.
            self._sync_buffers_from_primary_rank()

    return quantized, indices