Building MLP Block

Creating a MLP Block

In this section, we'll explore how to leverage network primitives to create complex neural network components. Specifically, we’ll focus on building a multilayer perceptron (MLP) block.

Addressing Core Functionality Gaps

Even with the foundational operations available, building deep learning models often requires more advanced mathematical functions such as square roots and exponentials, particularly in operations like calculating variance and the softmax function. The solution? Polynomial approximations.

Polynomial approximations allow us to approximate complex functions with high accuracy, substituting them into computations where homomorphic encryption restricts direct implementation. For example, we can approximate functions like GeLU (Gaussian Error Linear Unit) using a polynomial, enabling its application in an encrypted setting.

Polynomial Approximation for the GeLU Function

To implement the GeLU activation function, we can generate a polynomial approximation using numpy. Here’s an example that demonstrates how to approximate the GeLU function:

# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt

# Generate data points for the GeLU function
x = np.linspace(5, 1000, 1000000)  # Range of x values to avoid sqrt(0)
def gelu(x):
    return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * x ** 3)))
y = gelu(x)

# Fit a polynomial to approximate the GeLU function
degree = 6
coeffs = np.polyfit(x, y, degree)
poly_approx = np.poly1d(coeffs)

# Calculate error metrics
abs_mean_error = np.abs(y - poly_approx(x)).mean()
max_error = np.abs(y - poly_approx(x)).max()
median_error = np.median(np.abs(y - poly_approx(x)))

print(f"Mean absolute error: {abs_mean_error}")
print(f"Max error: {max_error}")
print(f"Median error: {median_error}")

# Plot the true function and polynomial approximation
plt.plot(x, y, label="GeLU(x)", color="blue")
plt.plot(x, poly_approx(x), label=f"Polynomial approx (degree {degree})", color="red", linestyle="--")

# Labels and legend
plt.xlabel("x")
plt.ylabel("y")
plt.title("GeLU Function and Polynomial Approximation")
plt.legend()
plt.grid(True)
plt.show()

# Print polynomial coefficients for reference
print(f"Coefficients of the polynomial: {coeffs}")

In this code, we:

1. Define the GeLU function and generate x values for evaluation.

2. Fit a polynomial of a specified degree to approximate GeLU.

3. Calculate the error of this approximation, showing how closely the polynomial follows the actual function.

4. Plot both the true GeLU function and the polynomial approximation to visually confirm the fit.

Once we have this polynomial approximation, we can use CoFHE's built-in polynomial evaluator to apply it to encrypted data.

Applying the Polynomial Approximation with CoFHE

After identifying the polynomial, we can use CoFHE’s library functions to evaluate it on encrypted data:

import cofhe

# Initialize CoFHE
cofhe.init("path/to/config.json")

# Define the polynomial based on pre-calculated coefficients
gelu_poly = cofhe.init_polynomial([-2.80379477e-16, 9.59885734e-13, -1.30614629e-09, 9.07733733e-07, -3.53767845e-04, 1.01593887e-01, 2.58332175e+00])

# Encrypt input data
input_data = cofhe.encrypt(100)

# Evaluate polynomial on encrypted data
output = cofhe.solve_polynomial(gelu_poly, input_data)

# Decrypt and display result
print(cofhe.decrypt(output))

Building the MLP Block with CoFHE

With polynomial approximations of core functions, we can build a multi-layer perceptron (MLP) block using CoFHE primitives. Below is a code snippet demonstrating how to implement the MLP block:

import cofhe

# Initialize CoFHE
cofhe.init("path/to/config.json")

class MLPBlock:
    def __init__(self, input_size, hidden_size, output_size):
        # Initialize the layers
        self.linear1 = cofhe.init_linear(input_size, hidden_size)
        self.relu = cofhe.init_polynomial([-2.80379477e-16, 9.59885734e-13, -1.30614629e-09, 9.07733733e-07, -3.53767845e-04, 1.01593887e-01, 2.58332175e+00])
        self.linear2 = cofhe.init_linear(hidden_size, output_size)

    def forward(self, x):
        x = self.linear1(x)  # First linear layer
        x = self.relu(x)     # Apply ReLU approximation
        x = self.linear2(x)  # Second linear layer
        return x

# Initialize the MLP block
mlp_block = MLPBlock(input_size=768, hidden_size=512, output_size=256)

# Encrypt input data
input_data = cofhe.encrypt(np.random.rand(768))

# Forward pass through MLP
output = mlp_block.forward(input_data)

# Decrypt and display output
output = cofhe.decrypt(output)
print(output)

In this code:

1. Initialization: We set up CoFHE, defining each layer of the MLP block with a linear transformation followed by an activation function.

2. Polynomial Activation: We use our polynomial approximation of ReLU for the activation layer.

3. Encrypted Data Flow: Data flows through the MLP block entirely in its encrypted form, making use of the secure CoFHE primitives.