W-1.2 Architecture Specification

Executive summary

W-1.2 is Jhelum Labs’ planned second-generation experimental language model. It scales the laboratory’s W-1.1 work from approximately 124 million to 353 million trainable parameters while preserving a transparent, decoder-only causal Transformer design. The model is intended for research, controlled public experimentation, and future instruction-tuning work—not for high-stakes autonomous decisions.

Design objectives

• Improve language modeling quality over W-1.1 through greater depth, width and training data.
• Support a 4,096-token context window without changing the causal next-token objective.
• Keep inference practical through grouped-query attention and tied embeddings.
• Train on 8 billion tokens from a quality-controlled, predominantly human-written corpus.
• Preserve reproducibility through explicit configuration, checkpoints, manifests and evaluation artifacts.

At-a-glance specification

1. Model architecture

W-1.2 maps discrete token identifiers into 1,024-dimensional vectors, processes those vectors through 24 identical causal decoder blocks, applies a final RMSNorm, and projects the hidden states back to the 50,000-token vocabulary. The language-model head shares its weights with the token embedding matrix.

1.1 Top-level computation

• Input tokens → token embedding matrix.
• Twenty-four pre-normalized causal Transformer blocks.
• Final RMSNorm → tied vocabulary projection.
• Cross-entropy loss on the next token, with future tokens blocked by a causal mask.

1.2 Decoder block

Figure 1. Pre-normalized W-1.2 decoder block with two residual paths.

Each block applies RMSNorm before the attention and MLP sublayers. Residual connections provide short gradient paths through the full network. All linear projections use no bias unless later experiments demonstrate a clear benefit.

1.3 Parameter accounting

2. Causal grouped-query attention

Attention lets each token combine information from relevant earlier positions. W-1.2 uses 16 query heads but only 8 key/value heads. Two query heads therefore share each key/value head, reducing the key-value cache relative to full multi-head attention while retaining multiple query subspaces.

Figure 2. Original schematic of W-1.2 causal grouped-query attention.

2.1 Attention equation

Attention(Q, K, V) = softmax((QKᵀ / √dₕ) + M) V

Here, dₕ = 64 and M is the upper-triangular causal mask. Masked future positions receive negative infinity before softmax, giving them zero attention probability. A token can use its own current representation and all preceding tokens, but it cannot read any token that appears later in the sequence.

2.2 Rotary position encoding

RoPE rotates pairs of query and key features by position-dependent angles. This introduces relative positional structure directly into the attention calculation without adding a separate learned position vector to the token embedding.

3. Feed-forward and normalization design

3.1 SwiGLU MLP

SwiGLU(x) = Wdown[SiLU(Wgate x) ⊙ (Wup x)]

Each block expands the 1,024-dimensional residual stream to two 3,072-dimensional branches. One branch supplies a SiLU gate, the other supplies content, and their elementwise product is projected back to the model width.

3.2 RMSNorm

RMSNorm rescales hidden states using their root-mean-square magnitude and a learned per-channel scale. Pre-normalization is selected to improve optimization stability across the 24-layer stack.

4. Tokenization and data design

The W-1.2 tokenizer will be trained from the final cleaned corpus rather than inherited from W-1.1. A 50,000-entry BPE vocabulary is selected to improve coverage of natural language, scientific terminology and code while keeping the embedding table within the target parameter budget.

4.1 Corpus target

4.2 Quality controls

• License and provenance record for every source.
• Exact and near-duplicate removal before split creation.
• Language, encoding, corruption, boilerplate and low-information filters.
• Removal or masking of exposed personal data and obvious secrets.
• Document-level split assignment to prevent train/validation leakage.
• Tokenizer fertility, unknown-token behavior and domain coverage tests.

Figure 3. Reproducible data preparation, pretraining and evaluation pipeline.

5. Training configuration

Pretraining uses standard causal language modeling: for every sequence, input token t predicts token t+1. The preferred global token batch is 1,048,576 tokens per optimizer step. At that batch size, 8 billion training tokens correspond to approximately 7,630 optimizer steps.

Figure 4. Parameter allocation and nominal 8B-token training progression.

5.1 Suitable GPUs

• NVIDIA B200
• NVIDIA H200
• NVIDIA H100 80GB
• NVIDIA A100 80GB

The final selection should be validated using a short throughput and memory benchmark with the exact sequence length, batch size, attention implementation and optimizer configuration.

6. Checkpoint and recovery standard

A recoverable checkpoint must contain enough state to continue the same optimization trajectory after interruption. Weight-only snapshots are useful for inference but are not sufficient for exact training resumption.

• Model weights and architecture configuration.
• Optimizer state and learning-rate scheduler state.
• Current optimizer step and total tokens processed.
• Gradient-scaler state when applicable.
• Random-number-generator states and data-stream position.
• Validation history, best-checkpoint pointer and integrity hashes.

7. Evaluation plan

Evaluation must distinguish intrinsic language-model quality from downstream capability. Results should include confidence intervals, exact task versions, prompting settings, tokenizer and checkpoint hashes.

8. Post-pretraining phases

• Supervised instruction tuning on audited, high-quality instruction/response examples.
• Preference optimization using carefully reviewed chosen/rejected response pairs.
• Safety tuning and refusal calibration without destroying general helpfulness.
• Regression evaluation against the frozen pretrained and instruction-tuned baselines.

9. Risks and design controls

10. Release artifact layout

The recommended public package separates source code, configuration, tokenizer, evaluation, charts and model weights. Large model files may be distributed as release assets while the repository retains checksums and download instructions.

w-1.2/
├── config.json
├── tokenizer.json
├── model_card.md
├── src/
├── evaluation/
├── charts/
└── weights/
├── w-1.2.safetensors
└── SHA256SUMS

11. Acceptance criteria

• Parameter count matches the canonical configuration exactly.
• Tokenizer vocabulary, special-token IDs and model configuration agree.
• All 8B planned tokens are accounted for, or the release clearly states the actual count.
• No train/validation/test overlap is detected by the final audit.
• Final and best checkpoints load and generate deterministic smoke-test outputs.
• Evaluation results and limitations are published with reproducible settings.
• Weights, tokenizer, configuration and documentation pass SHA-256 verification.

References

1. Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
2. Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202.
3. Su, J. et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
4. Zhang, B. & Sennrich, R. (2019). Root Mean Square Layer Normalization. arXiv:1910.07467.
5. Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.
6. Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models. arXiv:2203.15556.

Document control