File size: 4,475 Bytes
ea51a3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | ---
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. |