Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Transformer Block Assembly › Numerical Stability Considerations

Michael BrenndoerferUpdated June 13, 20252 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 110Transformer Block Assembly › Numerical Stability Considerations

Reading section 12 of 16 · Numerical Stability Considerations · 2 minSection 12 of 16 · 2 min

Section 12 of 16 · 2 min read

Transformer Block Assembly › Numerical Stability ConsiderationsOne section · Sections in this chapter12 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Deep transformers can encounter numerical issues. Several practices help maintain stability:

Mixed precision normalization: Even when using float16 for most computation, normalization should use float32 to prevent underflow when variance is small:

In[38]:

Code

class StableRMSNorm:
    """RMSNorm with float32 computation for stability."""

    def __init__(self, dim, eps=1e-6):
        self.eps = eps
        self.weight = np.ones(dim)

    def __call__(self, x):
        # Store input dtype
        input_dtype = x.dtype

        # Compute in float32
        x_f32 = x.astype(np.float32)
        rms = np.sqrt(np.mean(x_f32**2, axis=-1, keepdims=True) + self.eps)
        x_norm = x_f32 / rms

        # Apply weight and cast back
        return (self.weight * x_norm).astype(input_dtype)

Residual scaling: For very deep networks, scaling residual contributions prevents representation explosion. The intuition is that each residual add accumulates magnitude. After $L$ layers without scaling, the representation magnitude can grow as $O(\sqrt{L})$ or worse. A common fix is to scale each residual contribution:

\mathbf{y} = \mathbf{x} + \alpha \cdot \Delta

where:

$\mathbf{x}$ : the input (residual path)
$\Delta$ : the sublayer output
$\alpha$ : a scaling factor, typically $\alpha = 1/\sqrt{2L}$ for a model with $L$ layers

In[39]:

Code

def scaled_residual_add(x, residual, scale=1.0):
    """
    Add residual with optional scaling.

    For deep networks (100+ layers), use scale = 1/sqrt(2*n_layers)
    to prevent representation magnitude from growing unboundedly.
    """
    return x + scale * residual

Attention score capping: Large attention logits before softmax can cause numerical issues. Some implementations cap logits:

In[40]:

Code

def stable_attention(Q, K, V, mask=None, max_logit=100.0):
    """Attention with logit capping for stability."""
    d_k = Q.shape[-1]
    scores = np.matmul(Q, K.transpose(-2, -1)) / np.sqrt(d_k)

    # Cap extreme values
    scores = np.clip(scores, -max_logit, max_logit)

    if mask is not None:
        scores = scores + mask

    attn = softmax(scores, axis=-1)
    return np.matmul(attn, V)