Piecewise Polynomial Layers and Other High Order Layers in PyTorch

This is a PyTorch implementation of my tensorflow repository and is more complete due to the flexibility of PyTorch.

Lagrange Polynomial, Piecewise Lagrange Polynomial, Discontinuous Piecewise Lagrange Polynomial, Fourier Series, sum and product layers in PyTorch. The sparsity of using piecewise polynomial layers means that by adding new segments the representational power of your network increases, but the time to complete a forward step remains constant. Implementation includes simple fully connected layers, convolution layers and deconvolutional layers using these models. This is a PyTorch implementation of this Discontinuous Piecewise Polynomial Neural Networks which was written almost a decade before the recent interest in KAN's, including huge number of extensions including continuous, Fourier series and convolutional neural networks... and many applications with varrying degrees of success. If you come from a computational physics background, this type of approach seems very obvious as ancient techniques like the finite element method are discretized this way.

Collab Notebook

Using simple high order layers Simple function approximation

Using simple high order MLP 2d function approximation

Idea

The idea is extremely simple, instead of a weight at the synapse we have a function F(x) that can be arbitrarily complex. As a practical matter I implement this by using multiple weights corresponding to each link, these weight are used as parameters of the function, and to make sure there is still some GPU efficiency, these weights are just coefficients of the basis functions. In most of this work, the n-weights describe the value of a piecewise polynomial on a regular grid (in the case of a piecewise polynomial) each of the n-weights can be updated independently. A Lagrange polynomial and Gauss Lobatto points are used to minimize oscillations of the polynomial. The same approach can be applied to any "functional" synapse, and I also have Fourier series synapses in this repo as well. Because the non-linearity is applied on the link, the node is simply a summation

In the image below each "link" instead of being a single weight, is a function of both x and a set of weights. These functions can consist of an orthogonal basis functions for efficient approximation.

<img src="plots/NetworkZoom.png" width=50% height=50% style="display: block; margin: 0 auto">

A small layer then looks like this, the values at the nodes are just summed.

<img src="plots/PiecewisePolynomialLayer.svg" width=50% height=50% style="display: block; margin: 0 auto">

A single neuron input output pair with a piecewise function is shown below. In the case where we use polynomials, Lagrange polynomials are being used so the values of the weights are identical to the value of the function at that point. The spacing is determined by chebyshev lobatto points, so there are always weights at the edge of each segment. In the case of discontinuous polynomial, the weights there are 2 weights for each interior segment edge.

<img src="plots/NeuronDrawing.svg" width=50% height=50% style="display: block; margin: 0 auto">

The image below shows the function passing through the weights when using lagrange polynomials. Note that there is no derivative continuity at the boundaries.

<img src="plots/NeuronDrawingWeights.svg" width=50% height=50% style="display: block; margin: 0 auto">

Why

Using higher order polynomial representations allow networks with much fewer total weights in certain cases. There is a well known phenomena in numerical modeling known as exponential convergence using spectral methods when using hp refinement, it's possible something like that can happen in neural networks as well.

Is this a KAN?

Actually a single layer piecewise polynomial KAN (which is actually 2 layers) is a special case of a 2 layer piecewise polynomial network, which is used in this repo. Therefore, a piecewise polynomial layer is actually "Half a KAN" so it's actually simpler - Often all you need is a single polynomial layer at the input followed by a standard MLP so having the piecewise polynomial layer is important. Other names that have been used in the past Deep FLANN (functional link artificial neural network).

Lagrange polynomials are widely used in finite element analysis and have the advantage that the value of the weight is actually the value of the function at that point in space. By limiting the weights you are limiting the maximum value of the function (the function may be higher than the weights in between the nodes). Also, when you go beyond the range of definition [-1,1] the polynomial is still defined using the last (or first) polynomial in the sequence, whether you want it defined that way at high polynomial order is another question. I mention a paper at the bottom where they do a linear extension beyond the range [-1,1] so values do not rise too fast - but normalization works as well.

Issues

What about instabilities due to steep gradients? Seems like you can get around those with various approaches, polynomial refinement is one (start with piecewise linear and than increase the polynomial order after it converges), the lion optimizer helps a lot as well, while sophia may be even better since it's second order.

The biggest issues I've experienced though are that it's slower than dense networks and certain operations can take up more memory which can cause major issues with models that already push the limits of your gpu. Now that KANs are popular, hopefully there will be enough people to address all these issues.

In general, with enough effort, it seems I can make them "work" for any place the classic ReLU network works and in certain situations they clearly work much better. They also do a great job of overfitting, which just means, I need more data. For problems where your inputs are positional, x and y..., they seem to be far better.

Finally, I believe these methods actually will benefit much more from (approximate) second order optimizers. I used those in my original implementation. Although there are plenty of second order optimizers out there, to date, pytorch does not have a standard one except LBFGS which has its own issues.

Fully Connected Layer Types

All polynomials are Lagrange polynomials with Chebyshev interpolation points.

A helper function is provided in selecting and switching between these layers

from high_order_layers_torch.layers import *
layer1 = high_order_fc_layers(
    layer_type=layer_type,
    n=n,
    in_features=784,
    out_features=100,
    segments=segments,
)

where layer_type is one of | layer_type | representation |--------------------|-------------------------| |continuous | piecewise polynomial using sum at the neuron | |continuous_prod | piecewise polynomial using products at the neuron | |discontinuous | discontinuous piecewise polynomial with sum at the neuron| |discontinuous_prod | discontinous piecewise polynomial with product at the neuron| |polynomial | single polynomial (non piecewise) with sum at the neuron| |polynomial_prod | single polynomial (non piecewise) with product at the neuron| |product | Product | |fourier | fourier series with sum at the neuron |

n is the number of interpolation points per segment for polynomials or the number of frequencies for fourier series, segments is the number of segments for piecewise polynomials, alpha is used in product layers and when set to 1 keeps the linear part of the product, when set to 0 it subtracts the linear part from the product.

Convolutional Layer Types

conv_layer = high_order_convolution_layers(layer_type=layer_type, n=n, in_channels=3, out_channels=6, kernel_size=5, segments=segments, rescale_output=rescale_output, periodicity=periodicity)

All polynomials are Lagrange polynomials with Chebyshev interpolation points. | layer_type | representation | |--------------|----------------------| |continuous(1d,2d) | piecewise continuous polynomial |discontinuous(1d,2d) | piecewise discontinuous polynomial |polynomial(1d,2d) | single polynomial |fourier(1d,2d) | fourier series convolution

Initializing of layers

The default initialization is to initialize each link to a random constant, i.e. all weights have the same value in a link. This seems to work pretty well, however, I also have linear random linear initialization (non constant). The implementation of the linear initialization is slower and I'm not sure it's actually better.

Here is a function that does this linear initialization for non convolutional layers (it can be found in networks.py)

def initialize_network_polynomial_layers(
    network: nn.Module,
    max_slope: float,
    max_offset: float,
    scale_slope: Callable[[float], float] = lambda input_size: 1,
)

h and p refinement

p refinement is taking an existing network and increasing the polynomial order of that network without changing the network output. This allow the user to train a network at low polynomial order and then use that same network to initialize a network with higher polynomial order. This is particularly useful since a high order polynomial network will often converge poorly without the right initialization,