Manufacturing Dataset Model does not perform at all

[meeen]

Hello,

I am working on a research problem regarding a manufacturing process that can be solved as a regression. The dataset consists of 4 Input columns (e.g. pressure, depth) like and one output all being floats.

My Neuronal Network is able to perform with r2 = 0.99

class TestModel(nn.Module):
    def __init__(self):
        super(TestModel, self).__init__()
        self.layer1 = nn.Linear(4, 64)
        self.layer2 = nn.Linear(64, 16)
        self.layer3 = nn.Linear(16, 8)
        self.output_layer = nn.Linear(8, 1)

    def forward(self, x):
        x = nn.ReLU()(self.layer1(x))
        x = nn.ReLU()(self.layer2(x))
        x = nn.ReLU()(self.layer3(x))
        x = self.output_layer(x)
        return x

I use 200 epochs, lr=0.001, mse.loss, adam optimizer, batch size of 16. How can I build some thing that performance good with pykan? It would be great because one of the goals of the research project is to describe the underlying rules of the process which a symbolic formula could accomplish.

model = KAN(width=[4,[5,5],1], grid=5, k=5, seed=0)

dtype = torch.get_default_dtype()

def train_acc():
    return torch.mean((torch.round(model(dataset['train_input'])[:,0]) == dataset['train_label'][:,0]).type(dtype))

def test_acc():
    return torch.mean((torch.round(model(dataset['test_input'])[:,0]) == dataset['test_label'][:,0]).type(dtype))

results = model.fit(dataset, steps=200, lamb=0.002, metrics=(train_acc, test_acc));
results['train_acc'][-1], results['test_acc'][-1]
 

I tried different configurations, with multiplication and without, different layers, but everything performant badly. How can I optimize my approach?

Thank you for your help