bilinear NN, attempt 1

bilinear_nn.py

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# Define a synthetic multi-linear function f(x, y, z, t) = x * y + z * t
def true_function(x, y, z, t):
    return x * y + z * t

# Generate synthetic data
np.random.seed(42)
torch.manual_seed(42)

num_samples = 1000
input_dim = 1

x = np.random.uniform(-10, 10, (num_samples, input_dim)).astype(np.float32)
y = np.random.uniform(-10, 10, (num_samples, input_dim)).astype(np.float32)
z = np.random.uniform(-10, 10, (num_samples, input_dim)).astype(np.float32)
t = np.random.uniform(-10, 10, (num_samples, input_dim)).astype(np.float32)

# Calculate the target output using the true function
target = true_function(x, y, z, t)

# Convert data to PyTorch tensors
x_tensor = torch.tensor(x)
y_tensor = torch.tensor(y)
z_tensor = torch.tensor(z)
t_tensor = torch.tensor(t)
target_tensor = torch.tensor(target)

# Define the MultiLinearNet from the previous example
class MultiLinearNet(nn.Module):
    def __init__(self, input_dim, embed_dim):
        super(MultiLinearNet, self).__init__()
        # Embedding layers for each input
        self.embed_x = nn.Linear(input_dim, embed_dim)
        self.embed_y = nn.Linear(input_dim, embed_dim)
        self.embed_z = nn.Linear(input_dim, embed_dim)
        self.embed_t = nn.Linear(input_dim, embed_dim)
        # Output layer to map embeddings to scalar output
        self.output_layer = nn.Linear(embed_dim, 1)
    
    def forward(self, x, y, z, t):
        # Embeddings
        E_x = self.embed_x(x)
        E_y = self.embed_y(y)
        E_z = self.embed_z(z)
        E_t = self.embed_t(t)

        # Bilinear interactions
        E_xy = E_x * E_y
        E_zt = E_z * E_t

        # Summation for multilinear output
        combined = E_xy + E_zt

        # Final output
        output = self.output_layer(combined)

        return output

# Initialize model, loss function, and optimizer
embed_dim = 5
model = MultiLinearNet(input_dim, embed_dim)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training loop
num_epochs = 500
losses = []

for epoch in range(num_epochs):
    # Forward pass
    outputs = model(x_tensor, y_tensor, z_tensor, t_tensor)
    loss = criterion(outputs, target_tensor)
    
    # Backward pass and optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Store loss for visualization
    losses.append(loss.item())

    # Print loss every 50 epochs
    if (epoch+1) % 50 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# Plot the training loss
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('MSE Loss')
plt.title('Training Loss Over Epochs')
plt.show()

# Test the model with a new sample
test_x = torch.tensor([[2.0]])
test_y = torch.tensor([[3.0]])
test_z = torch.tensor([[1.5]])
test_t = torch.tensor([[4.0]])

# Ground truth for comparison
true_value = true_function(2.0, 3.0, 1.5, 4.0)

# Model prediction
model.eval()
with torch.no_grad():
    prediction = model(test_x, test_y, test_z, test_t)

print(f"True Value: {true_value}, Model Prediction: {prediction.item()}")