Update README.md

README.md

@@ -6,7 +6,58 @@ base_model:
 - google-bert/bert-base-uncased
 ---
 
-How to use the model step-by-step:
+# 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**
+
 
 ---