ArchSpace-Collection/SiameseNorm-DepthAttention

OLMo 3 1B — SiameseNorm + Depth-Attention and Matched Baseline

This repository contains the five checkpoints from the four-stage OLMo 3 1B training pipeline for both:

• SiameseNorm + Depth-Attention, using Transformers remote code.
• The matched pure OLMo 3 baseline, using the official Transformers Olmo3ForCausalLM implementation.

The standalone model repositories are linked below. The five modified checkpoints are also grouped in the SiameseNorm-DepthAttention collection.

Checkpoints

Stage	Training sequence length	Model context capacity	SiameseNorm + Depth-Attention	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 artifacts are also mirrored in this repository under:

Variant	Hub subfolders
SiameseNorm + Depth-Attention	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 apply YaRN only to Full-Attention layers. SWA layers retain the original RoPE and a 4,096-token window.

Loading SiameseNorm + Depth-Attention

The modified checkpoints require trust_remote_code=True. SDPA is the release-validated BF16 backend.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "ArchSpace-Collection/SiameseNorm-DepthAttention"
subfolder = "olmo3/1b/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 checkpoints require transformers>=4.57.6,<5 and use the official OLMo 3 implementation without remote code.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "ArchSpace-Collection/SiameseNorm-DepthAttention"
subfolder = "olmo3/1b/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.

Preliminary matched evaluation

The modified model converged faster during the early portion of training, but the final eight-task objective-evaluation averages are nearly tied. The modified Instruct-SFT checkpoint scores 36.9, compared with 36.8 for the matched OLMo 3 baseline.

Benchmark	Matched OLMo 3 1B baseline	SiameseNorm + Depth-Attention	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	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

The 1B result supports an early-convergence benefit but does not yet establish a substantial final downstream-quality improvement. The modified 3B experiment is still training, and the 7B experimental plan is suspended.

Architecture

Both variants share the OLMo 3 1B backbone:

• 16 transformer layers
• hidden size 2,048
• intermediate size 8,192
• 16 query heads and 16 key/value heads
• 128-dimensional attention heads
• [SWA, SWA, SWA, Full] attention pattern
• 4,096-token sliding window
• reordered RMSNorm

The modified variant additionally enables SiameseNorm and sparse cross-layer Depth-Attention. The baseline disables both modifications.

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