[ARCH]Add SiameseNorm and DepthAttention architecture

[franksfc]

Architecture Proposal

#2

Implementation Details

This PR implements SiameseNorm + Depth-Attention on top of the OLMo 3 decoder architecture and includes the reproducible Ascend/MindSpeed training pipeline used for the current 1B artifacts.

Modified components

• Added separate base and siamese_depth OLMo 3 model variants.
• Added forward-local dual residual-stream state for SiameseNorm.
• Added sparse cross-layer Depth-Attention before ordinary token attention.
• Preserved OLMo 3 Q/K normalization, RoPE, MLP, tokenizer, causal masking, and the [SWA, SWA, SWA, Full] attention pattern.
• Added Stage 1 pretraining, Stage 2 mid-training, Stage 3 long-context training, and Stage 4 Think/Instruct SFT configuration and launch paths.
• Added checkpoint save/resume/transition validation, native inference, 65K KV-cache validation, PPL evaluation, and OLMES objective-evaluation entry points.
• Added Transformers remote-code inference files and a non-weight mirror of the published 1B Hugging Face repository under arch/SiameseNorm-DepthAttention.

Key implementation details

• SiameseNorm maintains two forward-local residual streams while sharing each attention and MLP computation. For layer l, the Post-Norm-like stream receives the depth-scaled update Δ / sqrt(2l), while the Pre-Norm-like stream receives the full update Δ.
• Depth-Attention is applied after Q/K normalization and RoPE, but before SWA or Full Attention. The current query selects among the current layer and sparse earlier-layer K/V states at the same token position. The default source stride is 8.
• The mixed value is passed into the original token-attention operation and written to the inference KV cache. Depth-Attention adds no learned parameters or separate persistent cache.
• Stage 3 and Stage 4 apply YaRN only to Full-Attention layers for 65,536-token context. SWA layers retain the original RoPE and a 4,096-token window.
• Base and modified models use separate, strictly validated checkpoint keyspaces.
• Checkpoint save, resume, topology migration, and Stage 1 → Stage 2 → Stage 3 → Stage 4 transitions are included in the pipeline.

Compatibility considerations

• The published 1B artifacts use standard Transformers remote-code format and were release-validated with BF16 SDPA.
• The native training path targets Megatron/MindSpeed on Ascend NPUs.
• Third-party revisions are pinned, while deployment-specific paths, cluster configuration, and credentials remain outside version control.
• Model construction is available for the repository's OLMo 3 1B, 3B, and 7B presets.
• Real end-to-end four-stage training has currently been completed only for the modified 1B model.

Current deviations and scope

• This is a draft implementation PR, not a completed architectural-effectiveness claim.
• The modified 1B four-stage checkpoints have been published.
• A preliminary comparison against the published OLMo 2 1B pure-SFT reference shows a positive aggregate trend, but this is not the final matched OLMo 3 baseline comparison.
• The matched 1B OLMo 3 baseline is still training, so final matched-token loss and downstream-evaluation results are still in progress.
• No completed 3B or larger-scale matched comparison is included yet.
• Proposal #43 specifies 1B/3B/7B experiments, while the repository merge checklist asks for 1B/3B/8B. This scale mismatch should be resolved with reviewers before the final experimental matrix is frozen.

Experimental Validation

Research Question 1 — Does the implementation preserve OLMo 3 training, checkpoint, and long-context inference semantics?

Hypothesis: SiameseNorm and Depth-Attention can be integrated without breaking the four-stage OLMo 3 pipeline, checkpoint transitions, SWA/Full-Attention behavior, or 65K inference.

Results & Analysis:

The modified 1B model completed and published the following artifacts:

Stage	Iteration	Context length	Status
Stage 1 pretraining	89,407	8,192	Published
Stage 2 mid-training	47,684	8,192	Published
Stage 3 long-context	11,921	65,536	Published
Stage 4 Think SFT	43,224	65,536	Published
Stage 4 Instruct SFT	3,252	65,536	Published

Findings: The current evidence supports implementation correctness and end-to-end operability at the 1B scale.

Research Question 2 — Does SiameseNorm + Depth-Attention improve model quality?

Hypothesis: At matched model size, data order, token count, optimizer, learning-rate schedule, batch size, random seed, and evaluation protocol, the modified architecture will improve training loss and downstream evaluation metrics without introducing instability.

Results & Analysis:

Preliminary evaluation shows a positive aggregate result. The current OLMo 3 1B Instruct-SFT checkpoint reaches a macro average of 36.9, compared with 34.2 for the published OLMo 2 1B pure-SFT reference, an improvement of +2.7 points across the eight reported benchmarks.

Benchmark	OLMo 2 1B pure SFT	Ours: OLMo 3 1B Instruct SFT	Difference
BBH	32.8	38.9	+6.1
DROP	33.8	31.3	-2.5
GSM8K	52.1	51.5	-0.6
IFEval	50.5	70.1	+19.6
MATH	13.2	10.0	-3.2
MMLU	36.4	40.6	+4.2
PopQA	12.7	9.6	-3.1
TruthfulQA	42.1	43.2	+1.1
8-task macro average	34.2	36.9	+2.7

The strongest preliminary gains are observed on IFEval (+19.6), BBH (+6.1), and MMLU (+4.2). The modified model improves on four of the eight individual benchmarks and raises the overall macro average despite regressions on several tasks.

This table is an early external-reference comparison, not the final controlled baseline experiment. The matched OLMo 3 1B baseline is still training, and final matched-token loss curves and downstream results will be added when that run finishes.

Findings: Preliminary results indicate a meaningful aggregate improvement, with especially strong gains in instruction following and selected reasoning/knowledge benchmarks. The final effectiveness conclusion remains pending until the matched OLMo 3 baseline and complete evaluation are available.

Research Question 3 — Does the result generalize beyond 1B?

Hypothesis: The combined architecture remains stable and beneficial at larger model scales.

Results & Analysis:

No completed 3B or 7B/8B matched validation is included yet.

Findings: Pending. This PR should remain in draft until the target model scales are reconciled and the required experiments are complete.

Archive

• WandB Report (training and evaluation logs): https://wandb.ai/archspace/SiameseNormDepthAttention/reports/Siamese-Norm-and-Depth-Attention-in-OLMo-3-1B--VmlldzoxNzYwMzAwNw
• HuggingFace Collection (current 1B Stage 1–4 checkpoints): https://huggingface.co/ArchSpace-Collection/SiameseNorm-DepthAttention

---

Reviewer Assessment (for repo reviewers)

Draft status is recommended. The implementation, 1B four-stage artifact path, and preliminary positive evaluation are available for review. The controlled matched-baseline experiment and multi-scale generalization are still in progress and should be completed before the final merge-readiness assessment.

Merge Checklist:

• The implementation PR is linked to an in-progress Architecture Proposal issue (#43).
• The implementation correctness of the proposed architecture is verified at 1B through four-stage training, checkpoint, and inference validation.
• The effectiveness of the proposed architecture is validated at 1B, 3B, and 8B model scales.
• The archive information is complete, including WandB reports and HuggingFace collections of model checkpoints.

[Knlife]

1. Please move the Megatron-LM/ training code to the /reproduce directory.
2. Provide the modeling code in HuggingFace format in the arch/ directory, and it may look like as follows,

ArchSpace
├── README.md
├── reproduce <- dir for reproducible training and evaluation codes
└── archs     <- dir for architectures
    ├── olmo3
    │   ├── configuration_olmo3.py
    │   └── modeling_olmo3.py
    ├── qwen3
    │   ├── configuration_qwen3.py
    │   └── modeling_qwen3.py
    └── olmo3-siamese-depth-attn
        ├── configuration_olmo3_siamese_depth.py
        └── modeling_olmo3_siamese_depth.py