Chemberta model produces random output for embeddings, it depends on seed

[GenerateIdea]

Dear Alexander Kroll

In the file https://github.com/bayer-int/pbt-dsa-gds-cmd-ProSmith/tree/main/code/preprocessing/smiles_embeddings.py there are the logits produced. However, if I load the Chemberta model twice, it produces different outputs. If I fix a seed it produces always the same output. However, the logits produced should be constant.

The layer head.dense.weight is not initialized with user-supplied weights and produce some random different embedding everytime we load it.

In the respective model file from "DeepChem/ChemBERTa-77M-MTR": "pytorch_model.bin" you can see that one of the layers has the correct size and naming "regression.dense.weight" with dimension 384x384. It seems not correctly set in the instantiated model. However, it can be easily assigned by using:

smiles_bert.lm_head.dense.weight= torch.nn.Parameter( model["regression.dense.weight"])

Please let us know if this is correct.

Currently I am running transformers=4.27.2 and had the same issues with newer versions.

[GenerateIdea]

Some code to reproduce it:
All model layers are saved in a json and later retrieved, and their difference is computed.

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
import json
import os
import numpy as np
import torch


def calculate_smiles_embeddings(all_smiles,  seed, no_of_embeddings=1000):

    n = len(all_smiles)
    parts = int(np.ceil(n / no_of_embeddings))

    for part in range(parts):
        smiles_reprs = {}
        smiles_list = all_smiles[part * no_of_embeddings: (part + 1) * no_of_embeddings]

        for k, smiles in enumerate(smiles_list):
            smiles_rep = get_last_layer_repr(smiles, seed )
            # smiles_rep.requires_grad = False
            print(smiles_rep)

def get_last_layer_repr(smiles, seed):
    torch.manual_seed(seed)
    SMILES_BERT = "DeepChem/ChemBERTa-77M-MTR"
    smiles_reprs = {}
    smiles_tokenizer = AutoTokenizer.from_pretrained(SMILES_BERT)
    smiles_bert = AutoModelForMaskedLM.from_pretrained(SMILES_BERT)



    tokenizer = smiles_tokenizer
    model = smiles_bert
    key = "logits"
    max_length = 256

    tokens = tokenizer(
        smiles,
        max_length=500,
        padding=True,
        truncation=True,
        return_tensors="pt")

    tokens["input_ids"] = tokens["input_ids"]
    tokens["attention_mask"] = tokens["attention_mask"]
    model.eval()
    with torch.no_grad():
        last_layer_repr = model(**tokens)[key]


    with open("sample{}.json".format(seed), "w") as outfile:
        a={k: v.detach().numpy().tolist() for k, v in model.state_dict().items()}
        json.dump(a, outfile)

    return last_layer_repr


def compare_weights(file1, file2):
    with open(file1, "r") as outfile:
       f1=json.load(outfile)
    with open(file2, "r") as outfile:
       f2=json.load(outfile)

    for key,value in f1.items():

        dif=np.sum(np.sum(np.abs(np.array(value)-np.array(f2[key]))))
        print("{} ({}):{}".format(key,np.array(value).shape,dif))


if __name__=="__main__":
    calculate_smiles_embeddings(["CCO"], seed=42, no_of_embeddings=1000)
    calculate_smiles_embeddings(["CCO"], seed=20 ,no_of_embeddings=1000)

    compare_weights("sample42.json", "sample20.json")

[AlexanderKroll]

Dear GenerateIdea,

Thank you for pointing this out. We have also recently realised that when ChemBERTa is loaded twice, the parameters of the last model layer change and thus the embeddings are not consistent between different initialisations of the model. The ChemBERTa model therefore needs to be fixed to a specific model that ensures consistent parameters between different model runs, which is crucial for the model to work. We have done this in our new GitHub repository "https://github.com/AlexanderKroll/ESP_ProSmith", which allows the use of our ESP ProSmith function. We have uploaded the specific ChemBERTa model and here is how to access it: https://github.com/AlexanderKroll/ESP_ProSmith/blob/main/code/smiles_embeddings.py

Of course, this raises the question of whether the use of these "random" weights in the last ChemBERTa layer won't have a significant impact on the model's performance. To check this, we reproduced the whole pipeline using MolFormer instead of ChemBERTa. Surprisingly, the model performance didn't change. This means that the ProSmith multimodal transformer network can learn all the necessary information from these partially random embeddings. I'm sorry for this problem and confusion, but it seems that as long as you fix the model version (or switch to another model such as MolFormer), it doesn't seem to affect model performance.

[GenerateIdea]

Thanks so much Alexander for the fast response, I will test this one as well !

Very excited

[GenerateIdea]

Dear Alexander,
Still seeing the same error, I installed as proposed in ESP_Prosmith the environment, changed the code slightly and had to download the model, based on the key you provided: https://huggingface.co/DeepChem/ChemBERTa-77M-MTR/tree/66b895cab8adebea0cb59a8effa66b2020f204ca

Maybe you can also take a look at my solution (it may be simpler), just setting back the respective un-initialized layer:

smiles_bert.lm_head.dense.weight= torch.nn.Parameter( model["regression.dense.weight"])

[GenerateIdea]

Sorry, didnt want to close it

[code4luck]

Sorry, didnt want to close it

Hello, when I use ESP_prosmith to evaluate the esp test dataset, but I find the result is so poor. Have you successfully tested the results of esp_prosmith?

[AlexanderKroll]

Dear @code4luck,

We have tested the code and the results should be consistent with what we published in the manuscript (with very small deviations). However, if you find that this is not the case, there is likely an issue in your pipeline. For example, how is the performance after fine-tuning the multimodal Transformer Network before using the xgboost models, and how is the performance of the different xgboost models? This may indicate where something is going wrong.

[code4luck]

Dear @code4luck,

We have tested the code and the results should be consistent with what we published in the manuscript (with very small deviations). However, if you find that this is not the case, there is likely an issue in your pipeline. For example, how is the performance after fine-tuning the multimodal Transformer Network before using the xgboost models, and how is the performance of the different xgboost models? This may indicate where something is going wrong.

Thank you for your reply. I actually used ESP_prosmith to predict the results on the ESP test set. I obtained the results from the 'Prediction score' column and calculated the metrics against the labels. I was using the specified environment configuration, the cached model weights, and the prediction script, but I found that the performance was not good. I don't know why.