Future Design Extensions to the Neural Network Framework

Matteo Salvarezza, Wolfram Research

#WolframTechConf

ABSTRACT

As the Neural Network framework is being migrated to a PyTorch backend, its underlying eager execution will unlock important features. We’ll explore how this new software stack will allow combining of symbolic network graphs, imperative layer execution and native Wolfram Language functions, making it easy to mix custom Wolfram Language code with PyTorch-powered layers and pre-trained models.

OUTLINE

Current design of the Neural Network framework

Static Layers & Containers

Integration with the rest of the language

Possible future extensions

Dynamic shapes

Dynamic branching & loops

Seamless language integration via Compiler and AutoDifferentiation

Higher order gradients

Physics-informed neural networks

Bonus: Performance improvements

Current design of the Neural Network framework

Static Layers & Containers

Neural network layers and containers need to have a static input shape set in order to run:

NetInitialize@LinearLayer[2,"Input"->{3}]

NetInitialize@NetChain[{LinearLayer[2],ElementwiseLayer[Ramp]},"Input"->{3}]

This property ultimately stems from using a static backend (MXNet) for computations. The net is translated to a static MXNet computational graph, which is compiled and later ran in a single MXNet call.

Running layers with unspecified shapes gives the illusion this is not true, but they silently set their shape under the hood when data is provided:

unspecLayer=ElementwiseLayer[Tanh]

unspecLayer[-1.4]

For this reason there’s a noticeable performance cost when running them:

specLayer=ElementwiseLayer[Tanh,"Input"->"Real"]

unspecLayer[-1.4]//RepeatedTiming
specLayer[-1.4]//RepeatedTiming

The first dimension can in some cases be set to “Varying”, which again gives the illusion of a non-static framework:

mappedLayer=NetInitialize@NetMapOperator[LinearLayer[2],"Input"->{"Varying",3}]

mappedLayer@{{1,2,3}}
mappedLayer@{{1,2,3},{4,5,6}}
mappedLayer@{{1,2,3},{4,5,6},{7,8,9}}

This is realized through a complex orchestration of static MXNet graphs, each operating on a fixed sequence length, sharing memory with each other.

Performance is similar in this case, thanks to caching of the static graphs:

fixedLayer=NetInitialize@NetMapOperator[LinearLayer[2],"Input"->{3,3}]

mappedLayer@{{1,2,3},{4,5,6},{7,8,9}}//RepeatedTiming
fixedLayer@{{1,2,3},{4,5,6},{7,8,9}}//RepeatedTiming

Current design of the Neural Network framework

Integration with the rest of the language

FunctionLayer takes a function that mixes ordinary WL code and neural network layers and turns it into a neural network layer which can be plugged into any net:

function=Function@Module[{x,y,z,k},
x=Tanh@LinearLayer[{10,10}][Ramp[#Input1]];
x=Ramp[x];
y=Cos[x]+Exp[#Input2/2]+#Input1;
x=AggregationLayer[Max][x];
z=Tanh[y];
k=Sin[x+1];
<|"Output1"->x,"Output2"->z|>
];

funcLayer=FunctionLayer[function]

NetGraph[{funcLayer,LinearLayer[20]},{1->2}]

It works by analyzing the function body and creating an equivalent NetGraph

funcLayer

NetGraph[funcLayer]

But it also has a few primitives which don’t translate to other layers:

FunctionLayer[Det]

FunctionLayer[Dimensions]

CompiledLayer takes a Function[] object (complete with compiler type annotations) and essentially uses FunctionCompile to produce low-level compiled code which can be plugged directly into the MXNet graph:

function=Function[{Typed[input,TypeSpecifier["NumericArray"]["Real32",1]]},
Tanh[input^2]*(Floor[input]-Mean[input])
];

c=CompiledLayer[function,"Input"->10]

NetInitialize@NetChain[{LinearLayer[10],c},"Input"->6]

Gradient computation and thus training are possible via the compiler auto differentiation capabilities:

c[Range[10],NetPortGradient["Input"]]

Possible future extensions

Dynamic shapes

A in a fully dynamic backend like PyTorch (which runs operators eagerly without pre-compiling a computational graph) layers don’t need to have a static input shape set before they can run, and layers with unspecified shapes an be ready to be evaluated efficiently:

ElementwiseLayer[Ramp]

Another similar feature is “runtime” reshaping and transposing, where the specification are inputs to the layers instead of static parameters:

ReshapeLayer[{2,4,3},"Input"->{12,2}]
ReshapeLayer["ReferenceShape","Input"->{12,2},"ReferenceShape"->{2,4,3}]
TransposeLayer[{2,1,3},"Input"->{7,2,4}]

The “runtime” versions would look like this (using PlaceholderLayer to represent them):

runtimeReshape=PlaceholderLayer["RuntimeReshape",InputPorts->{"Input"->Automatic,"Shape"->"Varying"}]
runtimeTranspose=PlaceholderLayer["RuntimeReshape",InputPorts->{"Input"->Automatic,"TransposeSpec"->"Varying"}]

Possible future extensions

Dynamic branching & loops

A fully dynamic backend can branch layer execution depending on the data being processed:

FunctionLayer @ Function[x,
	Module[{y, z},
		y = f[x];
		z = If[Total[y, Infinity] > 0,
			LinearLayer[2][y],
			LinearLayer[3][y]
		];
		...
	]
]

Loops with a number of iteration which depends on the input data are also a dynamic-only feature:

FunctionLayer @ Function[x,
	Module[{y, n},
		y = f[x];
		n = h[y];
		Do[y = g[y], n];
		...
	]
]

FunctionLayer @ Function[x,
	Module[{y},
		y = f[x];
		Do[
			y = g[y];
			If[h[y] < 0, Break[]], 
			maxIterations
		];
		...
	]
]

Possible future extensions

Seamless language integration via Compiler and AutoDifferentiation

One can imagine to perform eager execution of layers mixed with WL code that’s JIT-compiled as needed, enabling gradient computation via auto differentiation in the compiler (although GPU support would require a large effort). The result would be a mix of FunctionLayer and CompiledLayer, where one can define a net computation by writing plain WL code and not being limited to something that needs to translate every statement to a neural net primitive like FunctionLayer and can just run eagerly:

NetFunction[input,
Module[{h},
h=Ramp@ConvolutionLayer[16,{3,3}][input];
h=Ramp@EdgeDetect[h];
LinearLayer[256][h]
]
]

Possible future extensions

Higher order gradients

PyTorch supports arbitrarily higher order gradients which could be either requested from NetPortGradient or even computed inside a net. The neural net framework currently only supports first-order gradients to be returned by net evaluation or used in training:

net=NetInitialize@NetChain[{LinearLayer[2],Tanh},"Input"->{3}]

net[{1,2,3},NetPortGradient[{"Input"}]]

Higher order gradients could be computed in net evaluation by extending the syntax of NetPortGradient:

net[{1,2,3},NetPortGradient[{"Input"},2]]

More importantly, gradient could be computed as part of the forward pass through the net with a dedicated future layer. Let’s use PlaceholderLayer to represent it:

gradLayer=PlaceholderLayer["GradientLayer",<|"Order"->3|>,
InputPorts->{"Numerator"->"Real","Denominator"->{3}},
OutputPorts->{"Output"->{3}}
]

This net would compute the 3rd order gradient of the output produced by ElementwiseLayer[“Swish”] w.r.t. the input of the net, in the forward pass:

net=NetInitialize@NetGraph[
{LinearLayer[3],ElementwiseLayer["Swish"],gradLayer},
{NetPort["Input"]->{1,3},1->2->NetPort[3,"Denominator"]}
]

Backpropagating through the same net using NetPortGradient[{“Input”}] would produce the 4th order gradient. Using NetPortGradient[{“Input”}, n] would produce a gradient of order 3+n

Possible future extensions

Physics-informed Neural Networks

GradientLayer would be used to create the loss function for physics-informed neural networks, or more generally to train nets to solve differential equations.

Consider a damped harmonic oscillator:

m(((∂^2)x(t))/(∂(t^2)))+c((∂x(t))/(∂t))+k x(t)=0
x(0)=(Subscript[x, 0])

And model x(t) as a neural network, which takes the time t as an input and outputs the position x:

net=NetInitialize@NetChain[{10,Tanh,10,Tanh,{}},"Input"->{}]

We can train the net so that it becomes the solution of the above ODE by enforcing it in the loss function:

ℒ=((|m(((∂^2)x(t))/(∂(t^2)))+c((∂x(t))/(∂t))+k x(t)|)^2)+α((|x(0)-(Subscript[x, 0])|)^2)

The loss function implementation can be framed as a NetGraph which takes the output of the net x(t), the time t and the initial position as predicted by the net, x(0). This graph uses GradientLayer to compute the gradients:

(* physical constants *)
{m,k,c}={5,2.1,3.4};
(* initial condition *)
x0=-4.6;
(* loss weighting *)
alpha=4;

gradientLayer[order_]:=PlaceholderLayer["GradientLayer",<|"Order"->order|>,
InputPorts->{"Numerator"->"Real","Denominator"->"Real"},
OutputPorts->{"Output"->"Real"}
]

loss=NetGraph[
<|
"Velocity"->gradientLayer[1],
"Acceleration"->gradientLayer[1],
"ODE"->ThreadingLayer[(m*#1+c*#2+k*#3)^2&],
"InitialCondition"->ElementwiseLayer[(#1-x0)^2&]
|>,
{
{NetPort["Position"],NetPort["Time"]}->"Velocity",
{"Velocity",NetPort["Time"]}->"Acceleration",
{"Acceleration","Velocity",NetPort["Position"]}->"ODE"->NetPort["ODELoss"],
NetPort["InitialPosition"]->"InitialCondition"->NetPort["InitialConditionLoss"]
},
"Position"->"Real","InitialPosition"->"Real","Time"->"Real"
]

The full training architecture features two copies of the net sharing the weights, one computing x(t) and the other x(0):

net=NetInsertSharedArrays[net]

trainingNet=NetGraph[
{NetArrayLayer["Array"->0],net,net,loss},
{
1->2->NetPort[4,"InitialPosition"],
NetPort["Time"]->3->NetPort[4,"Position"]
}
]

The training data is just a list of time values, and the net will be trained using the desired weighted combination of losses:

times=Range[1,10,0.01]

(* fails because GradientLayer is not real yet, it's just a placeholder *)
NetTrain[trainingNet,times,LossFunction->{"ODELoss","InitialConditionLoss"->Scaled[alpha]}]

Bonus: performance improvements

See my talk on the 7th for more in depth coverage

Native GPU support on macOS (Metal)

Low precision types (16-bit floats, 8-bit quantized types...)