Bytes encoder: read BOS and bytes, not EOS

README.md

@@ -24,7 +24,8 @@ has an embedding, and the latent transformer sees one vector per word.
 
 ## Model Setup
 
-- **Bytes Encoder** - A transformer over the bytes of each word, used bidirectionally. Its first (BOS) output is the word embedding.
+- **Bytes Encoder** - A transformer over BOS and the bytes of each word (not EOS, which the bidirectional encoder
+  does not need), used bidirectionally. Its first (BOS) output is the word embedding.
 - **Image Encoder** - A transformer over the 16x16 patches of each rendered word
   ([NaViT](https://arxiv.org/abs/2307.06304)-style: native size, with learned row and column embeddings, packed
   without padding; used bidirectionally). Its CLS output is the word embedding.

benchmarks/flops.py

@@ -61,7 +61,7 @@ def step_flops(config: dict, batches: int = 20) -> dict[str, float]:
         for kind, encoded in [("model", words), ("hardware", list(dict.fromkeys(words)))]:
             flops = common
             if "bytes_encoder" in configs:
-                flops += transformer_flops(configs["bytes_encoder"], [n for _, n, _ in encoded])
+                flops += transformer_flops(configs["bytes_encoder"], [n - 1 for _, n, _ in encoded])  # BOS, bytes
             if "image_encoder" in configs:
                 flops += transformer_flops(configs["image_encoder"], [p + 1 for _, _, p in encoded])
             totals[kind] += flops

tests/test_model.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 import torch.nn.functional as F  # noqa: N812
+from utf8_tokenizer.tokenizer import PAD_TOKEN_ID
 
 pytest.importorskip("megatron.bridge", reason="Requires the NeMo container")
 
@@ -39,6 +40,18 @@ def test_forward_is_finite(model, processor):
     assert (losses[0, :5] > 0).all()
 
 
+def test_bytes_encoder_ignores_eos(megatron, tiny_config, processor):
+    """The bytes encoder reads BOS and bytes only: a word's embedding does not depend on its EOS."""
+    model = build_model(tiny_config, image_encoder=False)
+    batch = processor(["hello world"])
+    ids = batch["input_ids"].cuda()
+    with torch.no_grad():
+        embeds = model.encode_words(ids)
+        last = (ids != PAD_TOKEN_ID).sum(-1, keepdim=True) - 1  # Each word's EOS
+        other = ids.scatter(-1, last.clamp(min=0), 1) * (ids != PAD_TOKEN_ID)  # A different id in its place
+        assert torch.equal(model.encode_words(other), embeds)
+
+
 def test_attention_no_look_ahead(model, processor):
     """Changing a later word does not change the losses of earlier words."""
     a = word_losses(model, processor, ["the quick brown fox jumps"])

welt/inference.py

@@ -95,7 +95,7 @@ def _encode_words(self, words: list[str]) -> dict[str, torch.Tensor]:
             embeds.append(self._pooled(self.image_encoder, prompts, params))
         if self.bytes_encoder is not None:
             tokenized = self.processor.tokenize_words(words)
-            prompts = [{"prompt_token_ids": ids[mask.bool()].tolist()}
+            prompts = [{"prompt_token_ids": ids[mask.bool()][:-1].tolist()}  # Without EOS, as WeLTModel.encode_words
                        for ids, mask in zip(tokenized.input_ids, tokenized.attention_mask, strict=True)]
             embeds.append(self._pooled(self.bytes_encoder, prompts, params))
         embeds = self._map(torch.cat(embeds, dim=-1).to(self.device), "encoder").cpu()

welt/model.py

@@ -338,7 +338,10 @@ def encode_words(self,
             patches = input_patches.view(-1, input_patches.size(-1))[index]
             embeds.append(self.image_encoder(patches, shapes[rows]))
         if self.bytes_encoder is not None:
-            embeds.append(self.bytes_encoder(input_ids[rows], words_mask[rows]))
+            # The bytes encoder reads each word's BOS and bytes, not its EOS: the decoder's target, it tells the
+            # (bidirectional) encoder nothing its length does not
+            mask = words_mask[rows]
+            embeds.append(self.bytes_encoder(input_ids[rows], mask & (mask.cumsum(-1) < mask.sum(-1, keepdim=True))))
 
         embeds = torch.cat(embeds, dim=-1)
         word_embeds = embeds.new_zeros(B * L, embeds.size(-1))