Non-record: xIELU Piecewise Quadratic Activation + Per-Layer QK Gain Convergence

[mikeapedia]

Summary

Non-record research submission. I'm out of compute credits but wanted to share these results with the community since I think they show potential. Built on top of PR #1586 by @dexhunter — two activation/attention techniques with converged per-layer coefficients, plus a modified fused Triton kernel:

• xIELU activation (inspired by arxiv 2411.13010): Replaces leaky_relu(x, 0.5).square() with a piecewise quadratic torch.where(x > 0, ap*x² + bp*x, an*x² + bn*x) using 4 hardcoded per-layer coefficients discovered via convergence loop. Zero throughput overhead.
• Per-layer QK gain init: Converged from uniform 5.0 to per-layer values in the 2.0–3.0 range via 3-seed averaging. The model consistently prefers much softer attention than the 5.0 default.
• Symmetric resid_mix on both parallel lanes: PR Record: Per-Layer Adaptive GPTQ Clip + int7 Embeddings + MLR 0.026 — val_bpb 1.07493 (3-seed mean) #1586 only applies resid_mix (x0 skip blending) to the attention lane input. We apply it to both lanes so the MLP lane also gets the x0 residual signal.

Both learnable techniques use a convergence loop methodology: train with learnable per-layer scalars → harvest post-EMA values → hardcode as init → retrain with different seed → average across seeds until stable.

xIELU Activation

The original xIELU paper uses expm1(x) in the negative branch, but this caused a 20% throughput drop even inside torch.compile. We replaced it with a pure quadratic — same gradient shape, no transcendentals, zero overhead.

Converged coefficients (averaged over 4 seeds):

XIELU_AP = [0.103, 0.196, 1.415, 1.196, 1.485, 1.546, 1.337, 1.727, 1.495, 0.988, 0.917]
XIELU_AN = [0.39, 0.578, 0.363, 0.491, 0.536, 0.548, 0.579, 0.983, 1.058, 0.935, 0.845]
XIELU_BP = [0.126, 0.07, 0.0, 0.0, 0.0, 0.002, 0.017, 0.067, 0.005, 0.058, 0.568]
XIELU_BN = [0.785, 0.638, 0.405, 0.377, 0.314, 0.289, 0.313, 0.571, 0.42, 0.286, 0.52]

Key findings:

• beta_p collapsed to ~0 for layers 1–8 — the model wants pure x² on the positive side for middle layers
• Early layers (0–1) are nearly linear — alpha_p ≈ 0.1–0.2 (very low curvature)
• Deep layers want strong asymmetric curvature — alpha_p up to 1.73 vs alpha_n ≈ 0.5–1.0
• Layer 10 retains a linear term — beta_p ≈ 0.57, unlike all other layers

Per-Layer QK Gain

Converged from uniform 5.0 init over 3 seeds (42, 1337, 2024):

QK_GAIN_INIT_PER_LAYER = [2.3495, 2.8818, 2.7627, 2.8148, 2.7893, 2.8762, 2.5657, 2.7206, 2.6426, 2.2737, 1.9741]

Every layer dropped dramatically from 5.0 — the model wants much softer attention across the board (range 2.0–2.9). Later layers prefer slightly softer attention.

Fused Triton Kernel Changes vs PR #1586

PR #1586's linear_leaky_relu_square_kernel fuses matmul + leaky_relu(h, 0.5).square() into a single Triton kernel using TMA descriptors. We modified it to linear_xielu_kernel with the following changes:

Forward path — replaced the hardcoded activation:

# PR #1586 original (unfactored):
aux = where(h > 0, h, 0.5*h)
aux = aux * aux                           # separate square step

# Ours (xIELU, factored):
aux = h * where(h > 0, ap*h + bp, an*h + bn)  # single fused expression

The factored form h * (ap*h + bp) saves a multiply per element vs the expanded ap*h*h + bp*h — 2 multiplies + 1 add inside the where + 1 outer multiply, vs 3 multiplies + 1 add.

Backward path — updated the fused gradient accordingly:

# PR #1586 original:
grad *= where(h > 0, 2*h, 0.5*h)

# Ours (xIELU):
grad *= where(h > 0, 2*ap*h + bp, 2*an*h + bn)

Additional kernel changes:

• 4 extra scalar args (ap, an, bp, bn) per layer, passed from the hardcoded coefficient arrays
• Renamed FusedLinearLeakyReLUSquareFunction → FusedXieluMLPFunction with matching arg threading
• No change to tile sizes, TMA descriptor layout, or persistent kernel structure

Symmetric resid_mix on Both Parallel Lanes

PR #1586 applies resid_mix (x0 skip connection blending) only to the attention lane:

# PR #1586:
attn_read = mix[0] * lane0 + mix[1] * x0   # attention gets x0 blending
mlp_read  = lane1                            # MLP reads lane1 directly

We apply the same mixing to both lanes:

# Ours:
attn_read = mix[0] * lane0 + mix[1] * x0   # attention gets x0 blending
mlp_read  = mix[0] * lane1 + mix[1] * x0   # MLP also gets x0 blending

This gives the MLP lane access to the same residual signal, making the two lanes symmetric. The resid_mix weights are already learnable, so the model can still choose to suppress the x0 contribution if it's not helpful.

Methodology: Convergence loops

Both techniques use the same pattern:

1. Add learnable per-layer scalars (zero throughput impact)
2. Train and harvest post-EMA values
3. Hardcode as init, retrain with different seed
4. Average across seeds until max delta < 0.2
5. Lock coefficients for zero-overhead deployment

This is general-purpose — applicable to any per-layer hyperparameter. The key insight is that per-channel learnable params are too expensive (16% throughput drop from 4×2048 softplus ops per layer), but per-layer scalars have zero measurable overhead.

Results

All results are post-EMA, pre-quantization val_bpb on 8×H100. We did not run full evaluation or TTT. The xIELU ablation was done on an earlier training base (PR #1529), so we report relative improvement rather than absolute numbers since the baselines differ:

Technique	Δ val_bpb	Notes
xIELU activation (hardcoded coefficients)	−0.0029	vs same-base baseline, zero throughput overhead
Per-layer q_gain init (converged from 5.0)	−0.0009	vs uniform 5.0 init, val_bpb stabilized at 1.0756–1.0765 across 3 seeds

⚠️ Not verified on competition eval infrastructure — no eval, no TTT, no quantization roundtrip. I ran out of compute credits. Sharing for the community to build on.

Test plan

• Syntax check passes
• Training runs complete without errors on 1×H100 and 8×H100
• Zero throughput overhead vs baseline (hardcoded coefficients)
• xIELU convergence stable across 4 seeds (max delta < 0.25)
• QK gain convergence stable across 3 seeds (all values in 2.0–3.1 range)
• Full competition eval (no compute credits — someone else please verify!)

🤖 Generated with Claude Code

[mikeapedia]

@dexhunter sharing an experiment that showed promise but I didn't have the compute to finish testing.