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

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:

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.

Downloads last month: 82

Angshul
/

TransWave