Cognica-PoE-v1.0-1.3B-stage-ner

Paper: Product of Experts as Scalable Local Learning: Modular Construction at 1.3B Parameters (Jeong, 2026)

Not a production NER model. This release is empirical validation of PoE's post-hoc specialist construction for a structured-output task — paper Section 6.5 extended.

A 164 M-parameter named-entity-recognition SFT specialist stage trained directly on the frozen Cognica-PoE-v1.0-1.3B-base. The stage is a sibling of the chat/math/code/tool/summary stages — all six branch from the same base, not from each other.

The specialist is trained with a natural-language target format ("Entity" is a person.) rather than the conventional BIO tag list. At the 1.3 B base scale, generative NER is only reliable when the training format stays close to the model's native text distribution; structured-tag targets push the autoregressive decoder into out-of-distribution prefixes that fail to generate under greedy decoding.

What this artifact does demonstrate:

1. Natural-language reformulation enables generative NER at 1.3 B. With repetition_penalty >= 1.1 and no_repeat_ngram_size >= 3, the model extracts 2-4 entities per short sentence in the "Entity" is a [category]. pattern.
2. Assistant-only val bpb 1.97 on a 512-example held-out split across CoNLL-2003, FewNERD, and WikiNeural, converging monotonically from 4.68 over 1 329 steps.
3. Additive composition at inference. The cascade loader folds ancestor stages' lm_head_stage into the effective lm_head_base when chaining, so this stage composes with the other sibling specialists (chat, math, code, tool, summary) via logits = lm_head_base(x) + Sigma_k lm_head_stage_k(x).

Quick start

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "cognica/Cognica-PoE-v1.0-1.3B-stage-ner"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.eval()

BOS, USR_S, USR_E, ASS_S = 32759, 32760, 32761, 32762
text = "Barack Obama was born in Hawaii and served as the 44th President of the United States."
prompt = f"{text}\n\nList the entities."
ids = [BOS, USR_S] + tokenizer.encode(prompt) + [USR_E, ASS_S]
input_ids = torch.tensor([ids], device=model.device)

with torch.no_grad():
    out = model.generate(
        input_ids,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.3,
        top_k=50,
        repetition_penalty=1.3,
        no_repeat_ngram_size=3,
        pad_token_id=BOS,
    )
print(tokenizer.decode(out[0, len(ids):]))
# '"Hawaii" is a location. "United States" is a location. ...'

Decoding recommendations:

• repetition_penalty=1.3 and no_repeat_ngram_size=3 are both required. Without them, the decoder loops on a single entity.
• Keep max_new_tokens modest (40-80). The first few entities are the reliable ones; later tokens drift.
• Do not use greedy decoding — always sample with temperature in [0.2, 0.5].

Architecture

Component	Detail
Parent	cognica/Cognica-PoE-v1.0-1.3B-base (PoE alpha=0.0, d24, step 26430, val bpb 0.7209)
New transformer layers	4 appended at positions 24-27 (d24 -> d28)
Frozen layers	24 (all base layers)
Dual-head	Yes — additive specialist lm_head_stage (shape 32768 x 1536, zero-init at training start)
Final projection	logits = lm_head_base(x) + lm_head_stage(x)
Total params	1,547,699,986 (~1.55 B)
Trainable params at training	163,577,912 (~164 M, 10.6 %)
Shipped delta	28 tensors, 213,909,560 params, 408 MB (bf16 safetensors)
VE pattern	Preserved from base — 12 value-embeds at layers [1, 3, ..., 23]; new layers carry no VE

Training

Objective	Cross-entropy over assistant turns only (uniform weighting, no label smoothing)
Target format	Natural-language sentences: "Entity" is a [PERSON|ORGANIZATION|LOCATION|MISC category].
User prompt	{text}\n\nList the entities.
Tag schema	PERSON, ORGANIZATION, LOCATION, MISC (coarse 4-class, collapsed from source datasets)
Data	tomaarsen/conll2003 x 20 epochs + DFKI-SLT/few-nerd supervised x 5 epochs + Babelscape/wikineural train_en x 6 epochs -> 1,495,463 train convs (+ 512 held-out val)
Case augmentation	-> 1,505,623 conversations
Sequence length	2,048
Per-GPU batch	8 x 2,048
World size	4 (1 node x 4 x A100 80 GB)
Total batch size	65,536 tokens/step
Steps	1,329
Optimizer	MuonAdamW with per-group LR scaling
Matrix LR	3.0 x 10^-4
lm_head_stage LR	1.0 x 10^-4, weight decay 0.1
Init LR fraction	0.2
Warmup / warmdown	5 % / 50 %
Eval / save cadence	every 25 / 100 steps
Shipped checkpoint	step 1,329, val bpb 1.9722 (best-and-final)

Validation bpb trajectory (selected)

Step	val bpb
25	4.6791
100	4.3801
250	3.8106
500	3.1899
750	2.7118
1000	2.1835
1200	2.0034
1325	1.9722

Monotonic decrease throughout training — no mode collapse, no overfitting spike, best checkpoint is also the final step.

Stacking further stages

base_model_name_or_path supports chaining. Point a new stage repo's config at this repo and the cascade loader will resolve base -> stage-ner -> new stage transparently. The loader folds each ancestor's lm_head_stage into the effective lm_head_base at load time, so all specialist heads compose additively into the final projection (logits = lm_head_base + Sigma_k lm_head_stage_k). See the paper for the formal account of this construction.

Files

File	Purpose
config.json	Model + stage config
delta.safetensors	28-tensor stage delta (bf16, 408 MB)
modeling_cognica_poe.py	Cascade loader + _GPT with dual-head forward
configuration_cognica_poe.py	CognicaPoEConfig with stage fields
tokenization_cognica_poe.py	Byte-level tokenizer
tokenizer.pkl, tokenizer_config.json, special_tokens_map.json, token_bytes.pt	Tokenizer assets
convert_stage_delta.py	Converts a nanochat save_stage_delta .pt into delta.safetensors

Limitations — explicit list

• Research preview at 1.3 B. Entity detection is more reliable than tag classification. Expect frequent ORG<->MISC and LOC<->PERSON confusions.
• Generative NER, not sequence tagging. Output is "Entity" is a category. sentences, not token-labeled spans. Not a drop-in replacement for encoder-based NER (e.g. spaCy, CoNLL-trained BERT).
• Greedy decoding loops. repetition_penalty>=1.1 and no_repeat_ngram_size>=3 are mandatory.
• Later-token drift. First 1-4 entities are reliable; extended generation degrades. Use max_new_tokens <= 80.
• No RLHF / preference tuning / safety tuning.
• Assistant-only loss — validation bpb is measured on assistant tokens only; do not compare to the base's full-text train_val_bpb = 0.7209.

Citation

@article{jeong2026poe,
  title  = {Product of Experts as Scalable Local Learning: Modular Construction at 1.3B Parameters},
  author = {Jeong, Jaepil},
  year   = {2026},
  institution = {Cognica, Inc.},
  doi    = {10.5281/zenodo.19547653},
  url    = {https://doi.org/10.5281/zenodo.19547653}
}

@misc{cognica-poe-stage-ner-2026,
  title  = {Cognica-PoE-v1.0-1.3B-stage-ner: NER dual-head specialist (4-layer) over a PoE base (research preview)},
  author = {{Cognica, Inc.}},
  year   = {2026},
  howpublished = {\url{https://huggingface.co/cognica/Cognica-PoE-v1.0-1.3B-stage-ner}}
}

License

Apache 2.0 — see LICENSE and NOTICE. Same terms as the base model. Training datasets (CoNLL-2003, FewNERD, WikiNeural) each carry their own licenses and are acknowledged in NOTICE.