# helia_edge.layers.ema_residual_vector_quantizer

Residual Vector Quantizer with Exponential Moving Average codebook updates.

Replaces the gradient-based codebook loss with EMA updates to codebook
embeddings, following van den Oord et al. 2017 (VQ-VAE).  Only the
commitment loss is back-propagated; codebook vectors are updated via
running averages of assigned encoder outputs.

## helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer

`class` · `python`

```python
EmaResidualVectorQuantizer(
    num_levels: int,
    num_embeddings: int | Sequence[int],
    embedding_dim: int,
    beta: float = 0.25,
    ema_decay: float = 0.99,
    epsilon: float = 1e-05,
    **kwargs={},
)
```

Residual VQ with EMA codebook updates.

Instead of learning codebook embeddings via gradient descent (which requires
a codebook loss term), this layer maintains exponential moving averages of
cluster assignment counts and embedding sums.  Codebook vectors are derived
from these running statistics with Laplace smoothing for numerical stability.

Only the *commitment loss* is back-propagated through the encoder; the
straight-through estimator copies gradients from the decoder to the encoder
as in the standard VQ-VAE.

Input:  ``[..., D]``  (last dim = ``embedding_dim``)
Output: ``[..., D]``  (sum of per-level dequantized vectors)

Metrics (logged via ``metrics`` property):
    - ``rvq_l{l}_perplexity``, ``rvq_l{l}_usage``,
      ``rvq_l{l}_bits_per_index``
    - ``rvq_perplexity_mean``, ``rvq_usage_mean``,
      ``rvq_bits_per_index_sum`` (entropy lower bound)

:::note[Losses added per level]
- ``beta * ||stop(q_l) - r_l||^2``  (commitment only; no codebook
  gradient loss)

:::

Example:

```python
rvq = EmaResidualVectorQuantizer(
    num_levels=4,
    num_embeddings=64,
    embedding_dim=16,
    ema_decay=0.99,
)
y = rvq(z, training=True)          # forward + EMA update
y, indices = rvq(z, return_indices=True)  # also return codes
```

:::note[References]
- van den Oord, A., Vinyals, O. & Kavukcuoglu, K. (2017).
  *Neural Discrete Representation Learning*. NeurIPS.

:::

**Parameters**

| Name | Type | Default | Description |
| --- | --- | --- | --- |
| num_levels | int | Required | Number of residual VQ stages (``M >= 1``). |
| num_embeddings | int \| Sequence[int] | Required | Codebook size ``K`` per level (int or per-level list). |
| embedding_dim | int | Required | Latent dimensionality ``D``. |
| beta | float | 0.25 | Commitment loss coefficient. |
| ema_decay | float | 0.99 | EMA decay rate for codebook updates (``0.99``–``0.999`` typical). |
| epsilon | float | 1e-05 | Small constant for Laplace smoothing of cluster counts. |

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:18`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.M

`constant` · `python`

```python
M = int(num_levels)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:84`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.D

`constant` · `python`

```python
D = int(embedding_dim)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:85`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.Ks

`attribute` · `python`

```python
Ks
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:89`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.beta

`attribute` · `python`

```python
beta = float(beta)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:92`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.ema_decay

`attribute` · `python`

```python
ema_decay = float(ema_decay)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:93`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.epsilon

`attribute` · `python`

```python
epsilon = float(epsilon)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:94`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.metrics

`attribute` · `python`

```python
metrics
```

Expose per-level + aggregate metrics so ``Model.fit`` logs them.

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:264`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.build

`method` · `python`

```python
build(input_shape)
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:109`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.call

`method` · `python`

```python
call(
    x: keras.KerasTensor,
    training: bool = False,
    return_indices: bool = False,
) -> keras.KerasTensor | tuple[keras.KerasTensor, list[keras.KerasTensor]]
```

Quantize *x* through all residual levels.

**Parameters**

| Name | Type | Default | Description |
| --- | --- | --- | --- |
| x | keras.KerasTensor | Required | ``[..., D]`` latent to be quantized. |
| training | bool | False | If ``True``, run EMA codebook updates. |
| return_indices | bool | False | If ``True``, also return per-level flat indices. |

**Returns**

| Name | Type | Description |
| --- | --- | --- |
|  | keras.KerasTensor \| tuple[keras.KerasTensor, list[keras.KerasTensor]] | ``y`` or ``(y, indices_list)``: dequantized vector and optional |
|  | keras.KerasTensor \| tuple[keras.KerasTensor, list[keras.KerasTensor]] | per-level index tensors. |

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:195`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.encode

`method` · `python`

```python
encode(x: keras.KerasTensor) -> list[keras.KerasTensor]
```

Return list of per-level flat index tensors ``[N]`` (no gradients).

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:272`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.decode

`method` · `python`

```python
decode(indices_list: list[keras.KerasTensor], original_shape: tuple[int, ...]) -> keras.KerasTensor
```

Sum per-level code vectors from *indices_list* and reshape.

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:284`

### helia_edge.layers.ema_residual_vector_quantizer.EmaResidualVectorQuantizer.get_config

`method` · `python`

```python
get_config()
```

Source: `helia_edge/layers/ema_residual_vector_quantizer.py:296`