Linear probe results replication doesn't match with the paper

[Vsarangi]

Hi, I am trying to replicate the Merfish mouse brain linear probe classification part of the paper (Fig 4b). I used the "downstream_fine_tune.ipynb" and modified it based on method section of the paper.
However, the results I am getting are much lower than reported in the paper:

I suspect it's the hyper parameters. Could you please take a look at the config and let me know. I couldn't find any other details from the paper.

Thank you!

Here is my config:

config = {

'data_path': arg.df5_file,  # Path to your AnnData file
'technology_mean_path': arg.technology_mean,  # Path to technology mean file
'checkpoint_path': arg.checkpoint,  # Path to pre-trained model
'output_path': output_file,  # Where to save results
'output_dir': output_dir,  # Directory for checkpoints

# Training parameters
'batch_size': 256,
'max_seq_len': 1500,
'aux_tokens': 30,
'chunk_size': 1000,
'num_workers': 4,
'precision': 32,
'max_epochs': 1,
'lr': 1e-3,
'warmup': 1,
'gradient_clip_val': 1.0,
'accumulate_grad_batches': 1,

# Model parameters
'supervised_task': 'niche_classification',  # or whichever task
'extract_layers': [11],  # Which layers to extract features from
'function_layers': "mean",  # Architecture of prediction head
'dim_prediction': 1, # dim of the output vector
'n_classes': 33,  # only foor classification tasks
'freeze': True,  # Whether to freeze backbone
'reinit_layers': None,
'extractor': False,
'regress_distribution': False,
'pool': 'mean',
'predict_density': False,
'ignore_zeros': False,
'organ': 'brain',
'label': 'author_cell_type'  # The target variable to predict
}