# How would you implement weight sharing?

**URL:** https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537
**Category:** SameDiff
**Created:** [August 22, 2026, 12:49pm UTC](https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537 "2026-08-22T12:49:20Z")
**Posts on this page:** 4
**Page:** 1

<div class="post-metadata">

### Author: ![Lana](https://avatars.discourse-cdn.com/v4/letter/l/d2c977/32.png) [@Lana](https://community.konduit.ai/u/Lana)
#### Post date: [August 22, 2026, 12:49pm UTC](https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537/1 "2026-08-22T12:49:20Z")

</div>

I am implementing a famous model where you have in python

```auto
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?

> <https://github.com/deeplearning4j/deeplearning4j/issues/7408>
>
> For implementing more complex multi-output layers, such as GANs or adversarial a…utoencoders (an autoencoder with a discriminator on the reconstruction), it would be nice to have an easier way to share weights between layers in different models.
> 
> As per #791, the current way to implement GANs is to train the discriminator, then copy the weights into the GAN, then train the GAN to fool the discriminator and copy the weights into the generator. This is very clunky.
> 
> Contrast this with Keras's implantation of shared networks, where you simply create 3 models from the same layers, and can train each independently.
> 
> The main thing that makes this easier is that in Keras, if you use the same layer object in two different models, the weights are shared.
> 
> From what I've looked at, a possible implementation is to give every layer and vertex a "weight pool" object, that \`paramTable\` and \`params\` interact with, that is basically a wrapper for any parameters. This can be created and shared by the layer configs, and passed down to the implementations.
> 
> Then you could have a method like \`(config)Layer.copyAndShareWeights()\` which returns a new config with the same weight pool object.
> 
> Also, you could have things like getting a subset of a network with shared weights.
> 
> It would most likely have to interact with memory workspaces, distributed training, and serialization (probably just index pools).

---

<div class="post-metadata">

### Author: ![agibsonccc](https://yyz1.discourse-cdn.com/flex035/user_avatar/community.konduit.ai/agibsonccc/32/697_2.png) [@agibsonccc](https://community.konduit.ai/u/agibsonccc)
#### Post date: [August 23, 2026, 10:15am UTC](https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537/2 "2026-08-23T10:15:58Z")

</div>

@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:

```auto
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:

```auto
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);
    }
}

```

---

<div class="post-metadata">

### Author: ![Lana](https://avatars.discourse-cdn.com/v4/letter/l/d2c977/32.png) [@Lana](https://community.konduit.ai/u/Lana)
#### Post date: [August 24, 2026, 12:03pm UTC](https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537/3 "2026-08-24T12:03:51Z")

</div>

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](https://distill.pub/2020/growing-ca/)

> 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!

---

<div class="post-metadata">

### Author: ![Lana](https://avatars.discourse-cdn.com/v4/letter/l/d2c977/32.png) [@Lana](https://community.konduit.ai/u/Lana)
#### Post date: [August 26, 2026, 10:44am UTC](https://community.konduit.ai/t/how-would-you-implement-weight-sharing/3537/4 "2026-08-26T10:44:07Z")

</div>

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

While doing that I filed some documentation issues.

Cheers!