The concatenated output from 2D-DWT is passed to an MLP layer (two \(1\times 1\)  convolutional layers separated by a GELU non-linearity) having a multiplication factor more than one where channel mixing is performed by the MLP. Since wavelet transform reduces the image resolution by half, the GPU consumption and computations needed by the MLP significantly reduces in each layer. The image size reconciliation is performed using transposed convolutions (up-convolutions) which resizes the image back to the original input resolution. The kernel size and stride of deconvolutional layers were chosen such that the output has the same size as the input to WaveMix-Lite block. We chose deconvolutional layer rather than an inverse 2D-DWT because the former is much faster and consumes less GPU than the latter. The outputs from the deconvolutional layers are then passed through batch normalization. A residual connection [1]} is provided within each WaveMix-Lite block so that the model can be made deeper with a larger number of blocks, if necessary.
