Producing a confidence output to use in a weighted average layer

Question

I am trying to solve a problem where the input consists of an unordered set of observations (variable size), and the desired output is a single value describing some property of the observations as a whole.

The straight-forward approach would be to train a network to handle an individual observation and fit it to the desired output. When predicting, feed in all observations from the input, and average the result.

This works to an extent, however, the big issue is that this gives all input observations the same weight. For this reason, observations within the input set which do not contain meaningful information (garbage) negatively skew the final average. (For classification problems with one-hot outputs we could use the output's value as the confidence, but this is not the case here.)

My goal is therefore to somehow get the model to itself decide how much meaningful information there is in an observation, and then use this value as the weight of a weighted average of the per-observation outputs to compute the final output value.

To illustrate, here is a function which generates input data (0 or 1) along with a "confidence" value, which indicates the probability that the input data is false:

num_features  =    3 # presence (0 or 1), prediction (0 or 1), and confidence [0,1]
num_timesteps =   64 # N of observations, max number of items in the set
num_samples   = 1024 # training data size

def gen_data():
    for i in range(num_samples):
        result = float(random.getrandbits(1))

        num_populated_timesteps = random.randint(num_timesteps // 2, num_timesteps)

        for j in range(num_timesteps):
            if j < num_populated_timesteps:
                presence = 1
                confidence = random.uniform(0, 1)

                if random.uniform(0, 1) < confidence:
                    # Truth
                    prediction = result
                else:
                    prediction = 1 - result
            else:
                presence = 0
                prediction = 0
                confidence = 0

            inputs.append([presence, prediction, confidence])
        labels.append([result])

One solution I found which does work is simply to use an LSTM layer, which is successfully trained to produce the expected result minding the confidence/presence inputs. However, this is somewhat wasteful as the observations now need to be evaluated sequentially, whereas they really are an unordered set.

I have attempted to solve this as follows:

• Make the per-observation layers output two values, the output and the confidence (which will be used as weight).
• Repeat the per-observation layers for the maximum set size.
• Define a weighted average layer as a lambda layer. (The reason to use a weighted average and not a weighted sum is that the number of observations varies.)
• Use the two outputs from the per-observation layers to compute the weighted average.

Code which attempts to solve the above toy problem:

inputs = np.array(inputs, dtype=(np.float32)).reshape((num_samples, num_timesteps, num_features))
labels = np.array(labels, dtype=(np.float32)).reshape((num_samples, 1))

inputs = np.swapaxes(inputs, 0, 1)
inputs = [x for x in inputs]

L_d = Dense(2, activation='relu')

l_inputs = []
timestep_layers = []
for i in range(num_timesteps):
    l = Input(shape=(num_features))
    l_inputs.append(l)
    l = L_d(l)
    timestep_layers.append(l)

def weighted_average(x):
    weight_sum = reduce_sum([t[0, 1] for t in x])
    return reduce_sum([t[0, 0] * (t[0, 1] / (weight_sum + 0.0001)) for t in x])

l_output = Lambda(weighted_average)(timestep_layers)

model = Model(inputs = l_inputs, outputs = [l_output])

If I set the weights by hand to simply pass the value input as the output value, and the confidence input as the output weight (i.e. [[0,0],[1,0],[0,1]],[0,0]), then the model works as expected and produces correct output. However, actually training the model does not work that well.

• The division in weighted_average causes problems due to the possibility to divide by zero, therefore introduce NaNs. I hacked around this by adding a small value.

• Backpropagation seems to not be working well because even when the weights are set manually, the loss can shoot upwards after a few epochs. I suspect that the division is to blame here, and such a layer would need custom logic for backpropagation.

My question is how to implement a working weighted average layer, or if there is a better solution with a different approach that I have missed.

Answer 1

The code in the question does not work very well because TensorFlow automatic differentiation can't make sense of the weighted_average function.

Though it's possible to define a custom gradient with @tf.custom_gradient, a simpler and more robust solution is to instead use a softmax based weighted average layer:

@tf.function
def weighted_average(x):
    values  = x[:, :, 0]
    weights = x[:, :, 1]
    return tf.math.reduce_sum(values * tf.nn.softmax(weights), 1)

L_lam = Lambda(weighted_average)

This also avoids issues with dividing by zero.