mixle.represent.quantize module¶
VectorQuantizer – learn a discrete vocabulary IN the shared embedding space, don’t guess it upstream.
Discrete tokens, when you want them (compression, transfer, a fixed vocabulary), come after embedding, not before segmentation: fit a codebook to the continuous vectors and each vector’s nearest code is its token id. The codebook is a learned model (k-means / a mixture), so the vocabulary is inferred from data rather than assumed – and because every modality is embedded into the same space, one codebook is a cross-modal vocabulary (an image patch and a word can share a token id when they land near the same centroid).
fit/quantize/dequantize are the codec; straight_through gives the VQ-VAE gradient so the codebook
and the encoders can be trained end to end under a generative or downstream objective. This is the only place
discreteness lives – the segmenter and embedding stay vocabulary-free.
- class VectorQuantizer(num_codes, dim, *, seed=0)[source]
Bases:
objectA learned codebook over
R^dim: nearest-centroid quantization of embedding vectors into discrete ids.- fit(vectors, *, iters=25)[source]
Fit the codebook by k-means (Lloyd) on
vectors(n, dim)– the vocabulary is learned, not assumed.
- quantize(vectors)[source]
Nearest-code id for each vector – the discrete token stream
(n,).
- dequantize(ids)[source]
Codebook vectors for token ids
(n,)->(n, dim)(the reconstruction / de-tokenization).
- reconstruction_error(vectors)[source]
Mean squared quantization error – the codebook’s fidelity (a codebook-size / bitrate knob).
- straight_through(vectors)[source]
VQ-VAE straight-through estimator: return quantized vectors but pass gradients to
vectorsunchanged.Lets the encoders and (with a codebook-commitment loss) the codebook train end to end through the discrete bottleneck.
vectorsis a torch tensor(n, dim).