How would you implement weight sharing?

SameDiff
1

I am implementing a famous model where you have in python

with tf.GradientTape() as g:
    for i in tf.range(60):
      x = model(x)
    loss = tf.reduce_mean(loss_f(x))
  grads = g.gradient(loss, model.weights)

Python’s gradient tape will backpropagate through every call to the model. Basically BPTT with only the last output compared to the target label.
In dl4j, I have approximated it by making a ComputationGraph model where all the layers are repeated in order 60 times. I do a forward pass, set the activations, and backprop the gradients. Up to this point it’s like 60 sets of independent weights.
To make this into shared weights, my dirty solution was to add all the gradients, apply them to one set of weights, and copy that value 60 times on the 60 sets of weights.
But the results are very different from the python implementation. Basically I get no convergence past a depth of 8, let alone 60. My guess is that this is not how weight sharing should be implemented, even as a quick and dirty solution.

This is all I found on the topic in the github, but it’s really old and was never solved anyway. Any pointers?

2

@Lana can you clarify why you aren’t just using fit? Why more than one feedforward? The older api is more cumbersome as I pointed out to you. Your other post is using tensorflow as an example can I show you how to use samediff instead?

And yes: on the docs: we indeed DO need to do a pass yet. Feel free to open docs issues as you find them and I’ll correct them both in M2.1 and the new docs I’m working on.

Here’s dl4j:

package org.example;

import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.RNNFormat;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.RnnLossLayer;
import org.deeplearning4j.nn.conf.layers.recurrent.SimpleRnn;
import org.deeplearning4j.nn.gradient.Gradient;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.nd4j.common.primitives.Pair;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.impl.LossMSE;

/**
 * Legacy DL4J equivalent of a shared 60-step transform. A single SimpleRnn
 * layer owns one recurrent weight matrix and reuses it at every time step;
 * BPTT accumulates all time-step contributions into one RW gradient.
 */
public final class LegacyDl4jWeightSharingExample {

    private static final int WIDTH = 8;
    private static final int STEPS = 60;
    private static final int BATCH_SIZE = 32;
    private static final int EPOCHS = 100;

    private LegacyDl4jWeightSharingExample() {
    }

    public static void main(String args) {
        Nd4j.getRandom().setSeed(12345L);
        DataSet data = lastStepOnlyData();

        Result multiLayerResult = runMultiLayerNetwork(data.copy());
        Result graphResult = runComputationGraph(data.copy());

        System.out.println("Executioner: " + Nd4j.getExecutioner().getClass().getName());
        printResult("MultiLayerNetwork", multiLayerResult);
        printResult("ComputationGraph", graphResult);
    }

    private static Result runMultiLayerNetwork(DataSet data) {
        MultiLayerConfiguration configuration = new NeuralNetConfiguration.Builder()
                .seed(12345L)
                .dataType(DataType.FLOAT)
                .updater(new Adam(1e-3))
                .list()
                .layer(new SimpleRnn.Builder()
                        .nIn(WIDTH)
                        .nOut(WIDTH)
                        .activation(Activation.TANH)
                        .build())
                .layer(new RnnLossLayer.Builder(new LossMSE())
                        .activation(Activation.IDENTITY)
                        .build())
                .setInputType(InputType.recurrent(WIDTH, STEPS, RNNFormat.NCW))
                .build();

        MultiLayerNetwork network = new MultiLayerNetwork(configuration);
        network.init();
        initializeRecurrentParameters(
                network.getParam("0_W"),
                network.getParam("0_RW"),
                network.getParam("0_b"));

        Pair<Gradient, INDArray> gradients = network.calculateGradients(
                data.getFeatures(),
                data.getLabels(),
                data.getFeaturesMaskArray(),
                data.getLabelsMaskArray());
        double recurrentGradientNorm = gradients.getFirst()
                .getGradientFor("0_RW")
                .norm2Number()
                .doubleValue();
        double initialScore = network.score(data);

        for (int epoch = 0; epoch < EPOCHS; epoch++) {
            network.fit(data);
        }

        double finalScore = network.score(data);
        long recurrentParameterArrays = network.paramTable().keySet().stream()
                .filter(name -> name.endsWith("_RW"))
                .count();
        validate(
                "MultiLayerNetwork",
                recurrentGradientNorm,
                initialScore,
                finalScore,
                recurrentParameterArrays);
        return new Result(
                recurrentGradientNorm,
                initialScore,
                finalScore,
                recurrentParameterArrays);
    }

    private static Result runComputationGraph(DataSet data) {
        ComputationGraphConfiguration configuration = new NeuralNetConfiguration.Builder()
                .seed(12345L)
                .dataType(DataType.FLOAT)
                .updater(new Adam(1e-3))
                .graphBuilder()
                .addInputs("sequence")
                .addLayer(
                        "sharedRnn",
                        new SimpleRnn.Builder()
                                .nIn(WIDTH)
                                .nOut(WIDTH)
                                .activation(Activation.TANH)
                                .build(),
                        "sequence")
                .addLayer(
                        "loss",
                        new RnnLossLayer.Builder(new LossMSE())
                                .activation(Activation.IDENTITY)
                                .build(),
                        "sharedRnn")
                .setOutputs("loss")
                .setInputTypes(InputType.recurrent(WIDTH, STEPS, RNNFormat.NCW))
                .build();

        ComputationGraph graph = new ComputationGraph(configuration);
        graph.init();
        initializeRecurrentParameters(
                graph.getParam("sharedRnn_W"),
                graph.getParam("sharedRnn_RW"),
                graph.getParam("sharedRnn_b"));

        graph.setInput(0, data.getFeatures());
        graph.setLabel(0, data.getLabels());
        graph.setLayerMaskArrays(
                data.getFeaturesMaskArray() == null
                        ? null
                        : new INDArray[]{data.getFeaturesMaskArray()},
                new INDArray[]{data.getLabelsMaskArray()});
        graph.computeGradientAndScore();
        double recurrentGradientNorm = graph.gradient()
                .getGradientFor("sharedRnn_RW")
                .norm2Number()
                .doubleValue();
        graph.clearLayerMaskArrays();
        double initialScore = graph.score(data);

        for (int epoch = 0; epoch < EPOCHS; epoch++) {
            graph.fit(data);
        }

        double finalScore = graph.score(data);
        long recurrentParameterArrays = graph.paramTable().keySet().stream()
                .filter(name -> name.endsWith("_RW"))
                .count();
        validate(
                "ComputationGraph",
                recurrentGradientNorm,
                initialScore,
                finalScore,
                recurrentParameterArrays);
        return new Result(
                recurrentGradientNorm,
                initialScore,
                finalScore,
                recurrentParameterArrays);
    }

    private static DataSet lastStepOnlyData() {
        INDArray features = Nd4j.zeros(DataType.FLOAT, BATCH_SIZE, WIDTH, STEPS);
        INDArray initialState = Nd4j.randn(DataType.FLOAT, BATCH_SIZE, WIDTH).muli(0.25);
        features.get(
                NDArrayIndex.all(),
                NDArrayIndex.all(),
                NDArrayIndex.point(0)).assign(initialState);

        INDArray labels = Nd4j.zeros(DataType.FLOAT, BATCH_SIZE, WIDTH, STEPS);
        INDArray labelMask = Nd4j.zeros(DataType.FLOAT, BATCH_SIZE, STEPS);
        labelMask.getColumn(STEPS - 1).assign(1.0);
        return new DataSet(features, labels, null, labelMask);
    }

    private static void initializeRecurrentParameters(
            INDArray inputWeights,
            INDArray recurrentWeights,
            INDArray bias) {
        inputWeights.assign(Nd4j.eye(WIDTH).castTo(DataType.FLOAT));
        recurrentWeights.assign(Nd4j.eye(WIDTH).castTo(DataType.FLOAT).muli(0.99));
        bias.assign(0.0);
    }

    private static void validate(
            String name,
            double recurrentGradientNorm,
            double initialScore,
            double finalScore,
            long recurrentParameterArrays) {
        if (recurrentParameterArrays != 1L) {
            throw new IllegalStateException(
                    name + " should have exactly one recurrent parameter array, got "
                            + recurrentParameterArrays);
        }
        if (!Double.isFinite(recurrentGradientNorm) || recurrentGradientNorm <= 0.0) {
            throw new IllegalStateException(
                    name + " recurrent gradient must be finite and non-zero");
        }
        if (!Double.isFinite(finalScore) || finalScore >= initialScore) {
            throw new IllegalStateException(
                    name + " did not reduce its score: " + initialScore + " -> " + finalScore);
        }
    }

    private static void printResult(String name, Result result) {
        System.out.printf(
                "%s: shared RW arrays=%d, RW gradient norm=%.8f, score=%.8f -> %.8f%n",
                name,
                result.recurrentParameterArrays,
                result.recurrentGradientNorm,
                result.initialScore,
                result.finalScore);
    }

    private static final class Result {
        private final double recurrentGradientNorm;
        private final double initialScore;
        private final double finalScore;
        private final long recurrentParameterArrays;

        private Result(
                double recurrentGradientNorm,
                double initialScore,
                double finalScore,
                long recurrentParameterArrays) {
            this.recurrentGradientNorm = recurrentGradientNorm;
            this.initialScore = initialScore;
            this.finalScore = finalScore;
            this.recurrentParameterArrays = recurrentParameterArrays;
        }
    }
}

Here’s samediff which I’d recommend as a parallel to tensorflow:

package org.example;

import java.util.HashMap;
import java.util.Map;

import org.nd4j.autodiff.samediff.SDVariable;
import org.nd4j.autodiff.samediff.SameDiff;
import org.nd4j.autodiff.samediff.TrainingConfig;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.adapter.SingletonDataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;

/**
 * Standalone demonstration of backpropagation through 60 applications of one
 * shared SameDiff parameter set.
 */
public final class SharedWeightsBpttExample {

    private static final int WIDTH = 8;
    private static final int STEPS = 60;
    private static final int BATCH_SIZE = 32;
    private static final int EPOCHS = 100;

    private SharedWeightsBpttExample() {
    }

    public static void main(String args) {
        Nd4j.getRandom().setSeed(12345L);

        SameDiff sd = buildModel(WIDTH, STEPS);
        INDArray features = Nd4j.randn(DataType.FLOAT, BATCH_SIZE, WIDTH).muli(0.25);
        INDArray labels = Nd4j.zeros(DataType.FLOAT, BATCH_SIZE, WIDTH);

        Map<String, INDArray> placeholders = placeholders(features, labels);
        double initialLoss = scalarLoss(sd, placeholders);
        Map<String, INDArray> gradients =
                sd.calculateGradients(placeholders, "weights", "bias");
        double weightGradientNorm = gradients.get("weights").norm2Number().doubleValue();
        double biasGradientNorm = gradients.get("bias").norm2Number().doubleValue();

        sd.fit(
                new SingletonDataSetIterator(new DataSet(features, labels)),
                EPOCHS);

        double finalLoss = scalarLoss(sd, placeholders);

        System.out.println("Executioner: " + Nd4j.getExecutioner().getClass().getName());
        System.out.println("Unrolled steps: " + STEPS);
        System.out.println("Trainable parameter arrays: weights, bias");
        System.out.printf(
                "Initial gradient norms: weights=%.8f, bias=%.8f%n",
                weightGradientNorm,
                biasGradientNorm);
        System.out.printf(
                "Loss: %.8f -> %.8f after %d epochs%n",
                initialLoss,
                finalLoss,
                EPOCHS);

        if (!Double.isFinite(weightGradientNorm) || weightGradientNorm <= 0.0) {
            throw new IllegalStateException(
                    "Shared weight gradient must be finite and non-zero");
        }
        if (!Double.isFinite(biasGradientNorm) || biasGradientNorm <= 0.0) {
            throw new IllegalStateException(
                    "Shared bias gradient must be finite and non-zero");
        }
        if (!Double.isFinite(finalLoss) || finalLoss >= initialLoss) {
            throw new IllegalStateException(
                    "Expected training to reduce loss, but got "
                            + initialLoss
                            + " -> "
                            + finalLoss);
        }
    }

    public static SameDiff buildModel(int width, int steps) {
        SameDiff sd = SameDiff.create();

        SDVariable input = sd.placeHolder("input", DataType.FLOAT, -1, width);
        SDVariable labels = sd.placeHolder("labels", DataType.FLOAT, -1, width);

        // One deterministic, slightly contractive parameter set. The same
        // SDVariables are passed to all 60 applications below.
        SDVariable weights = sd.var(
                "weights",
                Nd4j.eye(width).castTo(DataType.FLOAT).muli(0.99));
        SDVariable bias = sd.var(
                "bias",
                Nd4j.zeros(DataType.FLOAT, 1, width));

        SDVariable state = input;
        for (int step = 0; step < steps; step++) {
            state = modelStep(sd, state, weights, bias, step);
        }

        SDVariable loss = sd.loss().meanSquaredError("loss", labels, state, null);
        loss.markAsLoss();

        sd.setTrainingConfig(new TrainingConfig.Builder()
                .updater(new Adam(1e-3))
                .dataSetFeatureMapping("input")
                .dataSetLabelMapping("labels")
                .build());

        return sd;
    }

    private static SDVariable modelStep(
            SameDiff sd,
            SDVariable input,
            SDVariable weights,
            SDVariable bias,
            int step) {
        SDVariable preActivation = sd.mmul(
                "matmul_" + step,
                input,
                weights).add("pre_activation_" + step, bias);
        return sd.math().tanh("state_" + step, preActivation);
    }

    private static Map<String, INDArray> placeholders(
            INDArray features,
            INDArray labels) {
        Map<String, INDArray> placeholders = new HashMap<>();
        placeholders.put("input", features);
        placeholders.put("labels", labels);
        return placeholders;
    }

    private static double scalarLoss(
            SameDiff sd,
            Map<String, INDArray> placeholders) {
        return sd.output(placeholders, "loss").get("loss").getDouble(0);
    }
}
3

Thank you for your answer! Really appreciate it.

why you aren’t just using fit?

The output layer needs to be a convolutional layer, and initially trying to use “fit” with MultiLayerConfiguration I got an error that you cannot make a convolutional layer be recognized as an output layer. So I gave up on fit…
But also the model is a combination of different layers with several outputs, several inputs, and some frozen layers that still need to let the backprop reach earlier layers. So I can’t just use fit. On the plus side, not being able to use the default implementations has taught me a lot about the theory and practice!

Why more than one feedforward?

The model needs to run recurrently on its own output for 60~100 steps with no external input, and the loss is only calculated at the last iteration. This is the model if you’re curious: Growing Neural Cellular Automata

Feel free to open docs issues as you find them and I’ll correct them both in M2.1 and the new docs I’m working on.

Thanks, I will.

Thank you for the code examples, I will try them out!

4

@agibsonccc I used your samediff example to rebuild the model. It worked. Thanks a lot!

While doing that I filed some documentation issues.

Cheers!