Embedding

Quantization-aware embedding — trainable float table, fake-quant every forward.

The training-time stand-in that lucid.quantization.prepare_qat installs in place of a float lucid.nn.Embedding. It keeps a trainable float lookup table but fake-quantizes that table on every forward, so the learned embeddings experience the int8 rounding error while training and settle into values that survive the eventual int8 storage. This is what lets a QAT-trained embedding match its float accuracy after conversion, where a naively-quantized table would shift every token vector by an uncontrolled rounding step.

Only the table is fake-quantized. Unlike the QAT lucid.nn.qat.Linear / conv layers, there is no activation observer here: the layer's output is a plain gather of rows that are already on the (fake-)quantized grid, so the output range is fully determined by the table and a second observer would be redundant. The module therefore carries a single weight_fake_quant and no activation_post_process.

Straight-through estimator (STE). Rounding a table entry is a step function whose derivative is zero almost everywhere; the fake-quant rounds in the forward pass but passes the gradient straight through to the float table in the backward pass, so the embeddings stay fully trainable. lucid.quantization.convert then reads the trained table and the weight observer's final qparams and bakes them into an inference lucid.nn.quantized.Embedding whose rows are stored as int8 codes.

On each forward the table is fake-quantized, then the rows for x are gathered:

where $s, z$ are the scale / zero-point the weight FakeQuantize derives from its observer and $q_{\min}, q_{\max}$ are the grid bounds. The straight-through unit derivative keeps the table trainable despite the non-differentiable round.