add the meta/register tokens used successfully in Hymba

Author: lucidrains

README.md

@@ -232,3 +232,12 @@ sampled = model.generate_text_only(text[:, :1], 1024)
     url     = {https://api.semanticscholar.org/CorpusID:273849947}
 }
 ```
+
+```bibtex
+@inproceedings{Dong2024HymbaAH,
+    title   = {Hymba: A Hybrid-head Architecture for Small Language Models},
+    author  = {Xin Dong and Y. Fu and Shizhe Diao and Wonmin Byeon and Zijia Chen and Ameya Mahabaleshwarkar and Shih-Yang Liu and Matthijs Van Keirsbilck and Min-Hung Chen and Yoshi Suhara and Yingyan Lin and Jan Kautz and Pavlo Molchanov},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:274166163}
+}
+```

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.6.5"
+version = "0.6.6"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -147,6 +147,15 @@ def inner(t: Tensor, *args, **kwargs) -> Tensor:
         return out
     return inner
 
+def pack_with_inverse(t, pattern):
+    packed, packed_shape = pack(t, pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        return unpack(out, packed_shape, inv_pattern)
+
+    return packed, inverse
+
 def pack_one_with_inverse(t, pattern):
     packed, packed_shape = pack([t], pattern)
 
@@ -1115,6 +1124,7 @@ def __init__(
         self,
         *,
         num_text_tokens,
+        num_register_tokens = 16,
         transformer: dict | Transformer,
         dim_latent: int | tuple[int, ...] | None = None,
         channel_first_latent: bool | tuple[bool, ...] = False,
@@ -1298,6 +1308,11 @@ def __init__(
         self.latent_to_model_projs = ModuleList(latent_to_model_projs)
         self.model_to_latent_projs = ModuleList(model_to_latent_projs)
 
+        # maybe register tokens (used in hymba, renamed from "meta" to register as "meta" was reserved from above already for the modality meta tag)
+
+        self.register_tokens = nn.Parameter(torch.zeros(num_register_tokens, dim))
+        nn.init.normal_(self.register_tokens, std = 0.02)
+
         # relative positions
 
         self.rotary_emb = RotaryEmbedding(transformer.dim_head)
@@ -2392,6 +2407,7 @@ def inner(pred_flow):
         if modality_positions.numel() == 0:
             modality_positions = F.pad(modality_positions, (0, 0, 0, 1))
 
+
         # sort the modalities tensor and sanitize, readying for noising of modalities
 
         modality_positions, sorted_indices = order_modality_positions_by_seq_offset(modality_positions)
@@ -2415,6 +2431,18 @@ def inner(pred_flow):
 
         tokens = einx.where('b n, b n d, b n d', is_any_modality, modality_tokens, text_tokens)
 
+        # handle maybe meta / register tokens
+
+        register_tokens = repeat(self.register_tokens, '... -> b ...', b = batch)
+
+        num_register_tokens = register_tokens.shape[-2]
+        seq_len += num_register_tokens
+
+        tokens, unpack_register_tokens = pack_with_inverse((register_tokens, tokens), 'b * d')
+        modality_positions[..., 1] += num_register_tokens
+
+        is_modalities = F.pad(is_modalities, (num_register_tokens, 0), value = False)
+
         # derive rotary positions
 
         rotary_positions = derive_rotary_positions_from_modality_positions(seq_len, modality_positions)
@@ -2455,6 +2483,10 @@ def inner(pred_flow):
             return_kv_cache = True
         )
 
+        # remove register tokens
+
+        _, embed = unpack_register_tokens(embed)
+
         # early return for embedding for decoding modality
 
         if return_embed: