Is it possible to learn the number of layers?

Question

Is it possible, in a transformer or other deep architecture, to include the number of layers as a parameter of the model so it could be learned?

In fact, I have a keras layer that I use to change the final layer without rebuilding the model, so I can just change a parameter between epochs (The original use was to try to train deep networks starting from shallower ones, increasing the number of layers after each epoch).

class LayerSelect(tf.keras.layers.Layer):
    def __init__(self,nlevels,**kwargs):
        super(LayerSelect,self).__init__(**kwargs)
        self.nlevels = nlevels
        self.range=tf.range(self.nlevels,dtype=tf.float32)
                 
    def build(self, input_shape):
        self.kernel=self.add_weight(shape=(1,),
                                    initializer=tf.keras.initializers.Constant(min(self.nlevels,14.0)/1.9),
                                    trainable=True, dtype=tf.float32,
                                    constraint=lambda x: tf.clip_by_value(x,1.0,self.nlevels))
       
    def call(self,inputs): 
        selector=tf.math.maximum([0.0], 1.0 - 1.0 *(self.range-self.kernel)**2 ) 
        final=tf.reduce_sum(inputs*selector,axis=-1)
        return final

The layer expects an stack of hidden layers to choose from:

allEncoders=tf.stack([encoder[level] for level in range(layers)],axis=-1)
finalEncoderRaw=adhoc.LayerSelect(layers)(allEncoders)

So that by calling set_weights during the training I can choose as output any layer, or a combination of two, being the layer variable a float and using a wider selector, say 1.0 - 0.25 *(self.range-self.kernel)**2

And as you can expect, if I set the weight to be trainable, the optimiser moves the variable. But it keeps either moving randomly some small percent or moving backwards towards smaller values. So it is possible that this approach is a dead end?

If not a way to patch this method, is there another successful method to train the number of layers without using meta-parameter (hyperparameter grids) farms?

Answer 1

I like the idea, but I fear this approach may be a dead end. I see a few problems:

Layers in front of (closer to the output than) the currently selected layer(s) don't affect the output, so they won't change and they can't learn to be good predictors for the true output.

Layers behind the currently selected layer won't be trained to predict the final output (they'll be trained to provide outputs are useful inputs for the selected layer) so switching to those layers is unlikely to improve the output.

And if another layer did provide a better approximation of the final output, the network is unlikely to switch to it unless it's next to the currently selected layer (i.e. the selection is likely to get stuck in a local minimum). For example, if layer 3 is selected, and 1 is a better prediction than 3, but 3 is better than 2, the network would stick with 3.

is there another successful method to train the number of layers without using meta-parameter farms?

I haven't heard of one, and that seems like something that would be widely shared if it worked well.

I have seen genetic algorithms for selecting neural network architecture, but as far as I know they don't perform better than grid search for choosing the number of layers.

Answer 2

It is always possible to use a Dense layer to allow the network to built its own menu of layers

allEncoders=tf.stack([encoder[level] for level in range(1+nlayers)],axis=-1)
layerSelect=tf.keras.layers.Dense(1,activation=None, use_bias=False,
    kernel_initializer=tf.keras.initializers.Constant(0.5),
    kernel_constraint=tf.keras.constraints.MinMaxNorm(axis=1))
almostFinalEncoder=tf.keras.layers.Reshape([-1,layerUnits])(layerSelect(allEncoders))
finalEncoder=tf.keras.layers.LayerNormalization()(almostFinalEncoder)

but one must be careful with the initializer, and the vector layerSelect.get_weights()[0]needs to be monitorised.

Generically this layer will not converge to a Kronecker delta, and it must keep some extra weights open just for the sake of overfitting.

Still, the evolution seems to have some information. See here the contourplot for a 64 layers, 256 units, transformer that has reached overfitting at epoch 18 but keeps learning slowly. It seems to prefer to increase weights of layers 1 to 24, and it moves as the learning progresses, decreasing the use of the first layers.

(Sorry the colour scheme is confusing. Basically we have a descending trend in the first layers, ascending in the middle layer, then descending again in the tail)

Generically most plots show a trend to zero the two first layers (or not raising them at all if the starting vector is [0.1,0.0,...] and the growing of some peaks that can be related to the structure of the network. One of such peaks can be due to finite size, of course we can not try with an infinite number of layers and I can not imagine how to compensate for the long tail.

But it does not converge to a delta, or a set of deltas. And for few epochs it is not granted that the best peak is the best layer. You have three nearby definitions that do not seem to coincide:

1. The layer with greatest weight in the Dense fusion of layers.
2. The layer with best result when evaluated with a Kronecker weight (w[l]=1, w[!l]=0) in the LayerSelect class.
3. The number of layers that give the best training in the usual, unfusioned, training.

Answer 3

The work had been done before, take a look at this paper. The author not only search for the number of layers but also the whole model architecture.

By using reinforcement learning, the system makes a loop, generates a model by a sequence of LSTM then validates the reward by the accuracy of the test set. It's a famous paper but not widely used because of the high complexity and huge computation cost.