Embedding RMSNorm

Summary: An additional RMSNorm layer applied immediately after the token embedding lookup, before the first transformer block — distinct from the standard pre-attention RMSNorm in each block. Used in Inkling to stabilize early-layer training dynamics at 975B scale.

Overview

Standard transformer blocks apply RMSNorm/LayerNorm before attention and MLP (Pre-Norm) or after (Post-Norm). Inkling adds a third normalization point: right after the embedding layer, normalizing the embedding vectors before they enter the deep stack.

Key Components

Position in Architecture

Input tokens
    ↓
Embedding lookup (token → vector)
    ↓
🔴 Embedding RMSNorm  ← NEW
    ↓
Transformer Block 1 (RMSNorm → Attn → RMSNorm → MLP)
    ↓
...
Transformer Block 66
    ↓
Output projection

Mathematical Form

# Standard embedding
x = Embed(token_ids)  # [B, L, d_model]

# Embedding RMSNorm (new)
x = RMSNorm(x)        # [B, L, d_model]

# Then standard Pre-Norm blocks
for block in blocks:
    x = x + Attention(RMSNorm(x))
    x = x + MLP(RMSNorm(x))

Parameters

• Learnable scale: γ ∈ ℝ^{d_model}
• (Optional) No bias term (RMSNorm convention)

Motivation

Early-Layer Gradient Stability

• At 975B params / 66 layers, embedding gradients must propagate through 66 blocks
• Raw embedding scales can vary wildly (e.g., frequent vs rare tokens)
• Early normalization prevents scale explosion/collapse before first block

Separation of Concerns

• Embedding RMSNorm: Normalizes input distribution to transformer stack
• Block Pre-Norm: Normalizes residual stream for each sublayer
• Different statistics, different purposes

Quantization Prep

• BF16/MXFP8/NVFP4 training benefits from controlled dynamic range at every stage
• Embedding RMSNorm bounds embedding output range before first linear layer

Variants / Extensions

• Learnable per-dimension scale (standard RMSNorm)
• Fixed scale (γ=1), just centering
• Conditional Embedding RMSNorm: Condition on token frequency / position
• Embedding LayerNorm: LayerNorm instead of RMSNorm (adds bias, mean centering)

Applications

• Inkling (975B, 66 layers, multimodal)
• Potentially: any ultra-deep transformer (>50 layers) with large vocab
• Multimodal models where embeddings from different modalities (text, image patches, audio tokens) have different scales

Historical Context

Year	Work	Norm Placement
2017	Transformer	Post-Norm (after sublayers)
2018	BERT	Post-Norm
2019	Pre-LN Transformer	Pre-Norm (before sublayers)
2020	RMSNorm (Zhang & Sennrich)	Pre-Norm, no mean centering
2022	DeepNorm	Scaled residual + Pre-Norm
2023	Various	Sandwich Norm, QK-Norm, etc.
2026	Inkling	Embedding RMSNorm + Block Pre-Norm (dual)

Related Concepts

• inkling-architecture
• rmsnorm
• transformer-architecture
• pre-norm
• deepnorm
• inkling

Sources

• inkling-raschka-blog: Architecture observations #2 — Additional RMSNorm directly after token embeddings, separate from pre-attention RMSNorm; explicit in HF config and Transformers impl