Skip to content

Ng2k/depression-transformer-lab

BranchesTags

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

7 Commits
configs
configs
 
 
src/depformer
src/depformer
 
 
tests/unit/utils
tests/unit/utils
 
 
.gitignore
.gitignore
 
 
LICENSE
LICENSE
 
 
README.md
README.md
 
 
pyproject.toml
pyproject.toml
 
 
requirements.txt
requirements.txt
 
 

Repository files navigation

MIL-Project — Depression Detection via Multiple Instance Learning

A research framework for automatic depression detection from clinical interview transcripts using Multiple Instance Learning (MIL) with transformer-based encoders. Built for the AI in Bioinformatics course at Unimore.

Table of Contents

Overview

The project frames depression detection as a Multiple Instance Learning problem. Each clinical interview is a bag, and each Q&A turn (or answer chunk) is an instance. The model learns to predict a bag-level label (depressed / not depressed) from a set of unlabelled instances, without requiring per-turn annotations.

Datasets supported:

ID Source Variant Splits
daic_woz_qa DAIC-WOZ Questions + Answers train/val/test
daic_woz_answers DAIC-WOZ Answers only train/val/test
edaic E-DAIC Answers only train/val/test

Label: PHQ-8 score > 10 → depressed (binary).

Backbones:

  • roberta — cardiffnlp/twitter-roberta-base-sentiment (768-dim)
  • mt5 — google/mt5-small (512-dim)

Pooling strategies:

  • gated — Gated Attention (ABMIL, Ilse et al. ICML 2018)
  • mha — Multi-Head Attention with global learned query
  • transmil — TransMIL with sinusoidal positional encoding + inter-chunk self-attention (NeurIPS 2021)

Classifiers: linear, mlp, moe (Mixture of Experts)

Architecture at a Glance

Interview Transcript
        │
        ▼
  [Preprocessing]          Parse TSV/CSV → Q&A turns → augment critical answers
        │
        ▼
  Preprocessed JSON         {participant_id, label, phq8_score, instances: [...]}
        │
        ▼
  [Tokenization]            Chunk instances → RoBERTa / mT5 tokenizer
        │
        ▼
  Tokenized Bags (.pt)      {input_ids: [N_chunks, 512], attention_mask: [...]}
        │
        ▼
  [Random Search]           K-fold CV on val set → best hyperparameters
        │
        ▼
  [Training]                K-fold on train+val → best model checkpoint
        │
        ▼
  [Test / Inference]        MC Dropout → AUROC, AUPRC, F1, confusion matrix

MILModel forward pass:
                                      ┌──────────────────────┐
 input_ids [N, 512]    ──► Backbone ──► chunk_embeddings      │
 attention_mask [N, 512]             │  [N, d_model]          │
                                     │         │               │
                                     │   ┌─────┴──────┐       │
                                     │   │  Pooling   │       │
                                     │   │ (gated /   │       │
                                     │   │  mha /     │       │
                                     │   │ transmil)  │       │
                                     │   └─────┬──────┘       │
                                     │         │               │
                                     │  bag_repr [d_model]     │
                                     │         │               │
                                     │   ┌─────┴──────┐       │
                                     │   │ Classifier │       │
                                     │   │ (linear /  │       │
                                     │   │  mlp / moe)│       │
                                     │   └─────┬──────┘       │
                                     │         │               │
                                     │  bag_logits [1, 1]      │
                                     │                         │
                                     │  instance_logits [N, 1] │ (classifier applied per-chunk too)
                                     └──────────────────────────┘

Prerequisites

  • Python 3.12+
  • CUDA 12.4+ (recommended; CPU fallback works but training is very slow)
  • pip or uv
  • Access to DAIC-WOZ and/or E-DAIC datasets (requires institutional registration)

Installation

# Clone the repository
git clone <repo-url>
cd depression-transformer-lab

# Create and activate a virtual environment
python -m venv .venv
source .venv/bin/activate        # Linux/macOS
.venv\Scripts\activate           # Windows

# Install the project and all dependencies
pip install -r requirements.txt
pip install -e .

This installs the package in editable mode and registers the five CLI entry points: preprocess, tokenize, random_search, train, test.

Environment Setup

Copy dev.env and fill in your dataset paths:

cp dev.env .env

dev.env variables:

Variable Default Description
PYTHON_ENV development Set to prod for JSON structured logs
LOG_LEVEL DEBUG Python log level (DEBUG, INFO, WARNING)
DATASET_ROOT ./datasets/ Root folder containing raw datasets
DAIC_WOZ_URL USC DCAPS portal Download URL for DAIC-WOZ
EDAIC_ROOT ./datasets/edaic/ Root folder for E-DAIC data
EDAIC_URL USC DCAPS portal Download URL for E-DAIC

The scripts load either dev.env or prod.env depending on PYTHON_ENV.

Dataset Acquisition

Both datasets require institutional access through the USC DCAPS portal.

DAIC-WOZ

  1. Register at the DCAPS portal
  2. Download and extract transcripts into ./datasets/daic_woz/transcripts/
  3. Place split CSV files into ./datasets/daic_woz/splits/

Expected split CSV columns: Participant_ID, PHQ8_Score, PHQ8_Binary

E-DAIC

  1. Register at the DCAPS portal
  2. Download and extract transcripts into ./datasets/edaic/transcripts/
  3. Place split CSV files into ./datasets/edaic/splits/

End-to-End Pipeline

All commands below assume you are in the project root with the virtual environment active.

Step 1 — Preprocess

Parse raw transcripts into MIL bags (JSON):

preprocess \
  --datasets daic_woz_qa daic_woz_answers edaic \
  --out-dir ./datasets

Output: ./datasets/<dataset_name>/<dataset_label>/{train,val,test}.json

Step 2 — Tokenize

Tokenize and chunk bags into .pt embedding files:

tokenize \
  --datasets daic_woz_qa daic_woz_answers edaic \
  --encoders roberta mt5 \
  --splits train val test \
  --chunk-size 5 \
  --out-dir ./outputs/embeddings

Output: ./outputs/embeddings/<model>/<dataset_id>/{train,val,test}.pt

--chunk-size 5 means 5 Q&A turns are grouped into one chunk (one bag instance).

Step 3 — Random Search (Hyperparameter Optimization)

Find optimal hyperparameters via stratified k-fold cross-validation on the validation split:

random_search \
  --datasets daic_woz_qa \
  --models roberta \
  --pooling transmil \
  --classifier linear \
  --pt-root ./outputs/embeddings \
  --out-dir ./outputs/models \
  --n-trials 30 \
  --epochs 30 \
  --k-folds 5

Output: ./outputs/models/<model>/<dataset>/<classifier>_<pooling>/best_params.json

Note on this experiment's hyperparameters: all 36 combinations in the published results share the same fixed best_params.json, rather than one tuned per combination. This was a deliberate choice to keep the comparison fair and reduce total compute. The random_search script is provided so that future experiments can tune hyperparameters independently per combination for a deeper, per-configuration optimisation.

Step 4 — Train

Train using the best hyperparameters found in step 3:

train \
  --datasets daic_woz_qa \
  --models roberta \
  --pooling transmil \
  --classifier linear \
  --pt-root ./outputs/embeddings \
  --out-dir ./outputs/models \
  --epochs 30 \
  --k-folds 5 \
  --use-amp

Output per fold:

  • best_model_fold_<k>.pt — fold checkpoints
  • best_model_overall.pt — best fold copied here
  • fold_results.json — per-fold metrics + 95% confidence intervals

Step 5 — Test

Evaluate the best model on the held-out test set:

test \
  --datasets daic_woz_qa \
  --models roberta \
  --pooling transmil \
  --classifier linear \
  --pt-root ./outputs/embeddings \
  --out-dir ./outputs/results \
  --n-passes 10

Output: ./outputs/results/<model>/<dataset>/<classifier>_<pooling>/metrics.json

Printed to stdout: accuracy, F1-macro, recall, AUROC, AUPRC, confusion matrix, classification report.

Reproducing Results

To exactly reproduce any published result:

  1. Use the same --pooling, --classifier, --models, --datasets combination.
  2. The global seed is hardcoded as DEFAULT_SEED = 42 and applied to Python, NumPy, and PyTorch RNGs before every script.
  3. Use --use-amp only if you have a CUDA GPU; CPU results will match without it.
  4. All 36 combinations in this experiment were trained with a single shared best_params.json (not tuned per combination). To reproduce exactly, use the best_params.json already present in each outputs/models/<model>/<dataset>/<clf>_<pool>/ directory — do not re-run random_search, as that would overwrite the shared parameters with combination-specific ones.
  5. To run a deeper experiment with per-combination hyperparameter tuning, run random_search for each combination first and then re-train. Results will differ from those reported here.
  6. The fold_results.json written by train contains the exact per-fold metrics and 95% CI used in the paper.

Model Performance

Models were evaluated on two held-out test sets: the DAIC-WOZ test set (47 participants: 14 depressed, 33 non-depressed) and the E-DAIC test set (55 participants: 17 depressed, 38 non-depressed), using Monte Carlo Dropout with 10 forward passes. Cross-validation (CV) metrics were computed over 5 stratified folds on the training + validation split; 95% confidence intervals use the t-distribution with n − 1 degrees of freedom.

Configuration space explored (36 total):

Axis Values
Backbone RoBERTa (cardiffnlp/twitter-roberta-base-sentiment, 768-dim), mT5 (google/mt5-small, 512-dim)
Pooling gated (ABMIL), mha (Multi-Head Attention), transmil (TransMIL)
Classifier linear, mlp, moe (Mixture of Experts)
Dataset daic_woz_qa (Q+A turns), daic_woz_answers (answers only), edaic

Test Set Results — RoBERTa on DAIC-WOZ QA

Classifier Pooling Accuracy F1 Macro Recall (Dep.) AUROC AUPRC TP FP TN FN
linear gated 0.8298 0.7873 0.6429 0.8961 0.8101 9 3 30 5
linear mha 0.7660 0.7353 0.7143 0.8766 0.7189 10 7 26 4
linear transmil 0.7021 0.6908 0.8571 0.8615 0.7989 12 12 21 2
mlp gated 0.8085 0.7756 0.7143 0.9113 0.8116 10 5 28 4
mlp mha 0.7447 0.7389 1.0000 0.9048 0.7947 14 12 21 0
mlp transmil 0.7660 0.7496 0.8571 0.8593 0.7171 12 9 24 2
moe gated 0.7447 0.6189 0.2857 0.8874 0.7409 4 2 31 10
moe mha 0.7872 0.7749 0.9286 0.8896 0.7312 13 9 24 1
moe transmil 0.8298 0.8105 0.8571 0.8961 0.7016 12 6 27 2

RoBERTa QA highlights:

  • Best F1 and accuracy: MoE + TransMIL — highest F1 (0.811) and accuracy (0.830) with only 2 missed depression cases (FN = 2).
  • Best AUROC: MLP + Gated — AUROC 0.911, only 5 false positives and 4 false negatives.
  • Best sensitivity: MLP + MHA — recall 1.000, zero missed cases (FN = 0); every depression case detected. Cost: 12 false positives.
  • Most specific: MoE + Gated — only 2 false positives, but very low recall (0.286, FN = 10) — only appropriate when false alarms carry high cost.
  • TransMIL pooling consistently achieves recall ≥ 0.857 for linear and MoE classifiers.

Test Set Results — RoBERTa on DAIC-WOZ Answers

Classifier Pooling Accuracy F1 Macro Recall (Dep.) AUROC AUPRC TP FP TN FN
linear gated 0.7872 0.7202 0.5000 0.8030 0.6644 7 3 30 7
linear mha 0.7021 0.6083 0.3571 0.7900 0.6292 5 5 28 9
linear transmil 0.7660 0.7257 0.6429 0.8550 0.7796 9 6 27 5
mlp gated 0.7447 0.6948 0.5714 0.8463 0.7082 8 6 27 6
mlp mha 0.7660 0.7257 0.6429 0.8117 0.6395 9 6 27 5
mlp transmil 0.7447 0.7157 0.7143 0.8074 0.6899 10 8 25 4
moe gated 0.7447 0.7299 0.8571 0.8398 0.7050 12 10 23 2
moe mha 0.8085 0.7952 0.9286 0.8874 0.7700 13 8 25 1
moe transmil 0.7660 0.7549 0.9286 0.8203 0.6526 13 10 23 1

RoBERTa Answers highlights:

  • Best overall: MoE + MHA — highest AUROC (0.887), F1 (0.795), accuracy (0.808), recall 0.929 (FN = 1).
  • Best recall (tied): MoE + MHA and MoE + TransMIL both achieve recall 0.929 with only 1 missed case.
  • QA vs Answers: RoBERTa on QA is consistently stronger — best AUROC 0.911 vs 0.887, best F1 0.811 vs 0.795. The question context in QA provides additional signal that improves test-set predictions.

Test Set Results — mT5 (both dataset variants)

Classifier Pooling Accuracy F1 Macro Recall (Dep.) AUROC AUPRC TP FP TN FN
linear gated 0.6383 0.5909 0.5000 0.6147 0.3850 7 10 23 7
linear mha 0.6596 0.6083 0.5000 0.6364 0.3934 7 9 24 7
linear transmil 0.5957 0.5766 0.6429 0.6948 0.5623 9 14 19 5
mlp gated 0.6596 0.6083 0.5000 0.5281 0.4595 7 9 24 7
mlp mha 0.6170 0.5594 0.4286 0.5649 0.3789 6 10 23 8
mlp transmil 0.6170 0.3816 0.0000 0.6364 0.3747 0 4 29 14
moe gated 0.6383 0.5909 0.5000 0.6905 0.4518 7 10 23 7
moe mha 0.6383 0.5583 0.3571 0.5368 0.3647 5 8 25 9
moe transmil 0.6596 0.6210 0.5714 0.6710 0.4530 8 10 23 6

mT5 highlights:

  • mT5 is substantially weaker than RoBERTa: best AUROC 0.695 vs 0.911 (RoBERTa QA), best F1 0.621 vs 0.811.
  • QA and Answers variants produce identical test-set predictions for mT5 — the tokenised test-set bags are the same after chunking regardless of format.
  • MLP + TransMIL collapses — predicts all participants as non-depressed (recall 0.000, TP = 0). This is a degenerate solution.
  • Linear + TransMIL achieves the highest recall (0.643) and AUROC (0.695) for mT5.
  • mT5's smaller embedding dimension (512 vs 768) and domain mismatch (multilingual model on English clinical text) likely explain the performance gap.

Cross-Validation Results — RoBERTa on DAIC-WOZ QA (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.728 [0.589, 0.868] 0.773 [0.653, 0.893] 0.636 [0.384, 0.888]
linear mha 0.724 [0.580, 0.868] 0.780 [0.673, 0.887] 0.567 [0.345, 0.788]
linear transmil 0.688 [0.532, 0.844] 0.716 [0.546, 0.886] 0.719 [0.536, 0.903]
mlp gated 0.706 [0.579, 0.832] 0.745 [0.640, 0.850] 0.681 [0.391, 0.970]
mlp mha 0.699 [0.592, 0.805] 0.737 [0.656, 0.819] 0.703 [0.378, 1.028]
mlp transmil 0.523 [0.391, 0.654] 0.695 [0.656, 0.734] 0.217 [−0.053, 0.487]
moe gated 0.702 [0.619, 0.786] 0.737 [0.640, 0.835] 0.683 [0.466, 0.901]
moe mha 0.706 [0.570, 0.842] 0.794 [0.715, 0.874] 0.450 [0.186, 0.714]
moe transmil 0.686 [0.569, 0.803] 0.737 [0.620, 0.854] 0.547 [0.431, 0.664]
Per-fold F1 details — RoBERTa QA
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.758 0.905 0.620 0.649 0.708
linear mha 0.792 0.863 0.619 0.590 0.754
linear transmil 0.741 0.836 0.642 0.500 0.721
mlp gated 0.716 0.863 0.590 0.650 0.708
mlp mha 0.748 0.789 0.681 0.563 0.713
mlp transmil 0.408 0.417 0.650 0.576 0.563
moe gated 0.748 0.788 0.626 0.650 0.701
moe mha 0.741 0.863 0.576 0.635 0.714
moe transmil 0.718 0.810 0.591 0.591 0.721

Cross-Validation Results — RoBERTa on DAIC-WOZ Answers (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.659 [0.509, 0.808] 0.696 [0.527, 0.865] 0.639 [0.440, 0.837]
linear mha 0.717 [0.579, 0.856] 0.773 [0.637, 0.909] 0.547 [0.387, 0.707]
linear transmil 0.680 [0.527, 0.834] 0.767 [0.658, 0.875] 0.450 [0.169, 0.731]
mlp gated 0.710 [0.584, 0.836] 0.773 [0.660, 0.887] 0.525 [0.329, 0.721]
mlp mha 0.663 [0.506, 0.819] 0.752 [0.625, 0.879] 0.425 [0.144, 0.706]
mlp transmil 0.705 [0.523, 0.887] 0.753 [0.588, 0.917] 0.567 [0.381, 0.752]
moe gated 0.667 [0.529, 0.805] 0.745 [0.654, 0.837] 0.472 [0.198, 0.746]
moe mha 0.654 [0.532, 0.777] 0.709 [0.603, 0.815] 0.575 [0.294, 0.856]
moe transmil 0.678 [0.518, 0.839] 0.731 [0.567, 0.895] 0.586 [0.295, 0.877]
Per-fold F1 details — RoBERTa Answers
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.517 0.810 0.594 0.619 0.754
linear mha 0.752 0.810 0.581 0.619 0.825
linear transmil 0.554 0.788 0.604 0.620 0.836
mlp gated 0.623 0.810 0.611 0.682 0.825
mlp mha 0.554 0.788 0.535 0.635 0.801
mlp transmil 0.589 0.850 0.563 0.650 0.873
moe gated 0.568 0.772 0.576 0.619 0.801
moe mha 0.580 0.704 0.562 0.626 0.801
moe transmil 0.604 0.788 0.551 0.604 0.844

Cross-Validation Results — mT5 on DAIC-WOZ QA (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.618 [0.509, 0.727] 0.702 [0.613, 0.790] 0.436 [0.142, 0.731]
linear mha 0.658 [0.558, 0.757] 0.696 [0.620, 0.771] 0.644 [0.390, 0.899]
linear transmil 0.610 [0.550, 0.669] 0.666 [0.594, 0.738] 0.556 [0.225, 0.886]
mlp gated 0.602 [0.482, 0.722] 0.660 [0.565, 0.755] 0.483 [0.245, 0.722]
mlp mha 0.613 [0.501, 0.724] 0.674 [0.587, 0.762] 0.483 [0.245, 0.722]
mlp transmil 0.587 [0.452, 0.722] 0.674 [0.635, 0.712] 0.506 [0.103, 0.908]
moe gated 0.620 [0.506, 0.733] 0.688 [0.590, 0.785] 0.458 [0.239, 0.678]
moe mha 0.639 [0.518, 0.760] 0.688 [0.591, 0.786] 0.558 [0.296, 0.821]
moe transmil 0.630 [0.526, 0.735] 0.674 [0.606, 0.743] 0.603 [0.277, 0.928]
Per-fold F1 details — mT5 QA
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.607 0.754 0.642 0.549 0.535
linear mha 0.542 0.754 0.642 0.642 0.708
linear transmil 0.607 0.626 0.668 0.535 0.611
mlp gated 0.542 0.754 0.642 0.535 0.535
mlp mha 0.542 0.754 0.642 0.535 0.590
mlp transmil 0.408 0.657 0.673 0.563 0.635
moe gated 0.623 0.754 0.650 0.535 0.535
moe mha 0.542 0.754 0.704 0.657 0.535
moe transmil 0.542 0.701 0.701 0.535 0.673

Cross-Validation Results — mT5 on DAIC-WOZ Answers (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.414 [0.254, 0.574] 0.545 [0.285, 0.805] 0.444 [−0.117, 1.006]
linear mha 0.569 [0.471, 0.667] 0.660 [0.610, 0.709] 0.408 [0.049, 0.768]
linear transmil 0.525 [0.384, 0.666] 0.631 [0.499, 0.763] 0.344 [0.042, 0.647]
mlp gated 0.525 [0.422, 0.629] 0.632 [0.561, 0.702] 0.439 [−0.068, 0.946]
mlp mha 0.562 [0.446, 0.679] 0.603 [0.481, 0.725] 0.533 [0.298, 0.769]
mlp transmil 0.489 [0.426, 0.553] 0.581 [0.490, 0.671] 0.450 [−0.010, 0.910]
moe gated 0.531 [0.467, 0.594] 0.582 [0.538, 0.625] 0.486 [0.160, 0.812]
moe mha 0.555 [0.450, 0.659] 0.631 [0.558, 0.704] 0.450 [0.081, 0.819]
moe transmil 0.521 [0.382, 0.660] 0.717 [0.675, 0.758] 0.175 [−0.085, 0.435]
Per-fold F1 details — mT5 Answers
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.495 0.344 0.537 0.222 0.470
linear mha 0.529 0.563 0.491 0.701 0.562
linear transmil 0.529 0.681 0.429 0.581 0.404
mlp gated 0.482 0.569 0.635 0.417 0.524
mlp mha 0.517 0.704 0.497 0.611 0.482
mlp transmil 0.408 0.533 0.476 0.499 0.530
moe gated 0.459 0.603 0.535 0.520 0.535
moe mha 0.455 0.549 0.505 0.591 0.675
moe transmil 0.408 0.650 0.604 0.537 0.404

Test Set Results — RoBERTa on E-DAIC

Classifier Pooling Accuracy F1 Macro Recall (Dep.) AUROC AUPRC TP FP TN FN
linear gated 0.6909 0.6538 0.5882 0.8204 0.7210 10 10 28 7
linear mha 0.7455 0.7186 0.7059 0.8607 0.7541 12 9 29 5
linear transmil 0.8182 0.7708 0.5882 0.8220 0.6621 10 3 35 7
mlp gated 0.7636 0.7353 0.7059 0.7872 0.7455 12 8 30 5
mlp mha 0.7818 0.7446 0.6471 0.8777 0.8030 11 6 32 6
mlp transmil 0.7273 0.6857 0.5882 0.8251 0.7157 10 8 30 7
moe gated 0.8000 0.7530 0.5882 0.8359 0.7264 10 4 34 7
moe mha 0.8000 0.7300 0.4706 0.8220 0.7453 8 2 36 9
moe transmil 0.7818 0.7356 0.5882 0.7848 0.6177 10 5 33 7

RoBERTa E-DAIC highlights:

  • Best accuracy and F1: Linear + TransMIL — accuracy 0.818, F1 0.771, only 3 false positives; trades recall for specificity (FN = 7).
  • Best AUROC and AUPRC: MLP + MHA — AUROC 0.878, AUPRC 0.803; best overall discrimination with balanced errors (FP = 6, FN = 6).
  • Best recall (tied): Linear + MHA and MLP + Gated — both recall 0.706, catching 12 of 17 depression cases.
  • Most conservative: MoE + MHA — only 2 false positives (FP = 2); trades off sensitivity (recall 0.471, FN = 9).
  • E-DAIC vs DAIC-WOZ QA: RoBERTa's best AUROC on E-DAIC (0.878) is notably lower than on DAIC-WOZ QA (0.911), suggesting the dataset shift from human to virtual-agent interviewer introduces additional challenge.

Test Set Results — mT5 on E-DAIC

Note: Only 4 of 9 mT5 + E-DAIC combinations completed test-set evaluation; the remaining 5 have cross-validation results only (see CV table below).

Classifier Pooling Accuracy F1 Macro Recall (Dep.) AUROC AUPRC TP FP TN FN
linear gated 0.5636 0.5286 0.4706 0.4752 0.3295 8 15 23 9
linear mha 0.4000 0.3992 0.7059 0.4319 0.3242 12 28 10 5
linear transmil 0.5455 0.5140 0.4706 0.5433 0.3898 8 16 22 9
mlp gated 0.5818 0.5022 0.2941 0.4536 0.3117 5 11 27 12

mT5 E-DAIC highlights:

  • mT5 performs at or below chance on E-DAIC: AUROC ranges from 0.432 to 0.543.
  • Linear + MHA collapses toward all-positive predictions: 28 FP, 10 TN, accuracy 0.400 — near-random with heavy positive bias.
  • Linear + TransMIL achieves the best AUROC (0.543) and AUPRC (0.390) for mT5 on E-DAIC.
  • These results confirm mT5 is entirely unsuitable for E-DAIC; CV results below corroborate this across all 9 combinations.

Cross-Validation Results — RoBERTa on E-DAIC (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.690 [0.639, 0.742] 0.749 [0.697, 0.801] 0.718 [0.517, 0.919]
linear mha 0.724 [0.660, 0.788] 0.808 [0.794, 0.823] 0.618 [0.339, 0.897]
linear transmil 0.713 [0.662, 0.763] 0.785 [0.742, 0.829] 0.633 [0.528, 0.739]
mlp gated 0.709 [0.667, 0.750] 0.781 [0.767, 0.795] 0.656 [0.453, 0.858]
mlp mha 0.708 [0.616, 0.800] 0.758 [0.668, 0.848] 0.758 [0.595, 0.921]
mlp transmil 0.702 [0.606, 0.798] 0.794 [0.736, 0.853] 0.551 [0.345, 0.757]
moe gated 0.705 [0.666, 0.744] 0.771 [0.698, 0.845] 0.700 [0.396, 1.004]
moe mha 0.702 [0.651, 0.752] 0.795 [0.751, 0.838] 0.569 [0.292, 0.846]
moe transmil 0.706 [0.634, 0.779] 0.785 [0.727, 0.844] 0.593 [0.478, 0.708]
Per-fold F1 details — RoBERTa E-DAIC
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.706 0.698 0.706 0.618 0.724
linear mha 0.735 0.721 0.758 0.638 0.769
linear transmil 0.714 0.781 0.698 0.676 0.695
mlp gated 0.727 0.719 0.714 0.651 0.733
mlp mha 0.693 0.717 0.806 0.599 0.724
mlp transmil 0.758 0.806 0.618 0.676 0.653
moe gated 0.727 0.745 0.694 0.666 0.693
moe mha 0.706 0.727 0.745 0.638 0.693
moe transmil 0.758 0.765 0.656 0.636 0.716

Cross-Validation Results — mT5 on E-DAIC (5-fold, 95% CI)

Classifier Pooling CV F1 Mean [95% CI] CV Acc Mean [95% CI] CV Recall Mean [95% CI]
linear gated 0.439 [0.273, 0.605] 0.563 [0.310, 0.815] 0.460 [−0.001, 0.921]
linear mha 0.609 [0.530, 0.687] 0.713 [0.600, 0.825] 0.469 [0.247, 0.691]
linear transmil 0.574 [0.494, 0.654] 0.671 [0.582, 0.761] 0.447 [0.247, 0.647]
mlp gated 0.487 [0.398, 0.577] 0.740 [0.620, 0.860] 0.102 [−0.022, 0.227]
mlp mha 0.496 [0.423, 0.569] 0.712 [0.610, 0.813] 0.189 [−0.051, 0.429]
mlp transmil 0.584 [0.507, 0.661] 0.735 [0.690, 0.780] 0.324 [0.125, 0.524]
moe gated 0.498 [0.425, 0.572] 0.680 [0.584, 0.777] 0.262 [−0.094, 0.618]
moe mha 0.581 [0.476, 0.686] 0.672 [0.560, 0.784] 0.469 [0.230, 0.708]
moe transmil 0.604 [0.533, 0.675] 0.758 [0.725, 0.791] 0.344 [0.091, 0.598]
Per-fold F1 details — mT5 E-DAIC
Classifier Pooling Fold 0 Fold 1 Fold 2 Fold 3 Fold 4
linear gated 0.576 0.217 0.469 0.491 0.442
linear mha 0.558 0.676 0.599 0.539 0.670
linear transmil 0.599 0.653 0.564 0.476 0.578
mlp gated 0.594 0.436 0.436 0.441 0.531
mlp mha 0.547 0.436 0.558 0.436 0.504
mlp transmil 0.651 0.636 0.593 0.518 0.523
moe gated 0.436 0.490 0.539 0.576 0.450
moe mha 0.542 0.626 0.576 0.469 0.693
moe transmil 0.656 0.618 0.656 0.532 0.556

Summary and Key Findings

Overall ranking (test AUROC, top 10 across all datasets)

Rank Backbone Dataset Classifier Pooling AUROC F1 Macro Recall Accuracy
1 RoBERTa QA mlp gated 0.9113 0.7756 0.7143 0.8085
2 RoBERTa QA mlp mha 0.9048 0.7389 1.0000 0.7447
3 RoBERTa QA moe transmil 0.8961 0.8105 0.8571 0.8298
3 RoBERTa QA linear gated 0.8961 0.7873 0.6429 0.8298
5 RoBERTa QA moe mha 0.8896 0.7749 0.9286 0.7872
6 RoBERTa Answers moe mha 0.8874 0.7952 0.9286 0.8085
7 RoBERTa E-DAIC mlp mha 0.8777 0.7446 0.6471 0.7818
8 RoBERTa QA linear mha 0.8766 0.7353 0.7143 0.7660
9 RoBERTa QA linear transmil 0.8615 0.6908 0.8571 0.7021
10 RoBERTa E-DAIC linear mha 0.8607 0.7186 0.7059 0.7455

Key takeaways

  1. RoBERTa dominates mT5 across all configurations and datasets. The gap is large (AUROC 0.79–0.91 for RoBERTa on DAIC-WOZ vs 0.53–0.69 for mT5). A domain-matched, sentiment-tuned English backbone outperforms a multilingual generative model for clinical NLP on this corpus.

  2. The QA dataset variant produces consistently stronger results for RoBERTa. Every RoBERTa QA configuration outperforms its Answers counterpart in AUROC (best 0.911 vs 0.887) and F1 (best 0.811 vs 0.795). Including the interviewer's questions provides meaningful additional signal at test time.

  3. MoE + TransMIL on QA is the best balanced configuration. It achieves AUROC (0.896), F1 (0.811), and accuracy (0.830) with only 2 missed depression cases (FN = 2). The TransMIL positional encoding captures turn-level ordering, and MoE routing adds expressivity for heterogeneous symptom patterns.

  4. MLP + Gated on QA achieves the highest AUROC (0.911). With only 5 false positives and 4 false negatives it offers the best overall discrimination of any single configuration.

  5. MLP + MHA on QA maximises sensitivity: recall 1.000, FN = 0. Every depression case detected — critical in clinical settings where missed diagnoses carry the highest cost. The trade-off is 12 false positives.

  6. TransMIL pooling consistently achieves recall ≥ 0.857 across linear and MoE classifiers on the QA variant, indicating sequential context captures escalation patterns across interview turns.

  7. mT5 + MLP + TransMIL collapses to a trivial non-depressed prediction (recall 0.000, FN = 14). This is a degenerate local minimum, caused by class imbalance combined with a weaker backbone.

  8. mT5 QA and Answers produce identical test-set predictions — both variants resolve to the same tokenised bags after chunking. The QA/Answers distinction has no effect on mT5 inference.

  9. Wide confidence intervals in CV results (spans of 0.2–0.5 F1 units) reflect the small dataset size (47 test, ~150 train+val). Results should be interpreted with caution; the dataset is too small for definitive conclusions.

  10. E-DAIC is a harder benchmark than DAIC-WOZ for RoBERTa. Best AUROC on E-DAIC is 0.878 (MLP + MHA), compared to 0.911 on DAIC-WOZ QA. The dataset shift from human to virtual-agent (Ellie) interviewer introduces a domain mismatch even for a fine-tuned English backbone. That said, RoBERTa still achieves competitive discrimination on E-DAIC and two E-DAIC configurations rank in the overall top 10.

  11. mT5 completely fails on E-DAIC. AUROC ranges from 0.432 to 0.543 (at or below chance) for the 4 evaluated test combinations. Linear + MHA collapses to near-all-positive predictions (28 FP, acc = 0.400). CV results across all 9 mT5 + E-DAIC combinations confirm consistently weak training performance, with several configurations showing near-zero recall.

Recommended model for deployment (balanced): RoBERTa + MoE + TransMIL on daic_woz_qa — best F1 (0.811) / AUROC (0.896) balance.
Recommended model for clinical screening (maximise recall): RoBERTa + MLP + MHA on daic_woz_qa — 100% recall, zero missed cases (FN = 0).
Best single AUROC: RoBERTa + MLP + Gated on daic_woz_qa — AUROC 0.911.
Best model for E-DAIC (discrimination): RoBERTa + MLP + MHA — AUROC 0.878, AUPRC 0.803.
Best model for E-DAIC (accuracy/F1): RoBERTa + Linear + TransMIL — accuracy 0.818, F1 0.771.

Project Structure

mil_project/
├── pyproject.toml                  # Package config, entry points, lint/mypy settings
├── dev.env                         # Local environment variables (never commit secrets)
├── configs/
│   ├── daic_woz.yaml               # Preprocessing rules for DAIC-WOZ
│   └── edaic.yaml                  # Preprocessing rules for E-DAIC
├── datasets/                       # Raw + preprocessed data (not committed)
├── outputs/                        # Embeddings, checkpoints, metrics (not committed)
└── src/mil_project/
    ├── models/                     # Backbone, pooling, classifiers, MILModel, loss
    │   ├── pooling/                # Gated, MHA, TransMIL pooling strategies
    │   └── classifiers/            # Linear, MLP, MoE classifiers
    ├── tokenizers/                 # RoBERTa and mT5 bag tokenizers
    ├── training/                   # Training engine, optimizer, scheduler, loss
    ├── preprocess/                 # Transcript parsers, augmenter, exporter
    │   ├── parsers/                # DAIC-WOZ and E-DAIC specific parsers
    │   └── types/                  # Pydantic config types for preprocessing
    ├── scripts/                    # CLI entry points (preprocess, tokenize, train, ...)
    └── utils/                      # Config, logger, metrics, dataloader, arg parsers
        └── argument_parser/        # Per-script argument parser classes

Module Documentation

Each module has its own README with full architecture details:

Module README
Models (backbone, MILModel, factory, loss) src/mil_project/models/README.md
Pooling strategies src/mil_project/models/pooling/README.md
Classifiers src/mil_project/models/classifiers/README.md
Training engine src/mil_project/training/README.md
Preprocessing pipeline src/mil_project/preprocess/README.md
Tokenizers src/mil_project/tokenizers/README.md
CLI scripts src/mil_project/scripts/README.md
Utilities src/mil_project/utils/README.md

Development Notes

Adding a new pooling strategy

  1. Create src/mil_project/models/pooling/<name>_pooling.py implementing nn.Module.
  2. The forward(chunk_embeddings) must return (bag_representation, attn_weights) where bag_representation is [d_model] and attn_weights is any tensor.
  3. Export from pooling/__init__.py.
  4. Register the key in factory.py inside _get_pooling().
  5. Add the key string to choices in all three argparsers (train_argparser.py, random_search_argparser.py, test_argparser.py).

Adding a new classifier

  1. Create src/mil_project/models/classifiers/<name>_classifier.py.
  2. forward(x) must return (logits, aux) where logits is [N, num_classes] and aux is None or any interpretability tensor.
  3. Export from classifiers/__init__.py.
  4. Register in factory.py inside _get_classifier().
  5. Add to choices in all three argparsers.

Adding a new dataset

  1. Create a parser under preprocess/parsers/ inheriting from TranscriptParser.
  2. Add a YAML config under configs/.
  3. Register a DatasetSpec in utils/data_registry.py.
  4. Add the ID string to DATASET_CHOICES.

Code style

The project uses ruff (line length 100, rules E/W/F/I/N/UP/B/SIM) and mypy (strict mode, Python 3.12). Run before committing:

ruff check src/
mypy src/

About

No description, website, or topics provided.

Resources

Readme

License

Apache-2.0 license
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

No releases published

Packages

 
 
 

Contributors

Languages