Finding a half-iterate approximation using neural networks

Question

I'm trying to find good approximate half-iterates, otherwise known as functional square roots. Given a function $g$, I want to find a function $f$ such that $g(x)=f(f(x))$ over some domain. There are a few ways to approach finding a functional square root including Newton series, matrix square roots of the Carleman matrix, and others. These are all well known, but they're kind of difficult to use and have all sorts of numerical problems.

I'd like to know if neural networks could be applied to the problem. Please bear in mind I only have a surface level understanding of neural networks - I just had a vague idea how they could be relevant here.

Suppose I am trying to find the $f$ such that $f(f(x))=\sin(x)$. I know we can train a neural network to learn a good approximation to $\sin(x)$. Is there a way to construct a neural network to learn an approximation to the half-iterate $f(x)$ ? We do not know the derivative of $f$, but we have $f'(x)f'(f(x))=\cos(x)$

If the network had one single input node, some middle layers, and a single output node, then maybe something like this would work, starting with a randomly initialized network:

1. for a given input $x_i$ and network state $N$, calculate the output $y_i$
2. use the output $y_i$ of step 1. as an input $x_i'$ of the same network to generate another output $y_i'$.
3. calculate the error $(\sin(x_i) - y_i')^2$ and use this to back-propagate and update weights for $N'$.

Any ideas, implementation, or pointers to literature would be very enlightening.

Answer 1

I am the original poster but I lost the account. I found a way to do this in PyTorch and wrote about it on Mathematica.SE here. It absolutely works for some functions like $\sin$ and even $\cos$. Other functions like $\sqrt{\cdot}$ and $\log$ are a problem and fail to converge. Perhaps these functions are fundamentally difficult or the network needs deeper or more exotic layers.

Hopefully somebody with more PyTorch experience can suggest improvements. The code below finds an approximation for the half-iterate of $\sin$ and dumps the results to csv:

import torch
import torch.nn as nn
import torch.optim as optim

from math import pi, sin, cos
import random
import csv


def targetfn(x):
    return sin(x)


class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.lin = nn.Linear(1, 20)
        self.lmid1 = nn.Tanh()
        self.lmid2 = nn.Linear(20, 20)
        self.lmid3 = nn.Tanh()
        self.lout = nn.Linear(20, 1)

    def forward(self, w):
        w = self.lin(w)
        w = self.lmid1(w)
        w = self.lmid2(w)
        w = self.lmid3(w)
        return self.lout(w)


def train():
    net = Net()
    print(net)

    optimizer = optim.SGD(net.parameters(), lr=0.01)
    criterion = nn.MSELoss()

    # init random
    net.zero_grad()
    outinit = net(torch.randn(1))
    outinit.backward(torch.randn(1))

    for i in range(100000):
        x = random.uniform(-2 * pi, 2 * pi)
        target = torch.tensor([targetfn(x)])
        y1 = net(torch.tensor([x]))
        net.zero_grad()
        optimizer.zero_grad()
        y2 = net(y1)
        loss = criterion(y2, target)
        loss.backward()
        optimizer.step()

    return net


def main():
    net = train()

    with open("hfn.csv", 'w', newline='') as csvfile:
        csvwriter = csv.writer(csvfile, delimiter=',')
        n = 2000
        xmin = -2 * pi
        xmax = 2 * pi
        step = (xmax - xmin) / n
        x = xmin
        for i in range(n):
            csvwriter.writerow([x, net(torch.tensor([x])).item()])
            x += step


if __name__ == '__main__':
    main()