mkurman/neuroblast-v4.1-450M

NeuroBLAST v4.1 450M

⚠️ WORK IN PROGRESS — EXPERIMENTAL RELEASE ⚠️

NeuroBLAST v4.1 is the latest iteration of the biologically inspired "cortical" architecture introduced in NeuroBLAST V3. It is based on mkurman/NeuroBLAST-V3-0.6B-SYNTH-EC-144B-TOK, with architecture modifications that reduce the parameter count to ~450M, and has been trained substantially further:

1. Continued pre-training on pleias/synth up to ~200B tokens.
2. Reasoning fine-tuning on a mix of SYNTH-like and R1-like reasoning data:
   • 50% medical
   • 25% chat
   • 15% math
   • 10% agentic

This checkpoint is not fully polished yet — expect rough edges — but it should serve as a solid starting point for further fine-tuning on your dedicated use case.

⚠️ Known Issue: Mixed Precision Training Does Not Work

Train this model in FP32 only.

The weights are so saturated that they barely move when anything other than FP32 precision is used. Attempts to fine-tune with BF16 / FP16 or mixed precision will result in little to no learning. If you plan to fine-tune, use full FP32 (and expect the associated memory/throughput cost). FP32 inference is also recommended.

Model Details

• Architecture: NeuroBLAST v4.1 (custom hybrid, modified from V3)
• Parameters: ~450M
• Vocab size: 65,538
• Base checkpoint: NeuroBLAST V3 0.6B SYNTH EC (144B tokens)
• Total pre-training tokens: ~200B (PleIAs/SYNTH)
• Fine-tuning: SYNTH-like + R1-like reasoning mix (medical / chat / math / agentic)
• Precision: FP32

As with V3, the architecture follows a three-stage cortical design — sensory, associative, and motor layers with hybrid attention/convolution blocks and deep residual bridges — see the V3 model card for the full architectural description. v4.1 introduces modifications on top of this design.

Usage

Note: You must use trust_remote_code=True, as this model uses custom modeling code.

import torch
from transformers import AutoTokenizer, TextStreamer, AutoModelForCausalLM

model_id = "mkurman/neuroblast-v4.1-450M"

tokenizer = AutoTokenizer.from_pretrained(model_id)

# IMPORTANT: load in FP32 — the model does not behave well in lower precision
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float32,
    device_map="cuda",
    trust_remote_code=True,
).eval()

streamer = TextStreamer(
    tokenizer, skip_prompt=False, decode_kwargs={"skip_special_tokens": False}
)

input_ids = tokenizer.apply_chat_template(
    [{"role": "user", "content": "what is hypertension?"}],
    tokenize=True,
    return_tensors="pt",
    add_generation_prompt=True,
)

with torch.no_grad():
    outputs = model.generate(
        input_ids=input_ids.to(model.device),
        max_new_tokens=256,
        streamer=streamer,
        use_cache=True,
        repetition_penalty=1.0,
    )

Intended Use & Limitations

• Best suited as a base for domain-specific fine-tuning (especially medical reasoning), not as a finished chat model.
• Outputs may be unpolished or inconsistent; thorough evaluation is recommended before any downstream use.
• Not suitable for clinical or medical decision-making. The medical data fraction reflects reasoning-style training only.
• Fine-tuning must be done in FP32 (see Known Issue above).

Acknowledgments

• Trained using Cloud TPUs provided by Google's TPU Research Cloud (TRC) program.
• Special thanks to Pierre-Carl Langlais and the PleIAs team for the SYNTH dataset.

Repo

GitHub: https://github.com/mkurman/neuroblast-v3