(How can this thing be back-propagable ? Out of curiosity, I would be interested to see how you compute the loss, if you don’t mind sharing the code.)
Another curious question : is the simplicity of the primitives (simple shapes) the reason why you chose such a small hidden size of 16 ?
Maybe a good idea would be to use some Conv2d layers to extract the image features, flatten the features, and then use some Conv3d layers to create the voxel image.
Your post reminds me of this paper : Unsupervised Learning of 3D Representations From Natural Images 26.
I don’t know if is fits your work exactly, but maybe you will find good ideas. They use a mix of 3D and 2D convolutions.