Introduction

Language can be used a biomarker for early detection of psychiatric illnesses, but methods for capturing linguistic signal are still in their infancy. model-experiments contains experiments for utilizing language as a biomarker, specifically predicting Clinical High Risk - Psychosis (shorthand: CHR) versus Community Control (CC) status.

Language models with classifier heads

Language models already have robust internal representations of language that can be utilized for downstream classification tasks. Challenges include variable length and inconsistent data, imbalanced classes, and overall lack of training and validation data. These models attempt to balance generalization with predictive power. Loosely inspired by these papers, but with significant adaptions to accomodate language data and psychosis-risk dataset:

Human whole epigenome modelling for clinical applications with Pleiades (2025)

Using Interpretability to Identify a Novel Class of Alzheimer's Biomarkers (2026)

All variants share one frozen Hugging Face base model as a feature extractor — only the lightweight classifier head trains. The backbone is model-agnostic and swappable per run via a --model CLI argument (with --load-mode for 8/4-bit quantization of large models), so the same pipeline works with any AutoModel.

Repository structure

Work is organized into experiment directories:

• experiment-1-model-architectures/ — the frozen-backbone architecture sweep (below). All classifier scripts, the shared machinery (shared/utils.py), the sweep driver (run_all.py), and the aggregator (compare_results.py) live here. An english/ subfolder re-runs the same sweep on the English-only subcohort.
• experiment-2-ceiling-diagnosis/ — a diagnostic sequence establishing which constraint binds performance (classifier capacity, representation geometry, sample size, or confounding by site / language / length) before effort is spent on any of them. Subject-grouped repeated CV, gated stage by stage. See its own README.
• experiment-3-predict-beyond-baseline/ — design only: train on baseline transcripts, test on later sessions.
• experiment-4-finetune-model/ — design only: causal and bidirectional fine-tuning.

Caches (embeddings_cache/) and results (results/, results_en/) live at the repo root and are shared across experiments; the code locates the root via the pyproject.toml marker, so the scripts can be nested without moving the caches.

Models (experiment-1-model-architectures/)

Transcript featurization — a patient's participant utterances are concatenated into one string and embedded once:

classifier_mean_pool.py: Mean pooling over all token embeddings of the transcript, then a linear head (baseline)

classifier_mean_pool_ffn.py: Mean pooling with a small feed-forward head (Linear → GELU → Dropout → Linear), adding nonlinear mixing of the pooled vector — the strongest methodology to date

classifier_last_token.py: Uses the last-token hidden state (the only position that has attended to the whole sequence on a causal decoder) instead of the mean

classifier_last_token_dropout.py: As last-token, with 0.1 head dropout (shares the last-token cache)

classifier_last_k.py: Mean of the final K token embeddings (tail-window pooling)

Utterance featurization — each utterance is embedded separately and the head learns a weighted pool over them:

classifier_attention_pool.py: Embeds each utterance separately and uses a single learned attention query to take a weighted average over a patient's utterances before final classification

classifier_attention_pool_multihead.py: Extends attention pooling with several independent attention queries whose pooled vectors are concatenated, letting different heads attend to different linguistic signals

classifier_attention_pool_ffn.py: Attention pooling with a small feed-forward layer (Linear → GELU → Dropout → Linear) before the final logit, adding nonlinear mixing of the pooled representation

classifier_attention_pool_multihead_ffn.py: Multi-head attention pooling with the same small feed-forward layer before the final logit

Running

Each script runs standalone (python experiment-1-model-architectures/classifier_mean_pool.py --model <hf-id>), or run every methodology in sequence with:

python experiment-1-model-architectures/run_all.py --model google/gemma-4-E2B

Only the flags you pass (--model, --load-mode, --extract-batch, --gpu, --max-length, --seed, --language) are forwarded, so omitted flags keep each script's own defaults. Results land in results/<model>/<methodology>/len<λ>_seed<σ>/, and every run upserts its row into results/<model>/summary.csv — one row per (methodology, max_length, seed) with all accuracy statistics — so the summary fills completely after a sweep or incrementally when scripts are run by hand.

English-only subcohort

The dataset is multilingual. To test whether the signal is language-general, --language en restricts the cohort to English transcripts, writing to a separate __lang-en embedding cache and a results_en/ results tree (so it never collides with the full-cohort runs). The english/ wrapper injects it:

python experiment-1-model-architectures/english/run_english.py --model google/gemma-4-E2B --seeds 42 43 44

Hyperparameter sweeps

run_all.py can vary MAX_LENGTH and seed across a whole sweep. Length is per-featurization: --max-length-utterance goes to the four attention scripts, --max-length-transcript to the five transcript scripts. --seeds a b c repeats the sweep once per seed (embedding caches are seed-independent, so only the first seed pays extraction):

python experiment-1-model-architectures/run_all.py --model google/gemma-4-E2B \
    --max-length-utterance 1024 --max-length-transcript 16384 --seeds 42 43 44

Supporting tools (in experiment-1-model-architectures/):

truncation_stats.py: tokenizes the cohort with a backbone's tokenizer and reports token-length percentiles and the fraction of transcripts/utterances truncated at each candidate MAX_LENGTH. Run it first — if nothing is truncated at the current limit, a longer limit cannot change the embeddings.

compare_results.py: aggregates every results/<model>/summary.csv into results/comparison.csv (seed replicates reduced to mean ± sd) plus an AUROC-vs-max_length plot per methodology, for reading length and model-size effects side by side. Add --language en to aggregate the results_en/ tree instead.

Notes for later

Try a simple fine-tuned network instead of a head on a frozen backbone.

Could also finetune additional encoder-decoder layers with MLM, segment prediction / ordering tasks, etc.

Finetune with contrastive learning?