Approximation of long sequence of layers by one layer

Consider the following situation : there is a deep neural network with a lot of layers, and in order to speed up the inference or for regularization purposes one would like to reduce the complexity of the model.

Layers from i up to j, effectively perform some nonlinear transformation: $$\mathbb{R}^{d_i} \rightarrow \mathbb{R}^{d_j}$$ Where $$d_i$$ and $$d_j$$ is the dimensionality of the input at $$i_{th}$$ and $$j_{th}$$ layers. Has there been any reseach on approximation of this nonlinear transform by some linear one - search of convolution layers with the activation function, which approximate the action of multiple layers as close as possible.

The loss function for this approximation can be simply an MSE error, where the inputs are the activations of $$\mathbb{R}^{d_i}$$ and outputs - activations $$\mathbb{R}^{d_j}$$. I see the possible problem, that mapping from one high dimensional space to another requires a lot od training data for reliable estimate, but, nevertheless I wonder, whether such an approach has been pursued in the literature?

The teacher/student model approach (at least as I understand it) could be used. It is normally used to replace all layers, but there is nothing stopping you applying it to a subset of layers.

First you create training data, by running your training data through the (trained) network, and recording the inputs to layer i, as your new training data (x) and the outputs from layer j as your desired answer (y).

You now simply (!) want to design a model that takes x and produces y. Ideally using fewer weights than layers i to j currently use.

If those layers were a fully-connected NN, then this is the idea behind autoencoders. If a Transformer then try with one attention + FCNN layer, and see how low the error will go. If it flatlines at an unacceptable error rate, maybe you will need to try two layers. If a CNN, then I am less sure it will work; my hunch would be would have better success keeping the same number of layers and convolution sizes, and just using fewer features in each layer.

Anyway, when you have trained this student model, you can plug it in to your original model to replace layers i to j. (Then maybe run some fine-tuning training.)

Note: if your layers have residual connections, you will have to think carefully about how they handled. Taking the i-1th output and feeding it to the j+1th layer, when it previously expected the output of the jth layer, may go badly. Especially if the dimension of the data is different.