Gradient Origin Network

This notebook is an attempt to implement this work.

Data gathering and processing

We’ll use the training data from the MNIST dataset.

mnist=ResourceObject["MNIST"];images=Keys@ResourceData[mnist,"TrainingData"];

For testing, we’ll only work with a small number of images.

samples=RandomChoice[images,5]





{width,height}=ImageDimensions@RandomChoice@images;

We create a long list of {i,j} pixel coordinates. Hopefully this pixel ordering is consistent with the built-in function ImageData.

indices = Tuples[{Range[width],Range[height]}];

Creating the input space

We first define a desired latent space dimension. We pick 32 here as in the original paper.

latentSpaceSize=32;

The latent space is a scaled version of the pixel space, concatenated with the latent space.I picked

as a scale factor for various reasons.

input[c_,l_]:=

c~Join~l

input[c_]:=input[c,ConstantArray[0,latentSpaceSize]]

SIREN showcase

We first demonstrate how SIRENs work. To do that, we first pick an image in our dataset.

testimage=RandomChoice@images

A SIREN is basically a MLP with sine activation functions:

siren=NetInitialize@NetChain[{width,Sin,width,Sin,width,1},"Input"2]

NetChain

	Input	vector (size: 2)
1	LinearLayer	vector (size: 28)
2	Sin	vector (size: 28)
3	LinearLayer	vector (size: 28)
4	Sin	vector (size: 28)
5	LinearLayer	vector (size: 28)
6	LinearLayer	vector (size: 1)
	Output	vector (size: 1)

Null

FE`NeuralNetworks`Private`NetChain`part$$1547497692814685212763562077646997957074



Here we don’t use the latent space yet, so we can train the network straight away. We still have to scale both the input and the output, though.

trainedNet=NetTrainsiren,

indicesPartition[Flatten[2#-1&/@ImageData[testimage]],1];

We then reconstruct the image, comparing the reconstruction with the original :

testimageImage(#+1)2&@PartitiontrainedNet/@

indices,width



If the two images look similar, then the pixel ordering is consistent with ImageData and there was sufficient training.

Building the initial network

We’ll use the same network structure as in the paper : four hidden layers with 256 features each.

net=NetChain[256//{#,Sin,#,Sin,#,Sin,#,Sin,1}&,"Input"2+latentSpaceSize]//NetInitialize

NetChain

	Input	vector (size: 34)
1	LinearLayer	vector (size: 256)
2	Sin	vector (size: 256)
3	LinearLayer	vector (size: 256)
4	Sin	vector (size: 256)
5	LinearLayer	vector (size: 256)
6	Sin	vector (size: 256)
7	LinearLayer	vector (size: 256)
8	Sin	vector (size: 256)
9	LinearLayer	vector (size: 1)
	Output	vector (size: 1)

Null

FE`NeuralNetworks`Private`NetChain`part$$3174178607452633547859483776475157337520



Training

This part is quite messy and difficult to explain, but it seems to work.

G=Table[net[input[c],NetPortGradient["Input"]][[3;;]],{c,indices}];F=Flatten@Table[net[input[c]],{c,indices}];d=Table[F-Flatten[2#-1&/@ImageData[img]],{img,samples}];dG=d.G;

trainedNet=NetTrain[net,(Flatten[Outer[input[#2,#1]&,dG,indices,1,1],1]Partition[Flatten@Table[Flatten[2#-1&/@ImageData[img]],{img,samples}],1])]

NetChain

	Input	vector (size: 34)
1	LinearLayer	vector (size: 256)
2	Sin	vector (size: 256)
3	LinearLayer	vector (size: 256)
4	Sin	vector (size: 256)
5	LinearLayer	vector (size: 256)
6	Sin	vector (size: 256)
7	LinearLayer	vector (size: 256)
8	Sin	vector (size: 256)
9	LinearLayer	vector (size: 1)
	Output	vector (size: 1)

Null

FE`NeuralNetworks`Private`NetChain`part$$5420193645325968557997852523203560707771



Table[Image[(#+1)/2&@Partition[trainedNet/@(input[#,dG[[i]]]&/@indices),width]],{i,Range@Length@samples}]





Image[(#+1)/2&@Partition[trainedNet/@(input[#,(dG[[1]]+dG[[2]])/2]&/@indices),width]]