Using Neural Networks to solve Ordinary Differential Equations | by Caio Davi | Towards Data Science

The idea of solving an ODE using a Neural Network was first described by Lagaris et al. The insight behind it is basically training a neural network to satisfy the conditions required by a differential equation. In other words, we need to find a function whose derivative satisfies the ODE conditions. This article will be going through the underlying mathematical foundations of this concept, and then we will implement it using TensorFlow. All the contents are also available in a Google Collab Notebook here.
Let’s say we have an ODE system, given by:
Hence, we can understand the differential operation as a function on the domain t with a known initial condition u(0)=u0. As we know, Neural Networks are known as universal approximators. We will take advantage of this property of Neural Networks to use them to approximate the solution of the given ODE:
Also, we may agree that the derivative of NN(t) will give us a similar equation:
So, if NN(t) is really close to the true solution, then we could say that its derivative is also close to the derivative of the true solution, i.e.:
Thus, we can turn this condition into our loss function. We have the given derivative function f(u,t) and we can calculate the Neural Network derivative NN′(t) at each step. This motivates the following loss function (which is the mean squared error of the two values):
You may remember the initial condition, we still need to handle that. The most straight-forward way would do this by adding an initial condition term to the cost function. It would look like this:
While that would work, it may not be the best approach. We all know the crucial importance of the loss function on the training of the Neural Network, and we also know that the number of terms on this function will impact directly the stability of our training. More terms on the loss function would (usually) imply unstable training. To avoid it, we can encode the initial condition into the loss in a more efficient way. Let’s define a new function and use it instead of directly using the neural network:
It’s easy to see that g(t) will always satisfy the initial condition since g(0) will lead to tNN(t)=0, leaving just the initial condition on the expression. This way, we can train g(t) to satisfy the ODE system instead of the Neural Network. Then, it will automatically be a solution to the derivative function. We can incorporate this new idea into our loss function:
We are about to implement the described method in python using the TensorFlow library. In order to have a better understanding of the method, we will use a low-level design, avoiding a number of possible optimizations provided by the library. Our focus, at this moment, is to clearly understand and implement the ODE-solver Neural Network. For this reason, we will also choose a simply ODE:
We can easily solve this problem by integrating both sides of the solution, leading to u+C=x2+C, and after fitting C to obey the initial condition we have u=x2+1. Nevertheless, instead of solve it analytically, let’s try to solve using Neural Nets.
For this example, we will create an MLP Neural Net with 2 hidden layers, sigmoid activation functions, and a gradient descent optimizer algorithm. Other topologies may also be used, this is just an example. We encourage you to try different approaches to this method.
Defining Variables
f0 = 1inf_s = np.sqrt(np.finfo(np.float32).eps)learning_rate = 0.01training_steps = 5000batch_size = 100display_step = 500# Network Parametersn_input = 1     # input layer number of neuronsn_hidden_1 = 32 # 1st layer number of neuronsn_hidden_2 = 32 # 2nd layer number of neuronsn_output = 1    # output layer number of neuronsweights = {'h1': tf.Variable(tf.random.normal([n_input, n_hidden_1])),'h2': tf.Variable(tf.random.normal([n_hidden_1, n_hidden_2])),'out': tf.Variable(tf.random.normal([n_hidden_2, n_output]))}biases = {'b1': tf.Variable(tf.random.normal([n_hidden_1])),'b2': tf.Variable(tf.random.normal([n_hidden_2])),'out': tf.Variable(tf.random.normal([n_output]))}# Stochastic gradient descent optimizer.optimizer = tf.optimizers.SGD(learning_rate)
Defining the Model and Loss Function
# Create modeldef multilayer_perceptron(x):  x = np.array([[[x]]],  dtype='float32')  layer_1 = tf.add(tf.matmul(x, weights['h1']), biases['b1'])  layer_1 = tf.nn.sigmoid(layer_1)  layer_2 = tf.add(tf.matmul(layer_1, weights['h2']), biases['b2'])  layer_2 = tf.nn.sigmoid(layer_2)  output = tf.matmul(layer_2, weights['out']) + biases['out']  return tf.nn.sigmoid(output)# Universal Approximatordef g(x):  return x * multilayer_perceptron(x) + f0# Given EDOdef f(x):  return 2*x# Custom loss function to approximate the derivativesdef custom_loss():  summation = []  for x in np.linspace(-1,1,10):    dNN = (g(x+inf_s)-g(x))/inf_s    summation.append((dNN - f(x))**2)  return tf.sqrt(tf.reduce_mean(tf.abs(summation)))
Notice that we don’t have any parameters on our loss function. Usually, the loss function would compare the prediction with the actual data. In our case, we don’t need data points.
Since we know the ODE function ruling this model, we can calculate the expected value at each point x. You may also notice that we are always calculating the loss using 10 points in the interval {-1,1}. It might not be enough for all possible functions, but given the simplicity of our example, I think we are gonna be ok.
The Neural Network Derivative is expressed by the variable dNN in the code. We are just using the differentiation definition:
Train Function
def train_step():  with tf.GradientTape() as tape:    loss = custom_loss()  trainable_variables=list(weights.values())+list(biases.values())  gradients = tape.gradient(loss, trainable_variables)  optimizer.apply_gradients(zip(gradients, trainable_variables))# Training the Model:for i in range(training_steps):  train_step()  if i % display_step == 0:    print("loss: %f " % (custom_loss()))
Plotting the Results
# True Solution (found analitically)def true_solution(x):  return x**2 + 1X = np.linspace(0, 1, 100)result = []for i in X:  result.append(g(i).numpy()[0][0][0])S = true_solution(X)plt.plot(X, result)plt.plot(X, S)plt.show()
This will yield us the following plot:
It is quite an approximation, mostly if we remember that was no dataset used to training. It might be even better if we had used more collocation points to calculate the loss function, or if we keep the training a little longer. We invite the reader to try out the code and make these adjustments or even try other ODEs. You can find all the contents of this post in a Google Collab Notebook here. Have fun! And thank you for the reading.