ff_modifications.py

ff_modifications.py

import torch
import torch.nn as nn

class FeedForwardLayer(nn.Module):
    def __init__(self, d_model, nonlinearity):
        super(FeedForwardLayer, self).__init__()
        self.linear1 = nn.Linear(d_model, d_model*6)
        self.linear2 = nn.Linear(d_model, d_model)
        self.nonlinearity = nonlinearity() # literally any function that is six-to-one

    def forward(self, x):
        x = self.linear1(x)
        x = self.nonlinearity(x)
        x = self.linear2(x)
        return x

class AveragePoolingWithGelu(nn.Module):
    def __init__(self):
        super(AveragePoolingWithGelu, self).__init__()
        self.avg_pool = nn.AvgPool1d(6)

    def forward(self, x):
        pooled = self.avg_pool(x)
        gelu_activation = torch.nn.functional.gelu(pooled)
        return gelu_activation
    
class MaxPoolingWithGelu(nn.Module):
    def __init__(self):
        super(MaxPoolingWithGelu, self).__init__()
        self.max_pool = nn.MaxPool1d(6)

    def forward(self, x):
        pooled = self.max_pool(x)
        gelu_activation = torch.nn.functional.gelu(pooled)
        return gelu_activation
    
class SumPoolingWithGelu(nn.Module):
    def __init__(self):
        super(SumPoolingWithGelu, self).__init__()
        self.avg_pool = nn.AvgPool1d(6)

    def forward(self, x):
        pooled = self.avg_pool(x) * 6 # gross, but works
        gelu_activation = torch.nn.functional.gelu(pooled)
        return gelu_activation

# Could also have called this hexalinear, in honor of bilinear
class ProductPoolingWithGelu(nn.Module):
    def __init__(self):
        super(ProductPoolingWithGelu, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = x_1 * x_2 * x_3 * x_4 * x_5 * x_6
        gelu_activation = torch.nn.functional.gelu(pooled)
        return gelu_activation
    
class YinLU(nn.Module):
    def __init__(self):
        super(YinLU, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = torch.sin(x_1) * torch.exp(-torch.pow(x_2, 2)) * torch.tanh(x_3) * torch.nn.functional.gelu(x_4) * x_5 * x_6
        return pooled

# Sort of named in honor of bilinear layers
class BiSwigLU(nn.Module):
    def __init__(self):
        super(BiSwigLU, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = torch.nn.functional.silu(x_1) * x_2 * x_3 + torch.nn.functional.silu(x_4) * x_5 * x_6
        return pooled

class BiGeGLU(nn.Module):
    def __init__(self):
        super(BiGeGLU, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        
        pooled = torch.nn.functional.gelu(x_1) * x_2 * x_3 + torch.nn.functional.gelu(x_4) * x_5 * x_6
        return pooled
    
class PolyLU(nn.Module):
    def __init__(self):
        super(PolyLU, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        pooled = x_1 * x_2 + torch.pow(x_1, 2) * x_3 + torch.pow(x_1, 3) * x_4 + torch.pow(x_1, 4) * x_5 + torch.pow(x_1, 5) * x_6
        return pooled

class BiPolyLU(nn.Module):
    def __init__(self):
        super(BiPolyLU, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        bi = x_1 * x_2
        bi_squared = torch.pow(bi, 2)
        bi_cube = torch.pow(bi, 3)
        bi_quad = torch.pow(bi, 4)

        pooled = bi * x_3 + bi_squared * x_4 + bi_cube * x_5 + bi_quad * x_6
        return pooled

class GyLU(nn.Module):
    def __init__(self):
        super(GyLU, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        a = x_1 * x_2
        b = x_3 * x_4
        c = x_5 * x_6

        # It's a ring, not because it makes sense but because it is funny
        pooled = nn.functional.silu(a) * b + nn.functional.silu(b) * c + nn.functional.silu(c) * a
        return pooled
    
class GyLU2(nn.Module):
    def __init__(self):
        super(GyLU2, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        a = x_1 * x_2
        b = x_3 * x_4
        c = x_5 * x_6

        # Why not
        pooled = nn.functional.silu(a + b) * c
        return pooled

class GyLU3(nn.Module):
    def __init__(self):
        super(GyLU3, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        a = x_1 * x_2
        b = x_3 * x_4

        # Figure the third variable c probably isn't doing anything
        pooled = nn.functional.silu(a) * b
        return pooled
    
class InnerProductPooling(nn.Module):
    def __init__(self):
        super(InnerProductPooling, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        a = x_1 * x_2
        b = x_3 * x_4
        c = x_5 * x_6

        pooled = a + b + c
        return pooled
    
class FourierPooling(nn.Module):
    def __init__(self):
        super(FourierPooling, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        pooled = torch.sin(x_1*x_2) + torch.sin(x_1*x_3) + torch.sin(x_1*x_4) + torch.sin(x_1*x_5) + torch.sin(x_1*x_6)
        return pooled
    
class NoisyFourierPooling(nn.Module):
    def __init__(self, noise_scaling=1000):
        super(NoisyFourierPooling, self).__init__()
        self.noise_scaling = noise_scaling

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        variances = torch.var(x_reshaped, dim=3)
        randoms = torch.randn(x_reshaped.size())
        result = torch.mul(variances.unsqueeze(-1), randoms)/self.noise_scaling
        x_reshaped = x_reshaped + result
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        pooled = torch.sin(x_1*x_2) + torch.sin(x_1*x_3) + torch.sin(x_1*x_4) + torch.sin(x_1*x_5) + torch.sin(x_1*x_6)
        return pooled
    
class BiSwigLU2(nn.Module):
    def __init__(self):
        super(BiSwigLU2, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = torch.nn.functional.silu(x_1 * x_2) * x_3 + torch.nn.functional.silu(x_4 * x_5) * x_6
        return pooled

class BiGeGLU2(nn.Module):
    def __init__(self):
        super(BiGeGLU2, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        
        pooled = torch.nn.functional.gelu(x_1 * x_2) * x_3 + torch.nn.functional.gelu(x_4 * x_5) * x_6
        return pooled
    
class TriSwigLU(nn.Module):
    def __init__(self):
        super(TriSwigLU, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = torch.nn.functional.silu(x_1) * x_2 + torch.nn.functional.silu(x_3) * x_4 + torch.nn.functional.silu(x_5) * x_6
        return pooled

class TriGeGLU(nn.Module):
    def __init__(self):
        super(TriGeGLU, self).__init__()
    
    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        
        pooled = torch.nn.functional.gelu(x_1) * x_2 + torch.nn.functional.gelu(x_3) * x_4 + torch.nn.functional.gelu(x_5) * x_6
        
        return pooled


class YinLU2(nn.Module):
    def __init__(self):
        super(YinLU2, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)
        pooled = torch.sin(x_1) + torch.exp(-torch.pow(x_2, 2)) + torch.tanh(x_3) + torch.nn.functional.gelu(x_4) + x_5 * x_6
        return pooled
    

class GatedFourierPooling(nn.Module):
    def __init__(self):
        super(GatedFourierPooling, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        pooled = torch.sin(x_1*x_2) * nn.functional.gelu(x_3) + torch.sin(x_1*x_4) * nn.functional.gelu(x_5) + x_6 # bias, I guess
        return pooled

class GatedFourierPooling2(nn.Module):
    def __init__(self):
        super(GatedFourierPooling2, self).__init__()

    def forward(self, x):
        batch_size, sequence_length, embedding_dim = x.size()
        block_size = embedding_dim // 6
        x_reshaped = x.view(batch_size, sequence_length, 6, block_size)
        x_1, x_2, x_3, x_4, x_5, x_6 = torch.unbind(x_reshaped, dim=2)

        pooled = torch.sin(x_1*x_2) + torch.sin(x_1*x_3) + torch.sin(x_1*x_4) + torch.sin(x_1*x_5)
        return pooled * nn.functional.gelu(x_6)

# Example usage
def test_modules(modules, input_size):
    for module in modules:
        module_instance = module()  # Instantiate the module to get its class name
        print(f"Testing {module_instance.__class__.__name__}:")
        input_tensor = torch.randn(*input_size)  # Example input tensor
        feedforward_layer = FeedForwardLayer(input_size[-1], module)
        output_tensor = feedforward_layer(input_tensor)
        print(f"Output size: {output_tensor.size()}")  # Print the output tensor's size
        print()


input_size = (4, 10, 512)

modules_to_test = [
    AveragePoolingWithGelu,
    MaxPoolingWithGelu,
    SumPoolingWithGelu,
    ProductPoolingWithGelu,
    YinLU,
    BiSwigLU,
    BiGeGLU,
    PolyLU,
    BiPolyLU,
    GyLU,
    GyLU2,
    GyLU3,
    InnerProductPooling,
    FourierPooling,
    NoisyFourierPooling,
    BiSwigLU2,
    BiGeGLU2,
    TriSwigLU,
    TriGeGLU,
    YinLU2,
    GatedFourierPooling,
    GatedFourierPooling2,
]

test_modules(modules_to_test, input_size)