ArchSpace-Collection/SiameseNorm-DepthAttention

OLMo 3 1B and 3B — SiameseNorm + DepthAttention and Matched Baselines

This repository mirrors the four-stage OLMo 3 training pipelines for both the SiameseNorm + DepthAttention architecture and matched pure OLMo 3 baselines. Each model size contains five checkpoints: Stage 1, Stage 2, Stage 3, Stage 4 Think SFT, and Stage 4 Instruct SFT.

The individual repositories are grouped in the SiameseNorm-DepthAttention collection.

OLMo 3 3B checkpoints

Stage	Training sequence	Model context	SiameseNorm + DepthAttention	Matched baseline
Stage 1 pretraining	8,192	8,192	stage1	baseline-stage1
Stage 2 mid-training	8,192	8,192	stage2	baseline-stage2
Stage 3 long-context training	65,536	65,536	stage3	baseline-stage3
Stage 4 Think SFT	32,768	65,536	stage4-think	baseline-stage4-think
Stage 4 Instruct SFT	32,768	65,536	stage4-instruct	baseline-stage4-instruct

The same 3B artifacts are mirrored here under:

Variant	Hub subfolders
SiameseNorm + DepthAttention	olmo3/3b/stage1, stage2, stage3, stage4/think, stage4/instruct
Matched baseline	olmo3/3b/baseline/stage1, stage2, stage3, stage4/think, stage4/instruct

OLMo 3 1B checkpoints

Stage	Training sequence	Model context	SiameseNorm + DepthAttention	Matched baseline
Stage 1 pretraining	8,192	8,192	stage1	baseline-stage1
Stage 2 mid-training	8,192	8,192	stage2	baseline-stage2
Stage 3 long-context training	65,536	65,536	stage3	baseline-stage3
Stage 4 Think SFT	32,768	65,536	stage4-think	baseline-stage4-think
Stage 4 Instruct SFT	32,768	65,536	stage4-instruct	baseline-stage4-instruct

The 1B artifacts retain their original layout:

Variant	Hub subfolders
SiameseNorm + DepthAttention	olmo3/1b/stage1, stage2, stage3, stage4/think, stage4/instruct
Matched baseline	olmo3/1b/baseline/stage1, stage2, stage3, stage4/think, stage4/instruct

Stage 3 and Stage 4 use the frozen 65,536-token configuration. YaRN applies to Full-Attention layers; SWA layers retain the original RoPE and a 4,096-token window.

Loading SiameseNorm + DepthAttention

The modified checkpoints require transformers>=4.57.6,<5 and trust_remote_code=True.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "ArchSpace-Collection/SiameseNorm-DepthAttention"
subfolder = "olmo3/3b/stage4/instruct"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    subfolder=subfolder,
    trust_remote_code=True,
    fix_mistral_regex=False,
)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    subfolder=subfolder,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
)

Loading the matched baseline

The baseline uses Transformers' official Olmo3ForCausalLM implementation and does not require remote code.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "ArchSpace-Collection/SiameseNorm-DepthAttention"
subfolder = "olmo3/3b/baseline/stage4/instruct"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    subfolder=subfolder,
    fix_mistral_regex=False,
)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    subfolder=subfolder,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
)

fix_mistral_regex=False is intentional and preserves the tokenizer behavior used during training.

Matched 3B Instruct SFT evaluation

On the final eight-task objective evaluation, SiameseNorm + DepthAttention improves the matched OLMo 3 3B baseline average from 48.23 to 51.57.

Benchmark	Matched OLMo 3 3B baseline	SiameseNorm + DepthAttention	Difference
BBH	52.93	56.72	+3.79
DROP	47.84	53.38	+5.54
GSM8K	65.58	69.29	+3.71
IFEval (loose)	72.83	77.08	+4.25
MATH (flex)	35.67	36.80	+1.12
MMLU	57.17	58.94	+1.77
PopQA	14.53	14.57	+0.04
TruthfulQA	39.28	45.83	+6.55
8-task macro average	48.23	51.57	+3.35

Preliminary matched 1B evaluation

The final eight-task 1B objective-evaluation averages are nearly tied: 36.9 for SiameseNorm + DepthAttention and 36.8 for the matched baseline.

Benchmark	Matched OLMo 3 1B baseline	SiameseNorm + DepthAttention	Difference
BBH	37.3	38.9	+1.6
DROP	36.2	31.3	-4.9
GSM8K	53.5	51.5	-2.0
IFEval (loose)	63.6	70.1	+6.5
MATH (strict)	8.0	10.0	+2.0
MMLU	44.8	40.6	-4.2
PopQA	11.3	9.6	-1.7
TruthfulQA	40.0	43.2	+3.2
8-task macro average	36.8	36.9	+0.1

Architectures

Size	Layers	Hidden	Intermediate	Q/KV heads	Head dimension
1B	16	2,048	8,192	16 / 16	128
3B	16	3,328	13,312	16 / 16	208

Both sizes use the [SWA, SWA, SWA, Full] attention pattern, a 4,096-token sliding window, and reordered RMSNorm. The modified variant additionally enables SiameseNorm and sparse cross-layer DepthAttention; the matched baseline disables both modifications.

These Hugging Face artifacts are intended for inference and generation. Exact continuation of the native distributed training objective should use the MindSpeed/Megatron training pipeline.