SwiGLU: Swish-Gated Linear Unit

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

class Swish(nn.Module):
    
    
    def __init__(self, beta=1.0):
        super().__init__()
        self.beta = beta
    
    def forward(self, x):
        return x * torch.sigmoid(self.beta * x)

class SwiGLU(nn.Module):
    
    
    def __init__(self, dim_input, dim_hidden, bias=True):
        super().__init__()
        
        # Two linear projections: one for gating, one for value
        self.W_gate = nn.Linear(dim_input, dim_hidden, bias=bias)
        self.W_value = nn.Linear(dim_input, dim_hidden, bias=False)  # No bias on value in modern implementations
        
    def forward(self, x):
        
        # Gate pathway with Swish activation
        gate = self.W_gate(x)
        gate = F.silu(gate)  # SiLU is equivalent to Swish with beta=1
        
        # Value pathway (linear)
        value = self.W_value(x)
        
        # Gating: element-wise multiplication
        output = gate * value
        
        return output

class FeedForwardSwiGLU(nn.Module):
    
    
    def __init__(self, dim, hidden_dim=None, multiple_of=256, dropout=0.0):
        super().__init__()
        
        # Hidden dim is typically 2.66x input dim in SwiGLU (to maintain parameter count)
        # Standard: hidden_dim = 4 * dim for GELU, ~8/3 * dim for SwiGLU
        if hidden_dim is None:
            hidden_dim = 4 * dim
            # Adjust to multiple_of for hardware efficiency
            hidden_dim = int(2 * hidden_dim / 3)  # SwiGLU uses 2/3 for comparable params
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)  # Gate projection
        self.w2 = nn.Linear(dim, hidden_dim, bias=False)  # Value projection
        self.w3 = nn.Linear(hidden_dim, dim, bias=False)  # Output projection
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        
        # SwiGLU: Swish(w1(x)) * w2(x)
        gate = F.silu(self.w1(x))  # SiLU/Swish activation
        hidden = gate * self.w2(x)  # Gating
        output = self.w3(hidden)
        output = self.dropout(output)
        
        return output

# NumPy implementation for understanding gradients
class SwiGLUNumPy:
    
    
    def __init__(self, input_dim, hidden_dim):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        
        # Xavier initialization
        self.W_gate = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim)
        self.W_value = np.random.randn(input_dim, hidden_dim) * np.sqrt(2.0 / input_dim)
        self.b_gate = np.zeros(hidden_dim)
    
    def swish(self, x):
        return x * (1 / (1 + np.exp(-np.clip(x, -500, 500))))
    
    def swish_derivative(self, x):
        sigmoid = 1 / (1 + np.exp(-np.clip(x, -500, 500)))
        return sigmoid + x * sigmoid * (1 - sigmoid)
    
    def forward(self, x):
        
        # Linear projections
        z_gate = np.dot(x, self.W_gate) + self.b_gate
        z_value = np.dot(x, self.W_value)
        
        # Swish activation on gate
        gate = self.swish(z_gate)
        
        # Gating
        output = gate * z_value
        
        cache = (x, z_gate, z_value, gate)
        return output, cache
    
    def backward(self, grad_output, cache, learning_rate=0.01):
        
        x, z_gate, z_value, gate = cache
        
        # Gradient through gating (element-wise multiplication)
        grad_gate = grad_output * z_value
        grad_z_value = grad_output * gate
        
        # Gradient through Swish
        grad_z_gate = grad_gate * self.swish_derivative(z_gate)
        
        # Gradients w.r.t. weights
        grad_W_gate = np.dot(x.T, grad_z_gate)
        grad_b_gate = np.sum(grad_z_gate, axis=0)
        grad_W_value = np.dot(x.T, grad_z_value)
        
        # Update weights
        self.W_gate -= learning_rate * grad_W_gate
        self.b_gate -= learning_rate * grad_b_gate
        self.W_value -= learning_rate * grad_W_value
        
        # Gradient w.r.t. input for backprop
        grad_x = np.dot(grad_z_gate, self.W_gate.T) + np.dot(grad_z_value, self.W_value.T)
        
        return grad_x

# Visualization and comparison
def visualize_activations():
    x = torch.linspace(-5, 5, 1000)
    
    # Different activations
    relu = F.relu(x)
    gelu = F.gelu(x)
    swish = x * torch.sigmoid(x)
    sigmoid = torch.sigmoid(x)
    
    # Plotting
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    
    # Activations
    ax = axes[0, 0]
    ax.plot(x.numpy(), relu.numpy(), label='ReLU', linewidth=2)
    ax.plot(x.numpy(), gelu.numpy(), label='GELU', linewidth=2)
    ax.plot(x.numpy(), swish.numpy(), label='Swish', linewidth=2)
    ax.set_title('Activation Functions')
    ax.legend()
    ax.grid(True)
    ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
    ax.axvline(x=0, color='k', linestyle='--', alpha=0.3)
    
    # Derivatives
    ax = axes[0, 1]
    x_grad = x[:-1] + 0.5 * (x[1] - x[:-1])
    relu_grad = torch.diff(relu) / torch.diff(x)
    gelu_grad = torch.diff(gelu) / torch.diff(x)
    swish_grad = torch.diff(swish) / torch.diff(x)
    
    ax.plot(x_grad.numpy(), relu_grad.numpy(), label='ReLU'', linewidth=2)
    ax.plot(x_grad.numpy(), gelu_grad.numpy(), label='GELU'', linewidth=2)
    ax.plot(x_grad.numpy(), swish_grad.numpy(), label='Swish'', linewidth=2)
    ax.set_title('Derivatives')
    ax.legend()
    ax.grid(True)
    ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
    ax.axvline(x=0, color='k', linestyle='--', alpha=0.3)
    
    # SwiGLU behavior visualization
    ax = axes[1, 0]
    swiglu = SwiGLU(1, 1)
    x_2d = torch.linspace(-3, 3, 100).unsqueeze(1)
    with torch.no_grad():
        # Manual SwiGLU computation
        gate = torch.sigmoid(x_2d) * x_2d  # Swish
        value = x_2d  # Identity for visualization
        output = gate * value
    
    ax.plot(x_2d.numpy(), x_2d.numpy(), label='Input', alpha=0.5)
    ax.plot(x_2d.numpy(), gate.numpy(), label='Gate (Swish)', alpha=0.7)
    ax.plot(x_2d.numpy(), output.numpy(), label='SwiGLU Output', linewidth=2)
    ax.set_title('SwiGLU Gating Behavior')
    ax.legend()
    ax.grid(True)
    
    # Smoothness comparison
    ax = axes[1, 1]
    ax.plot(x.numpy(), np.abs(np.diff(relu.numpy(), prepend=relu[0].numpy())), label='ReLU |Δ|', alpha=0.7)
    ax.plot(x.numpy(), np.abs(np.diff(gelu.numpy(), prepend=gelu[0].numpy())), label='GELU |Δ|', alpha=0.7)
    ax.plot(x.numpy(), np.abs(np.diff(swish.numpy(), prepend=swish[0].numpy())), label='Swish |Δ|', alpha=0.7)
    ax.set_title('Smoothness (Change Magnitude)')
    ax.legend()
    ax.grid(True)
    
    plt.tight_layout()
    plt.savefig('swiglu_visualization.png', dpi=150)
    print('Visualization saved to swiglu_visualization.png')

# Test implementations
print('Testing SwiGLU implementations...')

batch_size = 4
seq_len = 128
dim = 512
x = torch.randn(batch_size, seq_len, dim)

# Test SwiGLU
swiglu = SwiGLU(dim, dim * 2)
output = swiglu(x)
print(f'SwiGLU input shape: {x.shape}')
print(f'SwiGLU output shape: {output.shape}')

# Test FeedForward SwiGLU
ffn = FeedForwardSwiGLU(dim)
output_ffn = ffn(x)
print(f'
FeedForward SwiGLU input shape: {x.shape}')
print(f'FeedForward SwiGLU output shape: {output_ffn.shape}')

# NumPy version
swiglu_np = SwiGLUNumPy(dim, dim * 2)
x_np = x[0, 0].numpy().reshape(1, -1)
output_np, cache = swiglu_np.forward(x_np)
print(f'
NumPy SwiGLU output shape: {output_np.shape}')

# Gradient test
grad_output = np.ones_like(output_np)
grad_x = swiglu_np.backward(grad_output, cache)
print(f'Gradient shape: {grad_x.shape}')

# Parameter count comparison
ffn_relu = nn.Sequential(
    nn.Linear(dim, 4 * dim),
    nn.ReLU(),
    nn.Linear(4 * dim, dim)
)

ffn_gelu = nn.Sequential(
    nn.Linear(dim, 4 * dim),
    nn.GELU(),
    nn.Linear(4 * dim, dim)
)

ffn_swiglu = FeedForwardSwiGLU(dim, hidden_dim=int(8 * dim / 3))

print(f'
Parameter comparison:')
print(f'FFN ReLU: {sum(p.numel() for p in ffn_relu.parameters())/1e6:.2f}M')
print(f'FFN GELU: {sum(p.numel() for p in ffn_gelu.parameters())/1e6:.2f}M')
print(f'FFN SwiGLU: {sum(p.numel() for p in ffn_swiglu.parameters())/1e6:.2f}M')

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

# Standard PyTorch SwiGLU (used in Llama, PaLM)
class LlamaSwiGLU(nn.Module):
    
    
    def __init__(self, dim, hidden_dim=None, multiple_of=256, dropout=0.0):
        super().__init__()
        
        # Llama uses 2/3 of 4*dim for SwiGLU to match parameter count with standard FFN
        if hidden_dim is None:
            hidden_dim = 4 * dim
            hidden_dim = int(2 * hidden_dim / 3)
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
        
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)  # Gate projection
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)  # Down projection
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)  # Value projection
        
    def forward(self, x):
        
        # Llama FFN: w2(silu(w1(x)) * w3(x))
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

# Hugging Face Transformers integration
try:
    from transformers.activations import ACT2FN
    from transformers.models.llama.modeling_llama import LlamaMLP
    
    # Llama MLP with SwiGLU
    llama_mlp = LlamaMLP(config=None, hidden_size=512, intermediate_size=1376)
    print('HuggingFace Llama MLP loaded')
    
except ImportError:
    print('HuggingFace transformers not available')

# Custom SwiGLU variants
class SwiGLUVariant(nn.Module):
    
    
    def __init__(self, variant='swiglu', dim=512, hidden_dim=1376):
        super().__init__()
        self.variant = variant
        
        self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
        self.w_value = nn.Linear(dim, hidden_dim, bias=False)
        self.w_out = nn.Linear(hidden_dim, dim, bias=False)
        
    def forward(self, x):
        
        if self.variant == 'swiglu':
            # Swish(xW) * xV
            gate = F.silu(self.w_gate(x))
        elif self.variant == 'geglu':
            # GELU(xW) * xV
            gate = F.gelu(self.w_gate(x))
        elif self.variant == 'reglu':
            # ReLU(xW) * xV
            gate = F.relu(self.w_gate(x))
        elif self.variant == 'glu':
            # sigmoid(xW) * xV
            gate = torch.sigmoid(self.w_gate(x))
        else:
            raise ValueError(f'Unknown variant: {self.variant}')
        
        value = self.w_value(x)
        return self.w_out(gate * value)

# Memory-efficient SwiGLU with fused operations
class MemoryEfficientSwiGLU(nn.Module):
    
    
    def __init__(self, dim, hidden_dim, use_checkpoint=False):
        super().__init__()
        
        # Combine gate and value into single projection then split
        self.w_in = nn.Linear(dim, 2 * hidden_dim, bias=False)
        self.w_out = nn.Linear(hidden_dim, dim, bias=False)
        self.hidden_dim = hidden_dim
        self.use_checkpoint = use_checkpoint
    
    def forward(self, x):
        
        if self.use_checkpoint and self.training:
            return torch.utils.checkpoint.checkpoint(self._forward, x)
        return self._forward(x)
    
    def _forward(self, x):
        
        # Single projection for both gate and value
        x_proj = self.w_in(x)
        gate, value = x_proj.split(self.hidden_dim, dim=-1)
        
        # Apply Swish to gate and multiply
        return self.w_out(F.silu(gate) * value)

# Integration with Transformer blocks
class TransformerBlockWithSwiGLU(nn.Module):
    
    
    def __init__(self, dim, num_heads, mlp_ratio=2.67, dropout=0.0):
        super().__init__()
        
        self.norm1 = nn.RMSNorm(dim) if hasattr(nn, 'RMSNorm') else nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
        self.norm2 = nn.RMSNorm(dim) if hasattr(nn, 'RMSNorm') else nn.LayerNorm(dim)
        
        # SwiGLU FFN
        hidden_dim = int(mlp_ratio * dim)
        self.ffn = LlamaSwiGLU(dim, hidden_dim)
        
    def forward(self, x, mask=None):
        
        # Attention with pre-normalization
        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x), attn_mask=mask)[0]
        
        # FFN with SwiGLU
        x = x + self.ffn(self.norm2(x))
        
        return x

# TensorFlow/Keras implementation
import tensorflow as tf

class SwiGLUTF(tf.keras.layers.Layer):
    
    
    def __init__(self, hidden_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden_dim = hidden_dim
    
    def build(self, input_shape):
        dim = input_shape[-1]
        self.w_gate = self.add_weight(
            shape=(dim, self.hidden_dim),
            initializer='glorot_uniform',
            trainable=True,
            name='w_gate'
        )
        self.w_value = self.add_weight(
            shape=(dim, self.hidden_dim),
            initializer='glorot_uniform',
            trainable=True,
            name='w_value'
        )
        self.w_out = self.add_weight(
            shape=(self.hidden_dim, dim),
            initializer='glorot_uniform',
            trainable=True,
            name='w_out'
        )
        super().build(input_shape)
    
    def call(self, inputs):
        gate = tf.nn.silu(tf.matmul(inputs, self.w_gate))
        value = tf.matmul(inputs, self.w_value)
        hidden = gate * value
        return tf.matmul(hidden, self.w_out)

# JAX/Flax implementation
try:
    import jax
    import jax.numpy as jnp
    from flax import linen as nn
    
    class SwiGLUFlax(nn.Module):
        hidden_dim: int
        
        @nn.compact
        def __call__(self, x):
            dim = x.shape[-1]
            
            w_gate = self.param('w_gate', nn.initializers.glorot_uniform(), (dim, self.hidden_dim))
            w_value = self.param('w_value', nn.initializers.glorot_uniform(), (dim, self.hidden_dim))
            w_out = self.param('w_out', nn.initializers.glorot_uniform(), (self.hidden_dim, dim))
            
            gate = jax.nn.silu(jnp.dot(x, w_gate))
            value = jnp.dot(x, w_value)
            hidden = gate * value
            return jnp.dot(hidden, w_out)
    
    print('JAX/Flax SwiGLU available')
    
except ImportError:
    print('JAX not available')

# Test all implementations
print('
Testing SwiGLU variants...')
dim = 512
batch_size = 2
seq_len = 64
x = torch.randn(batch_size, seq_len, dim)

# Test different GLU variants
variants = ['swiglu', 'geglu', 'reglu', 'glu']
for variant in variants:
    glu = SwiGLUVariant(variant, dim, int(dim * 2.67))
    output = glu(x)
    print(f'{variant.upper()}: output shape {output.shape}, mean {output.mean().item():.4f}')

# Memory efficient version
mem_eff = MemoryEfficientSwiGLU(dim, int(dim * 2.67))
output = mem_eff(x)
print(f'
Memory Efficient SwiGLU: output shape {output.shape}')

# Compare activations
swiglu_mod = LlamaSwiGLU(dim)
gelu_ffn = nn.Sequential(
    nn.Linear(dim, int(dim * 4)),
    nn.GELU(),
    nn.Linear(int(dim * 4), dim)
)

with torch.no_grad():
    out_swiglu = swiglu_mod(x)
    out_gelu = gelu_ffn(x)
    
print(f'
Activation statistics:')
print(f'SwiGLU: mean={out_swiglu.mean():.4f}, std={out_swiglu.std():.4f}')
print(f'GELU: mean={out_gelu.mean():.4f}, std={out_gelu.std():.4f}')