From e5377cdf96babb009cc8a9d340cca25851e5d80d Mon Sep 17 00:00:00 2001
From: autoresearch <autoresearch@local>
Date: Sat, 10 Oct 2026 02:36:20 -0400
Subject: [PATCH] BitEmbedding: the bit-bias embedding as a plain module
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PatchedBitEmbeddings wraps a HF model's embeddings behind a tieable façade weight, refreshed by hooks. Models that
own their embeddings (WeLT on Megatron) and torch.compile need the same math as a plain module: a table plus a
zero-initialized projection of each byte's bits, with folded_weight() as the equivalent table for export.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f
---
 tests/test_byte_embeddings.py     | 19 +++++++++++++++++++
 utf8_tokenizer/byte_embeddings.py | 22 ++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/tests/test_byte_embeddings.py b/tests/test_byte_embeddings.py
index 3027780..bcf1797 100644
--- a/tests/test_byte_embeddings.py
+++ b/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])
diff --git a/utf8_tokenizer/byte_embeddings.py b/utf8_tokenizer/byte_embeddings.py
index a23fee3..50219a8 100644
--- a/utf8_tokenizer/byte_embeddings.py
+++ b/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