Keras
deep learning
neural networks
model merging
machine learning

Keras weighted merge

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

A weighted merge in Keras means combining multiple tensors with explicit scalar or tensor weights instead of using a plain Add, Average, or Concatenate. The key design choice is whether those weights are fixed constants or trainable parameters. Keras does not expose a one-line built-in “weighted merge” layer for every case, but the functionality is easy to build with either simple tensor math or a custom layer.

Fixed Weighted Merge

If the weights are known in advance, the merge is just weighted addition.

python
1import tensorflow as tf
2from tensorflow import keras
3
4input_a = keras.Input(shape=(4,))
5input_b = keras.Input(shape=(4,))
6
7merged = 0.7 * input_a + 0.3 * input_b
8model = keras.Model(inputs=[input_a, input_b], outputs=merged)
9
10x1 = tf.constant([[1.0, 2.0, 3.0, 4.0]])
11x2 = tf.constant([[10.0, 20.0, 30.0, 40.0]])
12print(model([x1, x2]).numpy())

This is the simplest answer when the blend ratios are part of the design rather than something the network should learn.

Trainable Weighted Merge

If the model should learn how much each branch matters, a custom layer is the cleanest approach.

python
1import tensorflow as tf
2from tensorflow import keras
3
4class WeightedMerge(keras.layers.Layer):
5    def build(self, input_shape):
6        self.alpha = self.add_weight(
7            name="alpha",
8            shape=(),
9            initializer="zeros",
10            trainable=True,
11        )
12
13    def call(self, inputs):
14        a, b = inputs
15        weight = tf.sigmoid(self.alpha)
16        return weight * a + (1.0 - weight) * b
17
18input_a = keras.Input(shape=(4,))
19input_b = keras.Input(shape=(4,))
20merged = WeightedMerge()([input_a, input_b])
21model = keras.Model(inputs=[input_a, input_b], outputs=merged)

Using sigmoid constrains the learned mixing coefficient into the [0, 1] range, which makes the merge easier to interpret.

Why Constraining The Weight Helps

If you use an unconstrained trainable scalar directly, the layer can amplify or invert signals unexpectedly. Sometimes that is fine, but if your intention is a true convex blend between two sources, constraining the weight is better.

That is why sigmoid(alpha) is a common design choice for a two-branch merge.

More Than Two Inputs

For more than two inputs, the usual pattern is to learn a vector of logits and normalize them with softmax.

python
1class MultiWeightedMerge(keras.layers.Layer):
2    def build(self, input_shape):
3        self.logits = self.add_weight(
4            name="logits",
5            shape=(len(input_shape),),
6            initializer="zeros",
7            trainable=True,
8        )
9
10    def call(self, inputs):
11        weights = tf.nn.softmax(self.logits)
12        output = 0
13        for i, tensor in enumerate(inputs):
14            output = output + weights[i] * tensor
15        return output

This produces trainable weights that sum to 1, which is often what people mean by a weighted merge.

Shape Compatibility Matters

Weighted addition requires compatible tensor shapes. If branch outputs differ, you need to project them to a compatible shape before merging.

For example, if one branch ends with 64 features and another with 128, you cannot directly add them. You might first use dense layers to map both to the same width.

Weighted Merge Versus Attention

A weighted merge is not automatically the same as attention. Attention usually computes data-dependent weights that vary per example or per token. A trainable weighted merge layer often learns one global blending rule shared across the dataset.

That distinction matters if your model needs context-sensitive routing rather than a fixed learned blend.

Common Pitfalls

  • Expecting a built-in Keras layer to cover every trainable weighted merge case automatically.
  • Forgetting to align tensor shapes before weighted addition.
  • Using unconstrained trainable weights when the intended behavior is a bounded interpolation.
  • Calling a global trainable blend “attention” even when the weight does not depend on the input.
  • Overengineering the merge when a simple Add or Concatenate would be enough.

Summary

  • A weighted merge in Keras is usually just weighted tensor addition.
  • Fixed weights can be implemented directly with ordinary tensor math.
  • Trainable weighted merges are best expressed with a custom layer.
  • Use sigmoid or softmax when you want interpretable constrained weights.
  • Make sure branch outputs have compatible shapes before merging.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.