ml.models.quantization.lfq

Provides an implementation of Lookup-Free Quantization (LFQ).

LFQ is from the paper Language Model Beats Diffusion - Tokenizer is Key to Visual Generation, which purports to beat image generation using language models simply by using a high-quality tokenizer.

class ml.models.quantization.lfq.Losses(per_sample_entropy: torch.Tensor, batch_entropy: torch.Tensor, commitment: torch.Tensor)

Bases: object

per_sample_entropy: Tensor

batch_entropy: Tensor

commitment: Tensor

single() → Tensor

ml.models.quantization.lfq.euclidean_distance_squared(x: Tensor, y: Tensor) → Tensor

ml.models.quantization.lfq.entropy(prob: Tensor, eps: float = 1e-20) → Tensor

class ml.models.quantization.lfq.LookupFreeQuantization(*, dim: int | None = None, codebook_size: int | None = None, entropy_loss_weight: float = 0.1, commitment_loss_weight: float = 1.0, diversity_gamma: float = 2.5, num_codebooks: int = 1, codebook_scale: float = 1.0)

Bases: Module

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

mask: Tensor

zero: Tensor

codebook: Tensor

bits_to_codes(bits: Tensor) → Tensor

property dtype: dtype

forward(x: Tensor, inv_temperature: float = 1.0) → tuple[torch.Tensor, torch.Tensor, ml.models.quantization.lfq.Losses]

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.