BitEmbedding: the bit-bias embedding as a plain module

tests/test_byte_embeddings.py

@@ -4,6 +4,7 @@
 from transformers import BertForMaskedLM
 
 from utf8_tokenizer.byte_embeddings import (
+    BitEmbedding,
     PatchedBitEmbeddings,
     join_embedding_layers,
     patch_embedding_layers,
@@ -224,3 +225,21 @@ def test_bit_projection_with_inference_mode(self, model, sample_input):
         # Verify output shape and validity
         assert output.shape == (2, 16, patched_embeddings.embeddings.embedding_dim)
         assert torch.isfinite(output).all()
+
+
+def test_bit_embedding_starts_as_its_table_and_learns_bits():
+    weight = torch.randn(256, 4)
+    embedding = BitEmbedding(weight)
+    ids = torch.tensor([[0, 1, 255]])
+    assert torch.equal(embedding(ids), weight[ids])  # The bit projection starts at zero
+
+    embedding(ids).sum().backward()
+    assert embedding.weight.grad is not None
+    assert embedding.bit_proj.grad.abs().sum() > 0
+
+    with torch.no_grad():
+        embedding.bit_proj.fill_(1.0)
+    # Byte 255 has all 8 bits set, byte 0 none
+    assert torch.allclose(embedding.folded_weight()[255], weight[255] + 8)
+    assert torch.allclose(embedding.folded_weight()[0], weight[0])
+    assert torch.allclose(embedding(ids)[0, 2], embedding.folded_weight()[255])

utf8_tokenizer/byte_embeddings.py

@@ -144,6 +144,28 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
         # Façade already refreshed by pre-hook; forward is just a normal embedding lookup
         return torch.nn.functional.embedding(input_ids, self.weight)
 
+class BitEmbedding(nn.Module):
+    """
+    The bit-bias embedding of PatchedBitEmbeddings as a plain module: a byte embedding table plus an additive
+    (zero-initialized) projection of each byte's 8 bits, W = base + bits(256,8) @ bit_proj.T(8,D).
+    Without a tieable façade or hooks: for models that own their embeddings (e.g. Megatron) and torch.compile.
+    folded_weight() is the equivalent plain embedding table, e.g. to export.
+    """
+
+    def __init__(self, weight: torch.Tensor):
+        super().__init__()
+        assert weight.shape[0] == 256, "Expected a byte-level embedding table (256 rows)."
+        self.weight = nn.Parameter(weight.detach().clone())
+        self.bit_proj = nn.Parameter(torch.zeros(weight.shape[1], 8, dtype=weight.dtype))
+        self.register_buffer("bits", unpack_bits(torch.arange(256, dtype=torch.uint8)), persistent=False)
+
+    def folded_weight(self) -> torch.Tensor:
+        return self.weight + self.bits.to(self.weight.dtype) @ self.bit_proj.T.to(self.weight.dtype)
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.embedding(input_ids, self.folded_weight())
+
+
 def tie_embeddings(model):
     try:
         # HuggingFace Transformers <=v4