BananaMind/BananaMind-2.1-Pico-Preview

[image]

BananaMind-2.1-Pico-Preview

BananaMind-2.1-Pico-Preview is a 1,480,516-parameter decoder-only base language model trained from scratch by BananaMind on 2B tokens of FineWeb-Edu. It is the first public checkpoint of the BananaMind 2.1 architecture line and is released as a preview: the architecture, not the checkpoint, is the point of the release.

The model has three physical Transformer blocks but executes four blocks per forward pass — block 2 runs twice (L1 → L2 → L2 → L3). Alongside the loop it carries three architectural additions over BananaMind 2.0: value-subtraction cross-head attention (XSA), a detached embedding-refresh gate, and a hashed causal trigram embedding injected mid-stack.

This is a base model, not an instruction-tuned or chat model. Use continuation-style prompts and load the repository with trust_remote_code=True.

Release Status

The execution schedule is set by loop_mode and loop_passes in config.json, not hardcoded in the modeling file. The released checkpoint was trained and evaluated under partial with 2 passes.

Model Details

Architecture Overview

Tokens pass through a tied input embedding scaled by sqrt(hidden_size), then through the looped block stack, a final RMSNorm, and the tied language-model head. Each block is pre-normalized and contains causal grouped-query self-attention, the embedding-refresh gate, and a SwiGLU feed-forward network, with residual connections around each sublayer.

Partial looping. Only block 2 is reused. The first and last blocks run once each, so the model pays for three blocks of parameters and four blocks of compute. This sits between a fully looped stack (every block twice) and a plain three-block stack; the ablation below shows partial looping beating both.

Value-subtraction XSA. After standard grouped-query attention, each head's output has the component parallel to that head's current-token value vector projected out. Because four query heads share one KV head, the subtraction is applied per KV group across its four query heads, removing the self-value direction that otherwise dominates short-context attention output.

Embedding-refresh gate. Each block re-injects the original token embedding through a gated path. The gating signal is the block's own attention output, detached from the gradient graph, combined with a strictly causal depthwise convolution over that signal (left-padded, kernel size 9). The gate multiplies a projection of the original embedding, and the result is scaled by a zero-initialized scalar alpha, so the path starts as a no-op and the model learns how much lexical identity to restore at each depth.

Hashed causal trigram embedding. A 3,906-bucket table is indexed by an integer hash of the current token and the two preceding tokens, using only current and past positions. It is added once, before block 2, in every loop pass. At this scale the table is a third of all parameters — it is reported separately in parameter_summary() for that reason.

Parameter Breakdown

Per block: attention 41,024, refresh gate 50,689, SwiGLU 147,456, norms 256.

Evaluation

[image]

ARC Easy, ARC Challenge, PIQA, and HellaSwag use zero-shot acc_norm,none. ArithMark 3 uses length-normalized continuation accuracy. Base Bench is the public 350-item BananaMindBench 1.1 suite.

BananaMind-2-Micro is the closest architectural relative — same 128-wide hidden size, same 2,048-token vocabulary, same XSA refresh path, same Muon/AdamW split — at nine physical layers rather than three. Pico-Preview is roughly half the parameters trained on 37.5× fewer tokens, and still leads it on Base Bench Elo and weighted accuracy while trailing on the four-choice academic benchmarks and INT Index.

The INT Index chance-normalizes HellaSwag, the mean of ARC Easy and ARC Challenge, PIQA, and ArithMark 3:

N(s, c) = 100 * (s - c) / (100 - c)
INT = [N(H,25) + N((ARC_E+ARC_C)/2,25) + N(P,50) + 0.65*N(A3,25)] / 3.65

Compute Efficiency

Training compute is estimated consistently as 6 × parameters × training tokens.

Pico-Preview reaches within 0.65 INT of Micro for 1.3% of the training compute, which is the result the architecture is meant to demonstrate. Compute here uses the full 1,480,516-parameter count; against models with no n-gram table, the 980,547 core count gives 11.77 PFLOPs and 0.456 INT per PFLOP.

All values are self-reported evaluations and can vary with harness version, tokenizer handling, dtype, and scoring configuration.

Loop Mode Ablation

The same three physical blocks were trained under three execution schedules, holding parameters, data, and step count fixed. partial reuses block 2 once (4 executions), all runs the whole stack twice (6 executions), none runs each block once (3 executions).

Full looping reached a perplexity roughly 2 lower than partial looping early in training but scored worse on every downstream benchmark, and partial looping led full looping by 1.41 INT Index at the end. Adding a third pass over block 2 (L1 → L2 → L2 → L2 → L3) scored below two passes, so the returns from looping the middle block turn negative quickly.

All three modes are reachable from the released config by setting loop_mode and loop_passes.

Training Data

The full 2B-token budget is FineWeb-Edu, streamed raw with no curriculum, no reweighting, and no phase schedule.

Training Setup

The two-optimizer split follows standard Muon practice: Muon updates the 2D hidden-layer weight matrices, while the tied token embedding, the trigram table, RMSNorm weights, and the scalar parameters (alpha, the n-gram scale) are handled by AdamW.

Tokenizer

BananaMind-2.1-Pico-Preview uses a 2,048-token vocabulary sized for the 1M-parameter class — at this width a full 32K vocabulary would consume more parameters than the entire block stack. The same tokenizer as BananaMind-2-Micro.

Usage

Install the runtime dependencies:

pip install -U torch transformers safetensors

Load the custom architecture with remote code enabled:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "BananaMind/BananaMind-2.1-Pico-Preview"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = (
    torch.bfloat16
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    else torch.float32
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=dtype,
).to(device).eval()

prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        use_cache=False,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

For deterministic continuation scoring, use do_sample=False. For free-form sampling, a temperature from 0.6 to 0.8, top_p=0.9, and repetition_penalty=1.1 are reasonable starting points.

Switching Loop Modes

The execution schedule is a config field, so any of the three ablated modes can be run against the released weights:

model.model.config.loop_mode = "none"     # "all", "partial", or "none"
model.model.config.loop_passes = 2

Weights are shared across loop passes, so none and all load without modification — but the checkpoint was trained under partial, and the other modes will score below the numbers above.

Intended Use

BananaMind-2.1-Pico-Preview is intended for architecture research at the sub-2M scale: studying weight-sharing and looped-depth trade-offs, XSA and refresh-gate ablations, n-gram-augmented embeddings, and small-vocabulary training dynamics.

It is a research artifact. At 1.48M parameters trained on 2B tokens it is not a useful text generator.

Limitations

• —No KV cache. The model class forces use_cache=False, so generation recomputes the full prefix at every step. Generation cost is quadratic in sequence length and slow for long outputs.
• —This is a base model and does not follow instructions.
• —The 2,048-token vocabulary produces long token sequences and coarse subword segmentation compared with standard 32K vocabularies.
• —The 3,072-token context window limits long-document use.
• —At 1.48M parameters, factual recall and multi-step reasoning are minimal; HellaSwag and ARC Challenge sit near chance.
• —A third of the parameters live in the trigram table, so parameter-matched comparisons against models without n-gram embeddings should use the 980,547 core count.
• —The training data is English-only FineWeb-Edu, so no other language or domain is characterized.
• —The model has received no safety alignment and can produce inaccurate, biased, repetitive, or undesirable text.
• —Loading requires repository-provided custom Transformers code with trust_remote_code=True.

Do not rely on the model for medical, legal, financial, safety-critical, or other high-stakes decisions.

License

Released under the Apache License 2.0.