| --- |
| license: agpl-3.0 |
| language: |
| - en |
| library_name: transformers |
| datasets: |
| - allenai/c4 |
| tags: |
| - causal-lm |
| - byte-level |
| - raw-byte |
| - attention-free |
| - meyer |
| - wavelet |
| --- |
| |
|
|
| # TransWave |
|
|
|
|
| TransWave is a 96.13M-parameter attention-free causal raw-byte language model using learned block-local lifting and learned Meyer-structured multiresolution memory. |
|
|
|
|
| This repository contains the best checkpoint from the completed C4 English pretraining run. |
|
|
|
|
| ## Checkpoint |
|
|
|
|
| - Architecture: `causal_transwave_blocklifting_meyer_byte_lm_g4_v2` |
| - Parameters: 96,130,040 |
| - Input vocabulary: 256 raw byte values |
| - External BPE/SentencePiece tokenizer: none |
| - Trained C4 context: 20,480 model tokens |
| - Structural architecture maximum: 32,768 raw-byte tokens |
| - Physical training batch: 280 |
| - Training tokens consumed: 4,499,988,480 |
| - Optimizer steps: 785 |
| - Dataset: `allenai/c4:en` |
| - Objective: exact causal next-byte NLL |
| - Source commit: `5cae31b81e15160379ded79f32bcbb443d7db9f7` |
|
|
|
|
| One model token is one raw UTF-8 byte. |
|
|
|
|
| ## Validation result |
|
|
|
|
| The run exhausted its 4.5B-token budget before reaching the PPL=3 early-stop target. |
|
|
|
|
| - Best/final held-out NLL: 2.228725893 |
| - Best/final byte perplexity: 9.288024598 |
| - Validation tokens: 16,773,120 |
|
|
|
|
| ## Tokenization |
|
|
|
|
| There are two distinct layers and they should not be conflated. |
|
|
|
|
| `TransWaveByteTokenizer` is the Hugging Face input codec. It converts text to raw UTF-8 byte IDs 0...255 and back. It has no BPE merges and no learned external vocabulary. |
|
|
|
|
| The learned tokenizer is inside the neural network itself. `CTWBlockLiftingTokenizer` contains the trained shared byte embedding, three learned lifting stages and the learned cross-scale mixer. Those learned parameters are contained in `best_model.pt` and are loaded as part of `CTWModel`. |
|
|
|
|
| ## Architecture |
|
|
|
|
| Eight raw bytes form one block. A shared 256-entry embedding maps each byte to 96 channels. Three learned block-local lifting stages recursively produce eight multiresolution scales, giving a 768-dimensional block state. |
|
|
|
|
| Each of 12 layers combines a learned 512-block causal Meyer/Bernstein spectral FIR, learned full 8x8 scale mixing, learned spectral gates, a separate strictly causal DWConv7 path, RMS normalization and a SwiGLU feed-forward network of width 2816. |
|
|
|
|
| The global block representation is shifted by one block before decoding: block b is predicted from h[b-1], never h[b]. Within each target block, an 8-step GRUCell decoder of width 256 predicts bytes autoregressively while reusing the shared byte embedding for the classifier. |
|
|
|
|
| ## Load |
|
|
|
|
| ```python |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| |
| |
| repo = "Angshul/TransWave" |
| |
| |
| tokenizer = AutoTokenizer.from_pretrained( |
| repo, |
| trust_remote_code=True, |
| use_fast=False, |
| ) |
| |
| |
| model = AutoModelForCausalLM.from_pretrained( |
| repo, |
| trust_remote_code=True, |
| ).eval() |
| |
| |
| ids = tokenizer( |
| "TransWave operates on raw UTF-8 bytes.", |
| return_tensors="pt", |
| )["input_ids"] |
| |
| |
| # Exact scoring currently requires complete 8-byte blocks. |
| ids = ids[:, : (ids.shape[1] // 8) * 8] |
| |
| |
| with torch.no_grad(): |
| out = model(input_ids=ids, labels=ids) |
| |
| |
| print("NLL:", float(out.loss)) |
| print("PPL:", float(torch.exp(out.loss))) |
| ``` |
|
|
|
|
| ## Reference byte generation |
|
|
|
|
| The custom model exposes `generate_bytes`, an exact but deliberately simple reference generator: |
|
|
|
|
| ```python |
| prompt = tokenizer("The future of", return_tensors="pt")["input_ids"] |
| |
| |
| generated = model.generate_bytes( |
| prompt, |
| max_new_tokens=64, |
| do_sample=True, |
| temperature=0.8, |
| top_k=40, |
| ) |
| |
| |
| print(tokenizer.decode(generated[0].tolist())) |
| ``` |
|
|
|
|
| The reference generator recomputes completed-block context as needed; it is intended for correctness and checkpoint inspection, not optimized streaming throughput. |
|
|
|
|
| ## Context note |
|
|
|
|
| The architecture contains positional parameters for 4,096 eight-byte blocks = 32,768 raw-byte tokens. This checkpoint was trained on 20,480-token C4 contexts, so positions beyond 20,480 were not exercised in this C4 pretraining run. |
|
|
|
|
| ## Provenance |
|
|
|
|
| Implementation: https://github.com/AngshulMajumdar/TransWaveV1 |
|
|
|
|
| Paper: https://openreview.net/forum?id=3a22WBU5kT |
|
|
|
|
| The Hugging Face portability layer preserves the trained architecture. The only device-portability change in `transwave_core.py` is that the two explicitly-disabled CUDA autocast regions in the Meyer FFT path use the parameter device type, allowing the checkpoint to load and score on CPU as well as CUDA. |
