From 5b2e7be2647168c533730c6082515cb1677bf9e4 Mon Sep 17 00:00:00 2001 From: autoresearch Date: Thu, 15 Oct 2026 07:12:36 -0400 Subject: [PATCH] Bytes encoder: read BOS and bytes, not EOS EOS is the bytes decoder's target (where a word ends); the bidirectional encoder already knows each word's length, so it reads one position fewer per word. Training (WeLTModel.encode_words), vLLM inference (WeLTGenerator), and the FLOPs benchmark alike. Saves 6-9% of step FLOPs with byte models split evenly (e.g. 25.6 -> 23.4 TFLOPs per step for the 25M latent transformer matched in parameters), under 1% with decoder-heavy ones. Checkpoints trained before this change read EOS: retrain (or evaluate with the earlier code). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015Xm62qoeaTunFSU5GYCQ1f --- README.md | 3 ++- benchmarks/flops.py | 2 +- tests/test_model.py | 13 +++++++++++++ welt/inference.py | 2 +- welt/model.py | 5 ++++- 5 files changed, 21 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 14e2cc3..edef512 100644 --- a/README.md +++ b/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. diff --git a/benchmarks/flops.py b/benchmarks/flops.py index 63806c5..082d44c 100644 --- a/benchmarks/flops.py +++ b/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 diff --git a/tests/test_model.py b/tests/test_model.py index 5c0e96a..1d94da6 100644 --- a/tests/test_model.py +++ b/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"]) diff --git a/welt/inference.py b/welt/inference.py index 509ec74..8c09953 100644 --- a/welt/inference.py +++ b/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() diff --git a/welt/model.py b/welt/model.py index d92eb4e..bd6abf2 100644 --- a/welt/model.py +++ b/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))