LELU Activation Function: Proposal for PyTorch

[falseywinchnet]

Summary

We propose the introduction of a fixed scaling sigmoid activation function to be called LELU — the Logistic Error Linear Unit. LELU serves as a lower-cost, analytically consistent alternative to GELU that is derived from the fixed-point correspondence between the logistic sigmoid and its associated cumulative distribution function (CDF) and probability density function (PDF).

---

Motivation

The Gaussian Error Linear Unit (GELU) is designed around the Gaussian CDF, offering smoother activation transitions than ReLU or ELU. However, GELU requires relatively costly operations (erf/tanh + cubic terms) and is only an approximation of the true Gaussian gate.

In contrast, LELU is grounded in the Logistic CDF which is analytically simpler, fully differentiable, and exhibits similar curvature to the Gaussian but with computational advantages. The logistic sigmoid family offers a natural match to many normalized training distributions observed in deep networks.

---

Derivation Overview

LELU is defined analogously to GELU:

LELU(x) = x * 0.5 * (1 + tanh(a * (x)))

where a = π / (2√3) arises from the equivalence between the variance of the standard logistic distribution and that of the standard normal distribution.

---

Proposal

Add torch.nn.LELU and torch.nn.functional.lelu

Documentation and tests can mirror the structure of torch.nn.GELU, with additional notes on logistic variance equivalence and performance observations.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class LELU(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x: torch.Tensor) -> torch.Tensor:
            return x * torch.sigmoid((math.pi / math.sqrt(3.0) )* x)

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

[falseywinchnet]

import math
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt

# --- 1. Activation Function Definitions ---

class LELU(nn.Module):
    """
    User-provided LELU (Logarithmic Error Linear Unit, scaled SiLU).
    This implementation scales the input to SiLU, making the activation steeper,
    and then scales the output down to normalize the magnitude.
    """
    def __init__(self):
        super().__init__()
        # Scaling factor alpha = pi / sqrt(3) ~= 1.8138
        val = math.pi / math.sqrt(3.0)
        self.register_buffer("alpha", torch.tensor(val))
        # Use the standard, non-inplace SiLU module
        self._silu = nn.SiLU(inplace=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Scale input, apply SiLU, then scale output down
        return self._silu(x * self.alpha) / self.alpha

# --- 2. Model and Data Setup ---

class SimpleMLP(nn.Module):
    """A simple Multi-Layer Perceptron to test activation functions."""
    def __init__(self, activation_fn):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1, 64),
            activation_fn,
            nn.Linear(64, 64),
            activation_fn,
            nn.Linear(64, 1)
        )

    def forward(self, x):
        return self.layers(x)

def create_synthetic_dataset(num_samples=1000):
    """
    Creates a synthetic dataset for a regression task.
    The function is designed to be complex, with both low and high-frequency components,
    making it a good test for an activation function's ability to model details.
    The input data is normally distributed, which is a condition where GELU is
    theoretically grounded.
    """
    # Use a normal distribution for X, which is a theoretical basis for GELU
    X = torch.randn(num_samples, 1) * 2.5
    # A complex target function: sinc(x) + high-frequency sin(5x) + noise
    y = torch.sinc(X) + torch.sin(5 * X) * 0.2 + torch.randn(num_samples, 1) * 0.05
    return TensorDataset(X, y)

# --- 3. Training and Evaluation ---

# Hook to capture activation distributions
activation_inputs = {}
def get_activation_input_hook(name):
    def hook(model, input, output):
        # We capture the input to the activation function
        activation_inputs[name] = input[0].detach().cpu().numpy().flatten()
    return hook

def train_model(model, train_loader, epochs=50):
    """Generic training loop for a given model."""
    optimizer = torch.optim.Adam(model.parameters(), lr=0.005)
    criterion = nn.MSELoss()
    loss_history = []

    for epoch in range(epochs):
        epoch_loss = 0
        for X_batch, y_batch in train_loader:
            optimizer.zero_grad()
            y_pred = model(X_batch)
            loss = criterion(y_pred, y_batch)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        
        avg_loss = epoch_loss / len(train_loader)
        loss_history.append(avg_loss)
        if epoch % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.6f}")
            
    return loss_history

# --- 4. Main Execution and Visualization ---

def run_comparison():
    # --- Setup ---
    torch.manual_seed(42)
    np.random.seed(42)
    
    dataset = create_synthetic_dataset()
    train_loader = DataLoader(dataset, batch_size=64, shuffle=True)
    X_test = torch.linspace(-8, 8, 400).view(-1, 1)

    activations = {
        "GELU": nn.GELU(),
        "SiLU": nn.SiLU(),
        "LELU": LELU()
    }
    
    models = {}
    histories = {}
    predictions = {}

    # --- Training ---
    for name, act_fn in activations.items():
        print(f"--- Training model with {name} ---")
        models[name] = SimpleMLP(activation_fn=act_fn)
        
        # Add a hook to the *second* activation layer to see how data flows
        hook_handle = models[name].layers[1].register_forward_hook(get_activation_input_hook(name))

        histories[name] = train_model(models[name], train_loader)
        
        # Run one forward pass on test data to capture final activation distribution
        with torch.no_grad():
            models[name].eval()
            predictions[name] = models[name](X_test)
        
        hook_handle.remove() # Clean up the hook

    # --- Visualization ---
    plt.style.use('seaborn-v0_8-whitegrid')
    
    # Plot 1: Activation Functions and Their Derivatives
    fig1, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    # Define the plotting range without requiring gradients on the range itself.
    x_range = torch.linspace(-4, 4, 200)
    
    for name, act_fn in activations.items():
        # Clone the range and set requires_grad=True to make it a leaf tensor for this calculation.
        x_in = x_range.clone().requires_grad_(True)
        y_out = act_fn(x_in)
        y_out.sum().backward()
        
        ax1.plot(x_range.numpy(), y_out.detach().numpy(), label=name)
        # x_in.grad will now be correctly populated.
        ax2.plot(x_range.numpy(), x_in.grad.numpy(), label=f"{name}'")
        
    ax1.set_title("Activation Functions", fontsize=16)
    ax1.set_xlabel("Input (x)")
    ax1.set_ylabel("Output")
    ax1.legend()
    
    ax2.set_title("Derivatives", fontsize=16)
    ax2.set_xlabel("Input (x)")
    ax2.set_ylabel("Gradient")
    ax2.legend()
    fig1.suptitle("Comparison of Activation Function Shapes", fontsize=20, y=1.02)
    fig1.tight_layout()
    fig1.savefig("plot_1_activation_shapes.png", bbox_inches='tight')

    # Plot 2: Training Loss
    fig2, ax = plt.subplots(figsize=(12, 7))
    for name, history in histories.items():
        ax.plot(history, label=name, lw=2.5)
    ax.set_title("Training Loss (MSE)", fontsize=16)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")
    ax.set_yscale('log')
    ax.legend()
    ax.grid(True, which="both", ls="--")
    fig2.savefig("plot_2_training_loss.png", bbox_inches='tight')

    # Plot 3: Model Predictions
    fig3, ax = plt.subplots(figsize=(12, 7))
    X_plot, y_plot = dataset.tensors
    ax.scatter(X_plot.numpy(), y_plot.numpy(), alpha=0.2, label="Ground Truth Data", color='gray')
    for name, pred in predictions.items():
        ax.plot(X_test.numpy(), pred.numpy(), label=f"{name} Prediction", lw=3)
    ax.set_title("Model Predictions vs. Ground Truth", fontsize=16)
    ax.set_xlabel("Input (x)")
    ax.set_ylabel("Output (y)")
    ax.set_ylim(-1.5, 1.5)
    ax.legend()
    fig3.savefig("plot_3_model_predictions.png", bbox_inches='tight')
    
    # Plot 4: Activation Distributions
    fig4, axes = plt.subplots(1, 3, figsize=(20, 5), sharey=True)
    bins = np.linspace(-4, 4, 50)
    for i, (name, inputs) in enumerate(activation_inputs.items()):
        axes[i].hist(inputs, bins=bins, density=True, alpha=0.8, label=f'Input to {name}')
        axes[i].set_title(f"Distribution into '{name}' layer", fontsize=14)
        axes[i].set_xlabel("Activation Value")
    axes[0].set_ylabel("Density")
    fig4.suptitle("Distribution of Pre-Activations (After Linear Layer)", fontsize=20, y=1.03)
    fig4.tight_layout()
    fig4.savefig("plot_4_activation_distributions.png", bbox_inches='tight')
    
    plt.show()


if __name__ == "__main__":
    run_comparison()

--- Training model with GELU ---
Epoch 1/50, Loss: 0.142152
Epoch 11/50, Loss: 0.016233
Epoch 21/50, Loss: 0.012188
Epoch 31/50, Loss: 0.011637
Epoch 41/50, Loss: 0.010681
--- Training model with SiLU ---
Epoch 1/50, Loss: 0.143378
Epoch 11/50, Loss: 0.023087
Epoch 21/50, Loss: 0.019567
Epoch 31/50, Loss: 0.014602
Epoch 41/50, Loss: 0.011912
--- Training model with LELU ---
Epoch 1/50, Loss: 0.128378
Epoch 11/50, Loss: 0.016104
Epoch 21/50, Loss: 0.012230
Epoch 31/50, Loss: 0.011170
Epoch 41/50, Loss: 0.010679

In this problem. LELU demonstrates functional equivalence with GELU.

[falseywinchnet]

import torch
import torch.nn as nn
import math, time

# --- Define LELU ---
class LELU(nn.Module):
    def __init__(self):
        super().__init__()
        val = math.pi / math.sqrt(3.0)
        val2 = 1 / val
        self.register_buffer("alpha", torch.tensor(val))
        self.register_buffer("alpha_inv", torch.tensor(val2))
        self._silu = nn.SiLU(inplace=False)

    def forward(self, x):
        return self._silu(x * self.alpha) * self.alpha_inv

# --- Benchmarking utility ---
def benchmark(act1, act2, shape=(4096, 4096), device="cuda", repeats=50):
    x = torch.randn(shape, device=device, requires_grad=True)
    losses = []

    # Warm-up
    for _ in range(10):
        (act1(x).sum()).backward(retain_graph=True)
        x.grad.zero_()

    def time_op(act):
        torch.cuda.synchronize()
        start = time.time()
        for _ in range(repeats):
            y = act(x)
            loss = y.sum()
            loss.backward(retain_graph=True)
            x.grad.zero_()
        torch.cuda.synchronize()
        return (time.time() - start) / repeats

    t1 = time_op(act1)
    t2 = time_op(act2)
    print(f"Average time per forward+backward (shape={shape}, repeats={repeats}):")
    print(f"  {act1.__class__.__name__:<10}: {t1*1000:.3f} ms")
    print(f"  {act2.__class__.__name__:<10}: {t2*1000:.3f} ms")
    print(f"  Ratio (LELU / GELU): {t1 / t2:.3f}x")

# --- Run benchmark ---
device = "cuda" if torch.cuda.is_available() else "cpu"
benchmark(LELU().to(device), nn.GELU().to(device), device=device)

on T4 gpu:
Average time per forward+backward (shape=(4096, 4096), repeats=50):
LELU : 4.733 ms
GELU : 3.253 ms
Ratio (LELU / GELU): 1.455x

on M4 mac:
Average time per forward+backward (shape=(4096, 4096), repeats=50):
LELU : 24.975 ms
GELU : 31.138 ms
Ratio (LELU / GELU): 0.802x

owing to highly optimized kernels on NVIDIA, it is difficult to get a win- all ops are fused there.

# --------- Triton kernels (1D elementwise, safe mask) ---------
import os, math, time
import torch
import torch.nn as nn

try:
    import triton
    import triton.language as tl
except Exception as e:
    raise RuntimeError("Needs Triton. pip install triton") from e

ALPHA = math.pi / math.sqrt(3.0)

# --------- Optimized Triton kernels ---------

@triton.jit
def lelu_fwd_kernel(x_ptr, y_ptr, n_el, alpha, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < n_el

    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    
    # Direct computation in input dtype when possible
    z = x * alpha
    # Use tl.sigmoid for potentially optimized path
    sig = tl.sigmoid(z)
    y = x * sig
    
    tl.store(y_ptr + offs, y, mask=mask)

@triton.jit
def lelu_bwd_kernel(x_ptr, gy_ptr, gx_ptr, n_el, alpha, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < n_el

    x = tl.load(x_ptr + offs, mask=mask, other=0.0)
    gy = tl.load(gy_ptr + offs, mask=mask, other=0.0)

    # Compute in input dtype
    z = x * alpha
    sig = tl.sigmoid(z)
    
    # dy/dx = sig + x * sig * (1 - sig) * alpha
    # Refactor: sig + x * alpha * sig * (1 - sig)
    sig_comp = sig * (1.0 - sig)
    dydx = sig + x * alpha * sig_comp
    
    gx = gy * dydx
    tl.store(gx_ptr + offs, gx, mask=mask)

# --------- Optimized autograd function with tuned configs ---------
class TritonLELUFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, alpha: float):
        ctx.alpha = float(alpha)
        if x.device.type != "cuda":
            y = x * torch.sigmoid(ctx.alpha * x)
            ctx.save_for_backward(x)
            ctx.cpu = True
            return y
        
        x_c = x.contiguous()
        y = torch.empty_like(x_c)
        n = x_c.numel()
        
        # Tuned grid config: larger blocks, more warps for better occupancy
        BLOCK_SIZE = 2048
        NUM_WARPS = 8
        grid = lambda META: (triton.cdiv(n, META['BLOCK_SIZE']),)
        lelu_fwd_kernel[grid](x_c, y, n, ctx.alpha, BLOCK_SIZE=BLOCK_SIZE, num_warps=NUM_WARPS)
        
        ctx.save_for_backward(x_c)
        ctx.cpu = False
        return y.view_as(x)

    @staticmethod
    def backward(ctx, grad_out):
        (x,) = ctx.saved_tensors
        a = ctx.alpha
        go = grad_out
        
        if ctx.cpu or x.device.type != "cuda":
            z = a * x
            sig = torch.sigmoid(z)
            dydx = sig + x * (sig * (1 - sig)) * a
            return go * dydx, None
        
        x_c = x
        go_c = go.contiguous()
        gx = torch.empty_like(x_c)
        n = x_c.numel()
        
        # Same tuned config
        BLOCK_SIZE = 2048
        NUM_WARPS = 8
        grid = lambda META: (triton.cdiv(n, META['BLOCK_SIZE']),)
        lelu_bwd_kernel[grid](x_c, go_c, gx, n, a, BLOCK_SIZE=BLOCK_SIZE, num_warps=NUM_WARPS)
        
        return gx.view_as(go), None

class TritonLELU(nn.Module):
    def __init__(self, alpha=ALPHA):
        super().__init__()
        self.alpha = float(alpha)
    def forward(self, x):
        return TritonLELUFunc.apply(x, self.alpha)

def correctness_check(device, dtype):
    torch.manual_seed(0)
    x = torch.randn(1 << 20, device=device, dtype=dtype, requires_grad=True)  # 1M elems
    x_ref = x.detach().clone().requires_grad_(True)

    lelu = TritonLELU().to(device=device, dtype=dtype)

    # forward
    y_triton = lelu(x)
    y_ref    = x_ref * torch.sigmoid(ALPHA * x_ref)

    f_err = (y_triton - y_ref).abs().max().item()

    # backward: d/dx sum(y)
    y_triton.sum().backward()
    g_triton = x.grad.detach().clone()
    x.grad.zero_()

    y_ref.sum().backward()
    g_ref = x_ref.grad.detach().clone()

    g_err = (g_triton - g_ref).abs().max().item()

    return f_err, g_err

# --------- benchmark ---------
def sync(dev):
    if dev.type == "cuda":
        torch.cuda.synchronize()

def bench(act, x, repeats=50, warmup=10):
    for _ in range(warmup):
        (act(x).sum()).backward()
        x.grad.zero_()
    sync(x.device)
    t0 = time.perf_counter()
    for _ in range(repeats):
        (act(x).sum()).backward()
        x.grad.zero_()
    sync(x.device)
    return (time.perf_counter() - t0) / repeats

def run(shape=(4096, 4096), dtype=torch.float16):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"\nDevice={device}, dtype={dtype}, shape={shape}")

    # 1) correctness first (smallish vector)
    f_err, g_err = correctness_check(device, dtype)
    print(f"max |f_triton - f_ref| = {f_err:.3e},  max |g_triton - g_ref| = {g_err:.3e}")

    # 2) perf
    x = torch.randn(shape, device=device, dtype=dtype, requires_grad=True)
    acts = [
        ("GELU(exact)", nn.GELU(approximate="none").to(device=device, dtype=dtype)),
        ("GELU(tanh) ", nn.GELU(approximate="tanh").to(device=device, dtype=dtype)),
        ("LELU(triton)", TritonLELU().to(device=device, dtype=dtype)),
    ]
    base_name, base_mod = acts[0]
    base_t = bench(base_mod, x)
    print(f"{base_name:<14}: {base_t*1e3:7.3f} ms   (x1.000)")
    for name, mod in acts[1:]:
        t = bench(mod, x)
        print(f"{name:<14}: {t*1e3:7.3f} ms   (x{t/base_t:5.3f})")

if __name__ == "__main__":
    # Run FP16 tensor path then FP32
    run(dtype=torch.float16)
    run(dtype=torch.float32)

Device=cuda, dtype=torch.float16, shape=(4096, 4096)
max |f_triton - f_ref| = 3.906e-03, max |g_triton - g_ref| = 1.953e-03
GELU(exact) : 1.597 ms (x1.000)
GELU(tanh) : 1.574 ms (x0.985)
LELU(triton) : 1.573 ms (x0.985)

Device=cuda, dtype=torch.float32, shape=(4096, 4096)
max |f_triton - f_ref| = 4.768e-07, max |g_triton - g_ref| = 7.153e-07
GELU(exact) : 2.521 ms (x1.000)
GELU(tanh) : 2.525 ms (x1.001)
LELU(triton) : 3.102 ms (x1.230)

[falseywinchnet]

lelu amounts to x * sigmoid(scale * x)
with scale a = math.pi / math.sqrt(3.0)