Question about the latent layer used in pretrain model

[lishensuo]

Hi, I load the pretrain model that you have provided and find the lantent layer is VAELatentLayer. According the codes in CellPLM/latent/autoencoders.py, it seems the ordinary VAE encoder, not the Gaussian mixture prior distribution that mentioned in your paper. Could you please explain the question? Thanks a lot.

PRETRAIN_VERSION = '20231027_85M'
DEVICE = 'cuda:0'

from CellPLM.pipeline import load_pretrain

load_pretrain(pretrain_prefix=PRETRAIN_VERSION, # Specify the pretrain checkpoint to load
              overwrite_config={},
              pretrain_directory='../ckpt')

OmicsFormer(
  (embedder): OmicsEmbeddingLayer(
    (act): ReLU()
    (norm0): GroupNorm(4, 1024, eps=1e-05, affine=True)
    (dropout): Dropout(p=0.2, inplace=False)
    (extra_linear): Sequential(
      (0): Linear(in_features=1024, out_features=1024, bias=True)
      (1): ReLU()
      (2): Dropout(p=0.2, inplace=False)
      (3): GroupNorm(4, 1024, eps=1e-05, affine=True)
    )
    (pe_enc): Sinusoidal2dPE(
      (pe_enc): Embedding(10000, 1024)
    )
    (feat_enc): OmicsEmbedder()
  )
  (mask_model): MaskBuilder()
  (encoder): TransformerEncoder(
    (layers): ModuleList(
      (0-3): 4 x FlowformerLayer(
        (self_attn): Flow_Attention(
          (query_projection): Linear(in_features=1024, out_features=1024, bias=True)
          (key_projection): Linear(in_features=1024, out_features=1024, bias=True)
          (value_projection): Linear(in_features=1024, out_features=1024, bias=True)
          (out_projection): Linear(in_features=1024, out_features=1024, bias=True)
          (dropout): Dropout(p=0.01, inplace=False)
        )
        (_ff_block): Sequential(
          (0): Linear(in_features=1024, out_features=2048, bias=True)
          (1): GELU(approximate='none')
          (2): Dropout(p=0.2, inplace=False)
          (3): Linear(in_features=2048, out_features=1024, bias=True)
          (4): Dropout(p=0.2, inplace=False)
        )
        (dropout1): Dropout(p=0.2, inplace=False)
        (norm1): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
        (norm2): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
      )
    )
  )
  (latent): LatentModel(
    (layers): ModuleList(
      (0): PlaceholderLayer()
      (1): VAELatentLayer(
        (hid_2mu): Linear(in_features=1024, out_features=512, bias=True)
        (hid_2sigma): Linear(in_features=1024, out_features=512, bias=True)
      )
    )
  )
  (decoder): NBMLPDecoder(
    (layers): ModuleList(
      (0): Sequential(
        (0): Linear(in_features=512, out_features=1024, bias=True)
        (1): PReLU(num_parameters=1)
        (2): Dropout(p=0.2, inplace=False)
        (3): GroupNorm(4, 1024, eps=1e-05, affine=True)
      )
    )
    (covariate_layers): ModuleList(
      (0): ModuleDict(
        (batch): Sequential(
          (0): Embedding(12596, 1024)
          (1): PReLU(num_parameters=1)
          (2): GroupNorm(4, 1024, eps=1e-05, affine=True)
        )
        (dataset): Sequential(
          (0): Embedding(484, 1024)
          (1): PReLU(num_parameters=1)
          (2): GroupNorm(4, 1024, eps=1e-05, affine=True)
        )
      )
    )
    (out_layer): NB(
      (dec_1): Linear(in_features=1024, out_features=128, bias=True)
      (dec_mean): Sequential(
        (0): Linear(in_features=128, out_features=19374, bias=True)
        (1): MeanAct()
      )
      (dec_disp_act): DispAct()
    )
  )
  (objective): Objectives(
    (layers): ModuleList(
      (0): NBReconstructionLoss()
    )
  )
  (pre_latent_norm): PreLatentNorm(
    (norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)
  )
)