Llama-MoB-3.2-1B-MLP

A compact Mixture-of-Bernoullis output head from MoB_Head. This release contains head weights and a token-ID-indexed semantic codebook, for use with meta-llama/Llama-3.2-1B.

Built with Llama. This artifact is subject to the Llama 3.2 Community License.

Architecture

• Head family: mob.
• Constructor arguments: {"depth": 3, "h": 2048, "k": 128, "m": 17, "width": 3072}.
• Head parameters: 32,255,232.
• Vocabulary: 128,256 tokens.
• Serialization: fp32 safetensors, without optimizer state.

The backbone and tokenizer are separate dependencies. This is not a complete Transformers causal language model, and AutoModelForCausalLM cannot load it. Do not substitute a different tokenizer or token-ID ordering.

Training

The backbone is frozen. The head minimizes exact teacher-probability-weighted code NLL on the teacher's renormalized top-1024 support. No sampled targets or continuous target noise are used. Stability controls include vertex initialization, mixture entropy regularization and mixture-temperature annealing.

{
  "data": "FineWeb-Edu",
  "teacher_topk": 1024,
  "sequence_length": 512,
  "initial_schedule": "WSD",
  "initial_peak_lr": 0.0007,
  "polish_updates": 50000,
  "polish_peak_lr": 0.0002
}

Recorded evaluation

{
  "normalized_ppl_ratio": 2.0,
  "round_agreement": 0.774
}

Normalized perplexity uses a sum over valid vocabulary codes, not all hypercube vertices. Candidate recall and reranking agreement are separate measurements. See the repository's technical reports for datasets, position counts, precision, and comparison settings. The Qwen artifact is update 224k/244k; reported results are for that checkpoint, with no assumed completion of the schedule.

Usage

Install the MoB_Head source and dependencies, then:

from mob.hub import load_head
head, codebook, config = load_head("ryan1145/Llama-MoB-3.2-1B-MLP")
# h: final normalized hidden states, shape [positions, hidden_size].
# bit_logits, mixture_logits = head(h)

The loader checks weight/codebook hashes and loads no remote Python code. For independent decoding, use mob.mob_common code decoders. Cascade reranking also requires the matching native output-table rows.

Limitations

The head approximates its teacher and does not preserve its full distribution. Independent long-form generation can drift; cascade decoding is constrained by proposal recall. Speedups depend on hardware, vocabulary, batch and decode path. The project comparison found performance close to CFHead in several settings, but did not surpass its overall acceleration tradeoff.

License and attribution

Artifact license: llama3.2; see LICENSE and NOTICE. Original loading code is MIT-licensed in MoB_Head. Source backbone and dataset terms remain applicable. No backbone parameters or corpus text are included.