How to construct input dependent convolutional filter?

Question

I am constructing a convolutional variational autoencoder for images, starting out with mnist digits. Typically I would specify convolutional layers in the following way:

input_img = layers.Input(shape=(28,28,1))
conv1 = keras.layers.Conv2D(32, (3,3), strides=2, padding='same', activation='relu')(input_img)
conv2 = keras.layers.Conv2D(64, (3,3), strides=2, padding='same', activation='relu')(conv1) 
...

However, I would also like to construct a convolutional filter/kernel that is fixed BUT dependent on some content related to the input, which we can call an auxiliary label. This could be a class label or some other piece of relevant information corresponding to the input. For example, for MNIST I can use the class label as auxiliary information and map the digit to a (3,3) kernel and essentially generate a distinct kernel for each digit. This specific filter/kernel is not learned through the network so it is fixed, but it is class dependent. This filter will then be concatenated with the traditional convolutional filters shown above.

input_img = layers.Input(shape=(28,28,1))
conv1 = keras.layers.Conv2D(32, (3,3), strides=2, padding='same', activation='relu')(input_img)

# TODO: add a filter/kernel that is fixed (not learned by model) but is class label specific
# Not sure how to implement this?
# auxiliary_conv = keras.layers.Conv2D(1, (3,3), strides=2, padding='same', activation='relu')(input_img)

I know there are kernel initializers to specify initial weights https://keras.io/api/layers/initializers/, but I'm not sure if this is relevant and if so, how to make this work with a class specific initialization.

In summary, I want a portion of the model's weights to be input content dependent so that some of the trained model's weights vary based on the auxiliary information such as class label, instead of being completely fixed regardless of the input. Is this even possible to achieve in Keras/Tensorflow? I would appreciate any suggestions or examples to get started with implementation.

Answer 1

Here is one way of achieving this. This network is an autoencoder, with extra auxiliary_convs. The active convolution depends on the input image's class, since each convolution layer's output is multiplied with the one-hot encoded class input.

import matplotlib.pyplot as plt
import numpy as np

from tensorflow import keras
from tensorflow.keras import layers
import tensorflow.keras.backend as K

BN = layers.BatchNormalization

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

# Convert class vectors to binary class matrices
num_classes = y_train.max() + 1
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

# Build the model
inputs = [keras.Input(shape=x_train.shape[1:]), keras.Input(shape=y_train.shape[1:])]
dim = 8

x = inputs[0]
x = BN()(layers.Conv2D(dim, kernel_size=3, activation='relu')(x))
x = BN()(layers.Conv2D(dim, kernel_size=3, activation='relu')(x))

auxiliary_convs = [layers.Conv2D(dim, kernel_size=3, padding='same', activation='relu')
                   for _ in range(num_classes)]

x_auxs = [conv(x) * inputs[1][:,ix:ix+1,None,None]
          for ix, conv in enumerate(auxiliary_convs)]

x = K.sum(K.concatenate([x_aux[:,:,:,:,None] for x_aux in x_auxs]), axis=-1)
x = layers.AveragePooling2D(pool_size=2)(x)

x = BN()(layers.Conv2D(dim, kernel_size=3, activation='elu')(x))
x = layers.AveragePooling2D(pool_size=2)(x)

x = layers.Flatten()(x)
x = layers.Dense(16, activation='elu')(x)
x = layers.Dense(2, activation='tanh')(x)

y = x
y = BN()(layers.Dense(4, activation='elu')(y))
y = BN()(layers.Dense(8, activation='elu')(y))
y = BN()(layers.Dense(16, activation='elu')(y))
y = layers.Dense(np.prod(inputs[0].shape[1:]), activation='sigmoid')(y)
y = layers.Reshape(inputs[0].shape[1:])(y)

model = keras.Model(inputs, y)
model.summary()

It is a bit simpler if you don't sum auxiliary tensors together, but simply concat them:

x = K.concatenate(x_auxs)

However in this case the last Conv2D layer would have redundant parameters to train, since x has dim * num_classes dimensions after the K.concatenate.

If you don't want to train the input-dependent part you can freeze the layers, but I don't know why you would want to do that.

I also tested a variation of this for fun, having a constrained convolutional network but modifying the autoencoded feature based on the image's class:

inputs = [keras.Input(shape=x_train.shape[1:]), keras.Input(shape=y_train.shape[1:])]
act, dim, enc_dim, l2_reg = 'elu', 16, 2, 1e-1

x = inputs[0]
x = BN()(layers.Conv2D(8, kernel_size=3, activation=act)(x))

for _ in range(3):
    x = BN()(layers.Conv2D(dim, kernel_size=3, activation=act)(x))

x = layers.AveragePooling2D(pool_size=2)(x)

for _ in range(3):
    x = BN()(layers.Conv2D(dim, kernel_size=3, activation=act)(x))

x = layers.Flatten()(x)
for k in [6, 8, 4, 3]:
    x = BN()(layers.Dense(k, activation='elu')(x))

x = layers.Dense(enc_dim, activation='linear', use_bias=False,
                 activity_regularizer=keras.regularizers.l2(l2_reg))(x)

a = layers.Dense(enc_dim, activation='exponential', use_bias=False)(inputs[1])
b = layers.Dense(enc_dim, activation='linear')(inputs[1])

x = keras.activations.tanh(a * x + b)

y = x
for k in [3, 4, 8, 32, 128]:
    y = BN()(layers.Dense(k, activation='elu')(y))

y = layers.Dense(np.prod(inputs[0].shape[1:]), activation='sigmoid')(y)
y = layers.Reshape(inputs[0].shape[1:])(y)

model = keras.Model(inputs, y)
model.summary()

The scatter plot on the left shows the codes (aka. embeddings), and the image's class is varied on the decoded examples on the right. Higher the l2_reg is, the tighter the classes' distribution is. Although I don't know what is the utility of this network :D

Answer 2

Not a tensorflow expert but I may be able to offer some conceptual advice. Since you do not care to learn the filter, but instead want to fix a discrete set of possible values for a discrete set of cases, you can use tensor operations (i.e convolutions) rather than neural network layer operations. Essentially, in framework-agnostic pseudocode this would look like:

# layers with learned parameters
output1 = layers1(input)

# apply unlearned but changeable layer convolutions
kernel_val = kernel_val_selection_function(output1)
output2 = convolve_2D(output1,kernel_val)

# more layers with learned parameters
output3 = layers3(output2)

...

The function graph will treat kernel_val as a constant for purposes of backpropagation, so as long as your convolution operations are done within the framework used to create the function graph (i.e. tensorflow) you shouldn't have any problems with backprop.