Easiest way to train a neural-network with neurons that deviate from $f_{nl}(x \cdot A)$

Question

I want to model how a neural network would behave for a system of input-output devices that are only approximately similar to a neuron. I think I have a resonable plan for how to do this, but I'm looking for some help in seeing if this is something that needs to be completely coded from scratch, or if this is something that can use some out-of-the-box tools in ML that already exist.

So first I'll describe what I want in more detail:

For a neuron in a neural network, the neuron in the model can be described by the following picture:

the output of the neuron is $output = f_{nl}(x \cdot A)$, where the inputs X are dot-producted with the weights A and are then a nonlinear function $f_{nl}$ is applied.

Now I am considering a mathematical model that behaves similar to a neuron, but not exactly.

This new model systematically gives something that deviates from this nonlinear process. For example, consider an output that looks like

$output = f_{nl}(x \cdot A-\epsilon A x^2)$,

For some small value of $\epsilon$ this is approximately a neuron. Could such an input-output system be trained with the same back-propagation techniques to be able to reliably learn on a training set?

Furthermore, I am considering involves a real physical system that needs to be solved with a system of differential equations, it might not have an analytical representation, and cannot be represented by this simple example of $f_{nl}(x \cdot A-\epsilon A x^2)$ instead it is just a general relationship:

$output = f(x, A)$,

The function is continuous, but needs to be solved numerically. It's very similar to the behavior of a neuron, but not exactly. I would like to see if such a thing can be trained on. I am thinking that I can create a neural network to learn the behavior of one of these "almost neurons" and then chain these neurons together to make a neural network.

My plan is to perform a large set of simulations for a wide set of inputs to collect a set of data representing the input-output response. Then I will make a neural network that learns the behavior of one of these "almost neurons."

For example, I could train it with a single neuron:

Let's call the inputs to this neuron Q, which consist of the both the inputs and weights of the input-output function that looks approximately like a neuron. The output of this neuron is:

$output = f_{nl}(Q \cdot B)$

In this example I have a 1 layer NN that is trained to learn the input-output behavior of my nonlinear function f(x, A). In principle this doesn't have to be a 1-layer NN, and could be some neural netweork that I've trained to represent f(x, A) (which approximately looks like a typical neuron).

The values of B are learned via the training data based on how my input nodes and weights vary the output f(x, A). I will obtain the values of B by trainining them on the a large dataset describing the input-out relationship of f(x, A) (which is obtained by repeatedly solving some differential equations).

So now I have a trained neural network which describing my nonlinear function, $output = f(x, A)$. Now I want to fix these weights and chain a number of these nonlinear functions to form a neural network.

So I can take my trained output behavior and create a neural network based on chaining a bunch of these things in series, to look like:

The values of B are fixed-weights that have been predetermed by the previous training, and characterizes the input-output behavior each of the individual input-output devices that are chained to form the network. The values of A are the unfixed weights that will be trained for a general machine learning data set. So for instance this system of 3 input-output devices will be trained via gradient decent to see if it can learn on pictures of numbers.

In this case it seems like the tricky thing is setting up this neural network which has some fixed weights, while others are tunable.

So ultimately my question is what is the fastest or easiest way to do such a thing. I am pretty new to coding machine learning projects, so I was wondering if people knew if such a type of problem is possible to handle with an "out-of-the-box" solution, or if I will really need to write my own code from scratch. For example will I need to write my own backpropagation algorithms, or does this seem like something like (for example) Keras will be able to handle?

Answer 1

In your first example, i.e. $f_{nl}(x\cdot A-\epsilon Ax^2)$, since $Ax^2$ is continuous in the parameters $A$ you can implement this with common ML tools like tensorflow and pytorch, even with multiple neurons and in a batch fashion directly. It would be something like:

h1 = matmul(X, A)
h2 = matmul(A, square(X))
out = activation(h1 - epsilon * h2)

where epsilon I think is an hyper-parameter, i.e. something that is constant and you don't want to learn with back-prop.

So if you know the mathematical form of the novel "neuron" (or better said a layer) you want to implement, assuming also the function is differentiable, you can implement and train it easily.

But the second case, when you talk about a $f$ whose form is unknown, can be more difficult. If you have access to both inputs and output pairs, you can approximate $f$ the usual way by training a DNN regressor model: for example minimizing an $l_2$-loss or something more specific. Also, what you have described sounds like Neural ODE: I don't know much about it, but it's related with differential equations and, indeed, neural networks.