Factorized Embedding

Factorized token embedding with optional Conv1d pre-projection.

class src.model.embeddings.factorized_embedding.FactorizedEmbedding(*args: Any, **kwargs: Any)

Bases: Module

Factorized token embedding with optional Conv1d pre-projection.

Reduces the embedding lookup dimension to factorized_embedding_dim, then projects up to hidden_size via a linear (or BitLinear) layer. Optionally applies a 1D convolution over the embedding stream for local context smoothing before projection.

low_dim

Reduced embedding dimension.

use_conv

Whether the Conv1d pre-projection is active.

embedding

Token embedding lookup table.

conv

Optional Conv1d layer for local context smoothing.

proj

Linear (or BitLinear) projection from low_dim to hidden_size.

__init__(config: FrankensteinModelConfig)

Initialize factorized embedding from a FrankensteinModelConfig.

Parameters:

config – Model configuration. Reads factorized_embedding_dim, vocab_size, use_embedding_conv, embedding_conv_kernel, use_bitnet, and hidden_size.

forward(input_ids: torch.Tensor) → torch.Tensor

Embed token IDs and optionally apply Conv1d + projection.

Sequence length is forced to match input_ids after convolution to keep MLM labels aligned (even-kernel symmetric padding can shift length by +1).

Parameters:

input_ids – Integer token indices of shape (B, S).

Returns:

Tensor of shape (B, S, hidden_size).