ml.models.quantization.fsq

Provides an implementation of Finite Scalar Quantization (FSQ).

FSQ is a quantization approach which has a relatively small number of parameters compared with codebook learning. It was proposed in the paper Finite Scalar Quantization: VQ-VAE Made Simple.

This implementation is largely adapted from the lucidrains implementation which in turn tracks very closely with the original implementation.

ml.models.quantization.fsq.round_ste(z: Tensor) → Tensor

class ml.models.quantization.fsq.FiniteScalarQuantization(levels: list[int])

Bases: Module

Defines a finite scalar quantization module.

The original paper proposes the following number of levels, depending on the target codebook size:

Codebook size	Number of levels
2^8	8, 6, 5
2^10	8, 5, 5, 5
2^12	7, 5, 5, 5, 5
2^14	8, 8, 8, 6, 5
2^16	8, 8, 8, 5, 5, 5

Parameters:

levels – The number of levels. The product of the levels is the number of unique codes. The input to the module should be a tensor with shape (..., len(levels)).

Properties:

dim: The number of dimensions of the quantized tensor, i.e. the length

of the levels argument.

n_codes: The number of unique codes.

Inputs:

z: A tensor of shape (..., len(levels)).

Outputs:

quantized: A quantized tensor of shape (..., len(levels)). The

quantized values will be in the range [-1, 1].

Initializes internal Module state, shared by both nn.Module and ScriptModule.

implicit_codebook: Tensor

forward(z: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

quantize(z: Tensor) → Tensor

codes_to_indices(zhat: Tensor) → Tensor

indices_to_codes(indices: Tensor) → Tensor