---
title: TWEO Framework for Stable FP8 Training
url: https://www.emergentmind.com/topics/tweo-framework
type: topic
---

# TWEO Framework for Stable FP8 Training

Transformers Without Extreme Outliers (TWEO) is a framework addressing the problem of catastrophic activation outliers during low-precision training and quantization of transformer models. By introducing a non-invasive, data-independent regularizer targeting the mechanical origin of such outliers, TWEO enables stable full-model FP8 training for large language and vision transformers and supports fast, hardware-friendly quantization to 8-bit precision—all without architectural changes, complex mixed-precision engineering, or invasive algorithmic modifications [2511.23225].

## 1. Motivation and Problem Context

Transformer activations are known to produce massive outliers (|activation| ≫ 1000), a ubiquitous phenomenon in both large language models (LLMs) and vision transformers (ViTs). Native FP8 support (e.g., E4M3 format, dynamic range ≈ ±448) in modern hardware, while promising increased training throughput, is rendered unreliable by these outliers. The key failure modes are (i) post-training quantization (PTQ) collapse due to unusably large quantization ranges, and (ii) catastrophic failure of full FP8 training arising from accumulator overflow and exploded loss. Existing remedies—extensive mixed-precision engineering (restricting FP8 to certain modules), architectural hacks premised on token statistics, and invasive activation modifications—either misdiagnose outlier genesis, sacrifice generality, or increase implementation complexity.

## 2. Theoretical Origin of Activation Outliers

TWEO identifies the root cause of extreme activations as a mechanical artifact of weight-matrix structure in transformer MLP blocks, specifically strong colinearity between up- and down-projection layers, rather than data characteristics or special token frequencies. For a simplified linear MLP,

$$
y = B A x, \quad y \in \mathbb{R}^{d_1},\ A \in \mathbb{R}^{d_2 \times d_1},\ B \in \mathbb{R}^{d_1 \times d_2},\ x \in \mathbb{R}^{d_1}
$$

If $A$ is given by its SVD $A = \sum_{i=1}^r s_i u_i v_i^\top$, the $k$-th output is

$$
y_k = w^\top A x = \sum_{i=1}^r s_i (w^\top u_i) (v_i^\top x)
$$

where $w^\top$ is the $k$-th row of $B$. If both $w$ and $x$ align strongly with the same singular vectors, $y_k$ becomes arbitrarily large, even with random inputs—demonstrating data-independence. This is empirically validated, e.g., in a ViT-B model, giving predicted and observed maximum values within 0.5% relative error.

## 3. TWEO Loss Regularizer

To mitigate this issue, the TWEO framework introduces a simple regularizer that directly penalizes large post-residual activations after each transformer block. Given activations $A^{(l)}$ for block $l$, the total loss is

$$
\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \lambda(t) \mathcal{L}_{\text{TWEO}}
$$

where

$$
\mathcal{L}_{\text{TWEO}} = \frac{1}{L} \sum_{l=1}^L \mathbb{E}\left[ \left| \frac{A^{(l)}}{\tau + \epsilon} \right|^p \right]
$$

Parameter settings are: $\tau = 3$ (soft-threshold), $p = 4$ (penalty power), $\epsilon = 10^{-6}$ (stability), and weighting factor $\lambda(t) = 0.01$ (optionally cosine-annealed). This loss is negligible when activations are small but imposes a steep penalty for values exceeding $\tau$, eliminating extreme outlier tails without discontinuities or architectural coupling.

## 4. Training Algorithm and Implementation

TWEO integrates seamlessly into standard transformer training, requiring only the addition of the $\mathcal{L}_{\text{TWEO}}$ term at each block. All layers—including embeddings, LayerNorm, and output heads—are run under FP8 autocast with NVIDIA Transformer Engine (E4M3/E5M2 format, DelayedScaling, amax_history_len=16). Optimization (AdamW), learning rate schedules, and gradient handling remain as in the baseline. No further clipped-gradient, quantization tricks, or architectural tweaks are utilized.

Key pseudocode logic:

```python
for t in 1…num_steps:
    for each micro-batch:
        x = input batch
        activations = []
        h = x
        for l in 1…L:
            h_norm = LayerNorm_l(h)
            h_mlp = MLP_l(h_norm)  # in FP8 autocast
            h = h + h_mlp
            activations.append(h)
        L_task = compute_task_loss(h, labels)
        L_tweo = (1/L) * sum_over_l mean( |activations[l]|/(τ+ε) )^p
        λ = λ_schedule(t)  # e.g., 0.01→0
        L_total = L_task + λ * L_tweo
        L_total.backward()
        optimizer.step()
        scheduler.step()
```

## 5. FP8 Pre-Training Results

TWEO enables stable full FP8 pre-training across GPT-2 model scales (124M–7B) on OpenWebText, whereas standard FP8 baselines collapse. Training throughput improves by 36% compared to BF16, as no modules revert to higher-precision compute.

**Table: FP8 Pre-training Results (OpenWebText)**

| Model         | BF16 PPL | FP8 Baseline PPL | FP8+TWEO PPL | Peak Outlier |
|---------------|----------|------------------|--------------|--------------|
| GPT-2-124M    | 20.04    | 169.81 (collapse)| 19.26        | 17           |
| GPT-2-350M    | 16.77    | collapse         | 15.64        | 19           |
| GPT-2-774M    | 14.78    | collapse         | 13.89        | 18           |
| GPT-2-1.6B    | 13.84    | collapse         | 12.58        | 19           |
| GPT-2-3.0B    | –        | collapse         | 12.24        | 18           |
| GPT-2-7.0B    | –        | collapse         | 12.02        | 20           |

TWEO regularization clamps outlier magnitudes below 20 in all layers, and training/validation curves match the BF16 baseline.

## 6. Quantization: W8A8 and Residual Stream

TWEO suppresses activation extremes, making previously intractable static per-tensor 8-bit (W8A8) quantization practical. Symmetric AbsMax quantization is used, supporting per-tensor (T), per-channel (C, for weights), and per-token (K, for activations) strategies:

- $Q_b = 2^{b-1} - 1$
- $s = \max(|X|) + \epsilon$
- $X_q = \mathrm{round}(X/s \cdot Q_b)$
- $X_{deq} = X_q/Q_b \cdot s$

**Table: PTQ with AbsMax Static Quantization (GPT-2 Family)**

| Model / Method        | BF16 PPL | W8(T)A8(T) / W8(T)A8(K) | W8(C)A8(T) / W8(C)A8(K) |
|----------------------|----------|-------------------------|-------------------------|
| GPT-2-124M (default) | 20.04    | 86.60 / 22.02           | 82.94 / 21.49           |
| GPT-2-124M (+TWEO)   | 18.83    | 20.82 / 19.51           | 20.54 / 19.00           |
| GPT-2-350M (default) | 16.77    | 1451.4 / 19.76          | 1456.97 / 19.95         |
| GPT-2-350M (+TWEO)   | 15.18    | 16.50 / 15.53           | 16.32 / 15.40           |
| ...                  | ...      | ...                     | ...                     |

TWEO-trained models match or outperform BF16 PPL with even the fastest per-tensor PTQ.

Quantizing the residual stream ($x$) becomes feasible for the first time, outperforming prior "difficulty transfer" approaches (e.g., SmoothQuant, AWQ) that typically cannot quantize $x$ without collapse.

**Table: Residual-Stream Quantization (GPT-2-XL)**

| Method                | Residual Quant? | W8(C)A8(T) / W8(C)A8(K) |
|-----------------------|-----------------|-------------------------|
| SmoothQuant (default) | no              | 14.81 / 14.01           |
| SmoothQuant (default) | yes             | 1876.70 / 21.93         |
| AbsMax on TWEO        | yes             | 13.06 / 12.63           |
| SmoothQuant on TWEO   | yes             | 12.89 / 12.51           |

Vision models (e.g., ViT-B) also benefit: under W6A6 quantization, top-1 accuracy recovers to 66.4%, up from 7.4% baseline.

## 7. Ablations, Sensitivity, and Deployment

Extensive ablations show:

- Varying $\tau \in [2,5]$ and $p \in \{2,4,6\}$ consistently yields strong outlier suppression.
- All blocks' activation tails are equally suppressed; no per-layer tuning required.
- The regularizer scales across model depth (12–48 layers) and width (124M–7B).
- $\lambda(t)$ can be fixed ($0.01$) or cosine-annealed with similar results.

**Deployment recommendations:**

1. Insert the $\mathcal{L}_{\text{TWEO}}$ term after each transformer block with $\tau=3$, $p=4$, $\epsilon=10^{-6}$.
2. Use $\lambda(t)=0.01$ (optionally cosine-annealed to $0$).
3. Run the entire model—including embeddings, LayerNorm, and heads—under FP8 autocast.
4. No additional engineering, architectural changes, or quantization-specific tricks are necessary.

**Limitations and Potential Directions:**

- Scalability beyond 7B parameters remains untested.
- Fine-tuning legacy BF16 checkpoints for outlier removal is uninvestigated.
- Extending the approach to enable stable FP4 training and inference is a future prospect.

TWEO offers a theoretically justified, implementation-efficient remedy for transformer outliers, supporting FP8 and post-training W8A8 quantization with state-of-the-art accuracy and throughput [2511.23225].

Source: https://www.emergentmind.com/topics/tweo-framework