Three-Level Nested Byte Mamba-2

This repository contains a research checkpoint for a 2.478B-parameter causal byte model with three nested Mamba-2 resolutions. It predicts raw bytes rather than tokenizer IDs and was trained on a mixture of web/PDF text, serialized image-text examples, and serialized audio.

It is not a Transformers AutoModel checkpoint and is not instruction-formatted as a conventional chat model. Use the included cached inference script.

Checkpoint contents

The published weights are sharded SafeTensors containing only the 980 model tensors. The original optimizer, scaler, training phase, data cursor, dataset paths, source fingerprints, and other training-only checkpoint objects were removed.

• Parameters: 2,478,820,575
• Weight precision on disk: FP32
• Raw tensor size: 9,915,282,300 bytes
• Source checkpoint step: 889,000
• Recommended runtime precision: BF16
• Recommended placement: fine/decoder on cuda:0, level 2 on cuda:1, level 3 on cuda:2
• Last training count: 15GB Absoloutly undertrained

The source checkpoint step is documentation only; it is not embedded in the SafeTensors weights or inference configuration.

Latest validation results

The latest recorded validation event is step 890,000, one scheduled validation event after the packaged last.pt weight step.

The aggregate should not be interpreted as a pure language score: audio accounts for most evaluated bytes and has a substantially different entropy scale. For text use, the JSONL and Parquet rows are the relevant measurements.

Architecture

Byte vocabulary

There is no learned tokenizer:

PAD=0, BOS=1, EOS=2, UNK=3
raw byte 0..255 -> ID 4..259
vocabulary size = 260

UTF-8 text and serialized binary modalities therefore share one next-byte objective.

Three causal resolutions

1. Fine level: a local causal convolutional encoder and 6 Mamba-2 blocks operate at byte resolution. A learned causal boundary head closes variable pools between 1 and 96 bytes.
2. Level 2: 20 Mamba-2 blocks consume completed fine-pool states. A learned boundary head groups 4–16 completed fine pools.
3. Level 3: 30 Mamba-2 blocks consume completed level-2 states and group 2–16 level-2 pools.

Every Mamba block uses model width 2,000, Mamba-2 d_state=64, and head dimension 100. A pool can use only states already available in its causal prefix. A closure never revises an earlier prediction.

Fusion decoder

For each byte, the decoder concatenates four 2,000-dimensional signals:

• byte-local contextual state;
• current fine latent;
• latest level-2 latent;
• latest level-3 latent.

The 10,000-dimensional concatenation is normalized, projected through an 8,000-wide GELU fusion layer, reduced to width 2,000, and mapped to 260 next-byte logits. This enlarged decoder was added to avoid choking the information arriving from three recurrent resolutions.

Pool-density fallback

Fine pooling includes a rolling short-pool quota. Among the most recent 6,000 completed fine pools, at most 3,000 may be shorter than 6 bytes. When that quota fills, the next pool must reach the secondary minimum; short closures become eligible again as older short pools leave the rolling window. The quota counts completed pools, not raw bytes.

Delayed decoder controller

The checkpoint includes an optional hold/refresh/compress controller. Its output at time t can influence closure only at t+1:

decode byte t -> controller C[t] -> choose closure at t+1 -> decode byte t+1

The included cached inference path applies this without future leakage or a second full-model pass.

Parameter distribution

Pooling reduces sequence activations and recurrent update frequency, not layer-weight storage. This is why the deepest level remains the largest parameter group even though it updates least frequently.

Inference

Dependencies

Use Linux, CUDA, and versions of PyTorch, mamba-ssm, Triton, and causal-conv1d that are mutually compatible:

pip install -r requirements.txt

BF16 is strongly recommended. FP16 cached rollouts can become numerically unstable on some Mamba-2 builds.

Three-GPU inference

From the downloaded repository:

python infer_nested_model.py \
  --checkpoint . \
  --prompt "The history of state space models begins" \
  --max-new-bytes 512 \
  --precision bf16 \
  --fine-device cuda:0 \
  --nested-devices cuda:1 \
  --tertiary-device cuda:2 \
  --temperature 0.8 \
  --top-p 0.9

The script accepts --prompt-file for arbitrary byte prefixes, --output for raw generated bytes, --html-output for hierarchy-attribution output, and --image-output-dir to extract complete generated P6 images.

Single-GPU inference is supported when the GPU can hold the requested precision:

python infer_nested_model.py --checkpoint . --device cuda:0 \
  --prompt "Once upon a time" --max-new-bytes 256 --precision bf16

Stateful generation

Generation prefills the prompt once, then caches the convolution and SSM states for the fine, level-2, and level-3 stacks. New bytes advance those caches token by token; the entire prefix is not reprocessed for every generated byte.

Training mixture and modality representation

The training run mixed educational web text, PDF-derived text, image/question and instruction examples, and paired music/cover data from the following repositories:

These datasets are not redistributed in this repository. Their upstream terms continue to apply independently and are not replaced by this repository's license.

• Text and code are UTF-8 bytes.
• Images are complete RGB PPM byte sequences plus associated text.
• Audio uses 24 kHz EnCodec payloads with generation and detection objectives.
• Instruction and dialogue fields present in source records were serialized in full rather than using assistant-response-only loss.

This mixture makes the checkpoint experimental and general-purpose at the byte level; it does not guarantee strong image or audio generation quality.

Limitations

• This is custom research code, not an official Mamba or Transformers architecture.
• The model is not a safety-aligned chat assistant.
• Raw-byte sampling can produce invalid UTF-8, malformed images, or incomplete audio containers.
• Image training used small PPM rasters, limiting fine visual detail.
• Audio validation remains much weaker than text validation.
• The audio corpus includes third-party creator material whose source rights are retained. The model license does not grant rights to reproduce protected training content, lyrics, compositions, voices, or recordings.
• The current weights are FP32 and large; practical use generally requires BF16 casting.
• The delayed pooling controller is causal but makes exact routing inherently sequential.
• The latest CSV validation event is at step 890,000, while the packaged last.pt weights identify step 889,000; the table must therefore be read as the latest run validation, not an evaluation re-run performed directly on this exported artifact.

Intended use

Intended for research into byte-level modeling, hierarchical state-space models, adaptive causal pooling, long recurrent context, and mixed text/binary generation. Validate outputs independently before using them in downstream systems.

License

The model weights, model card, and visual assets are available under CC BY-NC 4.0. The Python inference source is available under Apache License 2.0. Training datasets and third-party content are not covered by either grant. See LICENSE.md for the precise repository scope and notices.

Repository files

• model-*.safetensors: inference-only model shards
• model.safetensors.index.json: tensor-to-shard map
• config.json: architecture-only inference configuration
• modeling_nested_mamba.py: custom model implementation
• nested_inference_tools.py: SafeTensors loading, state caching, and sampling
• infer_nested_model.py: command-line generator
• architecture.svg: architecture visualization
• LICENSE.md: weight, documentation, and code license scope