Multitasking bert for multilabel classification of 5 categories

Question

I built and finetuned 5 BioClinicalBERT-based models (finetuned bert) to predict labels for medical records for the following categories:

specialties = ["aud","den","oph","oto","psy","tbi"]

clinical_summaries = ["consultation", "hospital_discharge", "operative"]

diagnostics = ["imaging", "lab", "vision", "audio", "psychological", "misc"]

organs = ["abdominal_inguinal_femoral_hernias", "ankle", "artery_and_vein", "back", "bone", "breast", "central_nervous_system", "cranial_nerve", "elbow_and_forearm", "endocrine", "esophageal", "foot", "hand_and_finger", "heart", "hip_and_thigh", "kidney", "knee_and_lower_leg", "male_repro", "muscles", "neck", "nose_and_sinus",  "osteomyelitis", "peripheral_nerves", "rectum_and_anus", "shoulder", "skin", "stomach_and_duodenal", "thyroid_and_parathyroid", "urinary_tract", "wrist"]

diseases = ["amputations", "amyotrophic_lateral_sclerosis", "chronic_fatigue","cold_injury", "diabetes_mellitus", "diabetic_peripheral_neuropathy", "fibromyalgia", "gynecological",  "hairy_cell_leukemia", "headaches", "hemic_and_lymphatic_leukemia", "hepatitis_cirrhosis_and_liver", "hiv", "hypertension", "infectious_diseases", "intestines_nonsurgical", "intestines_surgical", "loss_smell_taste", "multiple_sclerosis", "narcolepsy", "nondegenerative_arthritis", "nutritional_deficiencies", "parkinsons", "peritoneal_adhesions", "persian_gulf_afghanistan_infections", "prostate_cancer", "respiratory", "scars", "seizure_disorders_epilepsy", "sleep_apnea", "systemic_lupus_erythematosus", "tuberculosis"]. 

The fine-tuned models were built with the following architecture:

# Define the model class
class BioClinicalBERTClass(torch.nn.Module):
    def __init__(self, num_labels):
        super(BioClinicalBERTClass, self).__init__()
        self.bert_model = BertForSequenceClassification.from_pretrained("emilyalsentzer/Bio_ClinicalBERT", num_labels=num_labels)

    def forward(self, input_ids, attention_mask, token_type_ids):
        output = self.bert_model.bert(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        pooled_output = output.pooler_output
        logits = self.bert_model.classifier(pooled_output)
        return logits 

Now am working on building a multitasking model to combine the 5 fine-tuned models, am not sure if it's more suitable to use the same architecture with a single output layer for all categories, or an architecture with task-specific output layers as follows:

class MultiTaskModel(torch.nn.Module):
    def __init__(self, num_specialties, num_clinical_summaries, num_diagnostics, num_organs, num_diseases):
        super(MultiTaskModel, self).__init__()
        self.shared_base = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")

        # Task-specific output layers
        self.specialties_output = torch.nn.Linear(self.shared_base.config.hidden_size, num_specialties)
        self.clinical_summaries_output = torch.nn.Linear(self.shared_base.config.hidden_size, num_clinical_summaries)
        self.diagnostics_output = torch.nn.Linear(self.shared_base.config.hidden_size, num_diagnostics)
        self.organs_output = torch.nn.Linear(self.shared_base.config.hidden_size, num_organs)
        self.diseases_output = torch.nn.Linear(self.shared_base.config.hidden_size, num_diseases)

    def forward(self, input_ids, attention_mask, token_type_ids):
        shared_output = self.shared_base(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        pooled_output = shared_output.pooler_output

        specialties_logits = self.specialties_output(pooled_output)
        clinical_summaries_logits = self.clinical_summaries_output(pooled_output)
        diagnostics_logits = self.diagnostics_output(pooled_output)
        organs_logits = self.organs_output(pooled_output)
        diseases_logits = self.diseases_output(pooled_output)

        return specialties_logits, clinical_summaries_logits, diagnostics_logits, organs_logits, diseases_logits

I need recommendations for the convenient architecture to perform the multilabel classification for all categories with a single model laveraging the finetuning knowledge of the 5 trained models. Please provide a detailed explanation about their weights and biases loading and use to require minimal fine-tuning making it understand the correlation between categories. Should models contribute equally to the new model or should their influence on parameters be based on the number of predicted category labels?