Variational Autoencoder (VAE) latent features

Question

I'm new to DL and I'm working on VAE for biomedical images. I need to extract relevant features from ct scan. So I created first an autoencoder and after a VAE. My doubt is that I don't know from which layer I can extract feautures. My personal idea is to use features extracted by layers that compute the mean and variance (before reparameterization trick), but I think that also the layer before these is suitable for the purpose. I left here code of the encoder part:

class Sampling(tf.keras.layers.Layer):
    """Uses (z_mean, z_log_var) to sample z, the vector encoding a digit."""
    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.keras.backend.random_normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

def Encoder():
    inp = tf.keras.Input(shape=(32,256,256,1)) # prima era 64

    #enc = tf.keras.layers.Conv3D(8, (2,2,2), activation = 'relu', padding = 'same')(inp)
    #enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(16, (2,2,2), activation = 'relu', padding = 'same')(inp)
    enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(32, (2,2,2), activation = 'relu', padding = 'same')(enc)
    enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same')(enc)
 
    enc = tf.keras.layers.Conv3D(64, (2,2,2), activation = 'relu', padding = 'same')(enc)
    enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)

    enc = tf.keras.layers.Conv3D(32, (2,2,2), activation = 'relu', padding = 'same')(enc)
    enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)
    #enc = tf.keras.layers.Flatten()(enc)
    enc = tf.keras.layers.Conv3D(16, (2,2,2), activation = 'relu', padding = 'same')(enc)
    enc = tf.keras.layers.MaxPooling3D((2,2,2), padding = 'same') (enc)
    '''
    # conv 2D 
    code = tf.keras.layers.Reshape((8,8,96)) (enc)
    code = tf.keras.layers.Conv2D(96,(2,2), activation = 'relu', padding = 'same')(code)
    code = tf.keras.layers.MaxPooling2D((2,2), padding = 'same') (code)
    '''
    
    # latentent code vae
    latent_code = tf.keras.layers.Flatten()(enc)
    latent_code = tf.keras.layers.Dense(256, activation='relu')(latent_code)
    latent_mu = tf.keras.layers.Dense(32, activation='relu')(latent_code) # èprima era 10
    latent_sigma = tf.keras.layers.Dense(32, activation='relu')(latent_code) # prima era 10
    # Reparameterization trick
    #z = tf.keras.layers.Lambda(sample_z, output_shape=(128,), name='z')([latent_mu, latent_sigma])
    z = Sampling()([latent_mu, latent_sigma])
    encoder = tf.keras.Model(inp, [latent_mu, latent_sigma, z ], name = 'encoder')
    
    #encoder = tf.keras.Model(inp, enc)#[latent_mu, latent_sigma, z ], name = 'encoder')
    return encoder
```

Answer 1

The earlier layers of neural network learn more low-level features, while deeper layers learn more complicated, abstract features (see figure from Albawi et al, 2017).

Autoencoders are build of two networks encoder that encodes the data in terms of some latent variables (usually of lower dimensionality, hence they can be used for dimensionality-reduction) and decoder that transforms the latent representation back into the initial representation of the data.

You can use any layer from autoencoder, depending on your needs. Since autoencoders are usually symmetric, usually for feature generation you would be taking layers from the encoder, or it's output. If you want to use autoencoder for dimensionality reduction, then you would rather take the output of the encoder that has smaller dimension (see figure below from the paper by Windrim et al, 2019).

On another hand, you could use autoencoder in a same way as any other neural network for transfer learning. In such case, you would train autoencoder and then extract some layer of it as features for other algorithm. In such case, if you wanted lower-level features, you would take earlier layers. If you wanted more abstract features, you would take higher layers.

Using the example from first picture, you would take first layer of network trained on human faces and use it for extracting some basic shapes and edges from images other than human faces. On another hand, if you needed to detect faces, you would take deeper layers of the network.

Answer 2

Adding to the elaborate answer of @Tim:

VAE z latent is stochastic z, meaning samples will be different for a same $x_i$ sample. In the ideal case you latent representation ($\mu$ or z) will contain meaningful information, these are the ones I would extract (btw in tensorflow you can extract multiple layers ;) ). $\sigma$ is established to act as a noise component.

To make sure that what you extract is useful what you can do is an interpretability test.

The Ct-scans input features are the pixels you have, BUT is there some other information you are not using ? For example since you don't give explicit labels do you have some a scan image of a sick patient ? Or could you select 10 images by hand with some specific feature to interpret a bit what neurons are triggered in the latent space? If so what you can do is a correlation test. This can be as simple as neyman pearson or a 2d histogram showing how correlated features are. What you want to achieve in this case is some sense of what is being used from the model to decide.

Unfortunately this is easy for cats/docs and harder for more complex datasets, but it's something that you need to do to not have a black box machine. Good luck!