piyush-555 / fmfastsim Goto Github PK

Change flattening of the 3D patches to Conv3D at the start (before the transformer blocks) and add Conv3DTranspose at the end (after the last blocks in the ‘decoder’). Deconv layers are known to be able to produce crisp images, and hopefully this will fix the blurring effect observed with the plain VAE.

Change the dense layers of the decoder in the VAE to transformer blocks. Expand dimensionality from the bottleneck up the same way the dimensionality reduction was done in the encoder.

Check data-preprocessing steps. Energy typically should be normalised either with log or possibly power-law. Start with log, and then explore power-law while looking into different power-law parameters as done here. Otherwise we can look into symbolic encoding, however, this might not be accurate enough given that the precision will be limited by the number of categories (therefore keep this option for later and proceed with the next steps).

https://link.springer.com/article/10.1140/epjc/s10052-022-10258-4

Change positional encoding from flat embedding to i,j,k coordinates embedding. E.g. not flatting first and embedding after, but going from 3 coordiates i,j,k to embedding.

Think about phi coordinate and how to deal with phi roll-over. E.g sin / cos wave encoding could be an option.

Weighted random sampling to sample more from denser parts.