JayLacoma/fmri_encoder_3.0model

Transformer-Based fMRI Encoder Model

This repository contains a Transformer-based model trained on neuroimaging datasets to classify conditions like Autism Spectrum Disorder (ASD) and ADHD, and to analyze brain activity during movie-watching. The model combines fMRI data with demographic features (age and gender) for binary classification tasks. Below is a detailed explanation of the datasets, model architecture, and training process.

Model Architecture

The model integrates multi-modal data and leverages a Transformer backbone for feature extraction. Below is a breakdown of its components:

1. Inputs

• fMRI ROI Data: High-dimensional features representing brain activity.
• Age Data: Numerical input passed through a Multi-Layer Perceptron (MLP).
• Gender Data: Binary input (male/female) embedded into a dense representation.

2. Transformer Backbone

• A pretrained Hugging Face Transformer (e.g., BERT) with:
  • Configurable number of attention heads, layers, and hidden size.
  • Dropout for regularization.
• Dynamically adjusted hyperparameters using AutoConfig.

3. Pooling Mechanisms

• Aggregates the Transformer’s sequence outputs into a single vector using:
  • Mean Pooling: Averages hidden states.
  • Max Pooling: Selects the maximum value for each feature.
  • Attention Pooling: Learns attention weights to emphasize important sequence elements.

4. Output

• A fully connected layer maps the pooled output to a scalar value for binary classification.

---

Training Process

Key Details:

• Loss Function: Binary Cross Entropy with Logits (BCEWithLogitsLoss), with class imbalance handled using positive weights.
• Optimizer: Ranger (combines RAdam and Lookahead for stable convergence).
• Learning Rate Scheduler: Cosine Annealing for gradual learning rate reduction.
• Gradient Clipping: Prevents exploding gradients with a clipping threshold of 1.0.
• Early Stopping: Stops training after 250 epochs without validation loss improvement.

Datasets Used:

1. ABIDE: Autism vs. control classification.
2. ADHD-200: ADHD vs. control classification.
3. Pixar Movie Dataset (Nilearn): Brain activity analysis during movie-watching.

Output:

The model’s state dictionary is saved as fmri_encoder_model.pth.

---

How to Use This Model

1. Import Required Libraries

At the beginning, you import all necessary libraries for the model's implementation and training.

from transformers import AutoModel, AutoConfig
import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingLR

---

2. Define the Transformer Model

The TransformerModel is implemented as a PyTorch nn.Module. It integrates a Hugging Face transformer (e.g., bert-base-uncased) and custom embedding layers for ROI, age, and gender features.

• Embedding Layers:

  • age_mlp processes age as a continuous input.
  • gender_embed embeds the binary gender feature.

• ROI Encoder: Encodes the ROI feature input using a small feedforward network.

• Transformer: Initializes a Hugging Face transformer with a custom configuration.

• Pooling Mechanism: Supports three pooling strategies:

  • Mean pooling
  • Max pooling
  • Attention pooling (using an additional attention mechanism).

class TransformerModel(nn.Module):
    def __init__(self, roi_input_dim, embed_dim, num_heads, num_layers, dropout_rate, pretrained_model_name="bert-base-uncased", pooling="mean"):
        super(TransformerModel, self).__init__()

        # Ensure embed_dim is divisible by num_heads
        if embed_dim % num_heads != 0:
            embed_dim = (embed_dim // num_heads) * num_heads
            print(f"Adjusted embed_dim to {embed_dim} to ensure divisibility by num_heads.")

        # Embedding layers for age and gender
        self.age_mlp = nn.Sequential(
            nn.Linear(1, embed_dim),
            nn.GELU(),
            nn.BatchNorm1d(embed_dim),
            nn.Dropout(dropout_rate),
            nn.Linear(embed_dim, embed_dim),
        )
        self.gender_embed = nn.Embedding(2, embed_dim)

        # ROI encoder
        self.roi_encoder = nn.Sequential(
            nn.Linear(roi_input_dim, embed_dim),
            nn.GELU(),
            nn.BatchNorm1d(embed_dim),
            nn.Linear(embed_dim, embed_dim),
            nn.Dropout(dropout_rate),
        )

        # Hugging Face Transformer Model
        config = AutoConfig.from_pretrained(pretrained_model_name)
        config.hidden_size = embed_dim
        config.num_attention_heads = num_heads
        config.num_hidden_layers = num_layers
        config.hidden_dropout_prob = dropout_rate
        config.attention_probs_dropout_prob = dropout_rate
        self.transformer = AutoModel.from_config(config)

        # Pooling mechanism
        assert pooling in ["mean", "max", "attention"]
        self.pooling = pooling

        if pooling == "attention":
            self.attention_pool = nn.Sequential(
                nn.Linear(embed_dim, embed_dim // 2),
                nn.Mish(),
                nn.Linear(embed_dim // 2, 1),
                nn.Softmax(dim=1),
            )

        # Output layer
        self.output_layer = nn.Linear(embed_dim, 1)

    def forward(self, roi_data, age, gender):
        age_embed = self.age_mlp(age.unsqueeze(-1))
        gender_embed = self.gender_embed(gender.long())
        roi_encoded = self.roi_encoder(roi_data)

        combined_input = torch.stack((roi_encoded, age_embed, gender_embed), dim=1)
        transformer_output = self.transformer(inputs_embeds=combined_input).last_hidden_state

        if self.pooling == "mean":
            pooled_output = transformer_output.mean(dim=1)
        elif self.pooling == "max":
            pooled_output = transformer_output.max(dim=1).values
        elif self.pooling == "attention":
            attention_weights = self.attention_pool(transformer_output)
            pooled_output = (attention_weights * transformer_output).sum(dim=1)

        return self.output_layer(pooled_output).squeeze(1)

---

3. Initialize the Model

The model is initialized using the best hyperparameters obtained (e.g., from an Optuna study or prior experimentation).

• Parameters:
  • roi_input_dim: Number of input features for the ROI data.
  • embed_dim, num_heads, num_layers: Transformer model hyperparameters.
  • dropout_rate: Dropout rate to prevent overfitting.
  • pretrained_model_name: Name of the Hugging Face pretrained model.
  • pooling: Pooling strategy to summarize the transformer outputs.

model = TransformerModel(
    roi_input_dim=roi_features.shape[1],
    embed_dim=best_params["embed_dim"],
    num_heads=best_params["num_heads"],
    num_layers=best_params["num_layers"],  # Optuna-tuned
    dropout_rate=best_params["dropout_rate"],
    pretrained_model_name="bert-base-uncased",  # Pretrained transformer model
    pooling="attention",  # Optuna-tuned pooling strategy
).to(device)

---

4. Support for Multi-GPU Training

If multiple GPUs are available, the model is wrapped in nn.DataParallel for parallel training.

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

---

5. Define Loss Function and Optimizer

• Loss Function: The BCEWithLogitsLoss is used for binary classification tasks, and class imbalance is handled by a computed pos_weight.

• Optimizer: Uses the Ranger optimizer with a learning rate from the best hyperparameters.

• Scheduler: A cosine annealing learning rate scheduler adjusts the learning rate over training.

# Compute class weights for imbalance
pos_weight = torch.tensor([len(new_y) / new_y.sum() - 1], dtype=torch.float32).to(device)
criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight)

# Optimizer and Scheduler
optimizer = optimizers.Ranger(model.parameters(), lr=best_params["lr"], weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=10, eta_min=0)

---

6. Load Pretrained Model Weights

The model can load previously trained weights to resume training or perform inference.

model.load_state_dict(torch.load("/kaggle/input/bert-encoder-fmri/pytorch/default/1/fmri_encoder_model.pth"))

---

7. Best Hyperparameters

These are the best hyperparameters used for model initialization and training.

best_params = {
    "embed_dim": 768,
    "num_heads": 32,
    "num_layers": 12,
    "dropout_rate": 0.119,
    "lr": 3.66e-5,
}

---