vector_quantize

→Tensor

vector_quantize(x: Tensor, codebook: Tensor)

source edit

Snap x to its nearest codebook entries, straight-through.

The functional core of lucid.nn.VectorQuantizer — van den Oord, Vinyals, and Kavukcuoglu, "Neural Discrete Representation Learning" (2017). Each row of x is replaced by the closest entry of codebook, and the result is routed through straight_through so the producer of x trains as though quantisation were the identity.

Parameters

xTensor

Field of shape (*, D); the trailing axis is the feature dimension.

codebookTensor

Codebook of shape (K, D).

Returns

Tensor

Shape (*, D), numerically equal to the selected entries and differentiable with respect to x.

Notes

The returned quantized carries no gradient to codebook — the straight-through path routes past it by construction. Training the codebook needs the separate term $\|\mathrm{sg}[x] - e\|_2^2$ , which lucid.nn.VectorQuantizer builds and returns alongside the commitment term. Calling this function directly and optimising only a reconstruction loss leaves the codebook frozen at its initialisation — a silent failure that looks like a model which simply will not learn.

Examples

>>> import lucid
>>> import lucid.nn.functional as F
>>> codebook = lucid.tensor([[0.0, 0.0], [1.0, 1.0]])
>>> x = lucid.tensor([[0.9, 1.1]], requires_grad=True)
>>> quantized, indices = F.vector_quantize(x, codebook)
>>> quantized.tolist(), indices.tolist()
([[1.0, 1.0]], [1])

Used by 1

lucid.nn.functional

>>> import lucid >>> import lucid.nn.functional as F >>> codebook = lucid.tensor([[0.0, 0.0], [1.0, 1.0]]) >>> x = lucid.tensor([[0.9, 1.1]], requires_grad=True) >>> quantized, indices = F.vector_quantize(x, codebook) >>> quantized.tolist(), indices.tolist() ([[1.0, 1.0]], [1])