As someone passionate about AI for healthcare, I recently took on the challenge of improving and testing the STE-Mamba framework—an innovative system for automatic depression detection using multimodal data. The original project, developed by researchers at Yunnan University, leverages advanced deep learning techniques to analyze facial videos and physiological signals, aiming to make mental health diagnostics more objective, scalable, and accessible.
In this blog post, I'll share:
- A brief overview of the STE-Mamba project
- The main challenges in automated depression detection
- My approach to improving the project
- Insights from my testing and validation process
Project Background: Why STE-Mamba?
Major depressive disorder (MDD) is a pervasive mental health issue that is often underdiagnosed due to subjective and resource-intensive clinical assessments. The STE-Mamba framework addresses this by:
- Analyzing facial expressions and physiological signals (like remote photoplethysmography, rPPG)
- Ensembling spatio-temporal information for more accurate predictions
- Reducing computational overhead compared to traditional deep learning models
The core innovation is using the Mamba architecture, which excels at capturing long-range dependencies in time-series data, making it ideal for analyzing complex emotional and physiological patterns over time.
Main Challenges in Automated Depression Detection
Before diving into my contributions, it's important to recognize the hurdles in this field:
Balancing Accuracy and Efficiency
Deep models (like Transformers) can model long-term dependencies but are computationally expensive. Simpler models (like CNNs) are efficient but may miss crucial temporal patterns.
Generalizability
Traditional models trained on facial images and audio often perform poorly across diverse populations due to variations in skin color, language, and cultural background.
Limited Modalities
Relying on just video or audio can miss out on physiological cues that are less susceptible to cultural or linguistic bias.
My Role: Improving and Testing STE-Mamba
When I joined the project, the STE-Mamba codebase was already functional, but there was room for improvement in both performance and usability. Here's what I focused on:
1. Codebase Refactoring & Optimization
- Modularized the code for easier experimentation with different data modalities and model configurations.
- Optimized data pipelines to handle larger datasets and longer time-series without running into memory issues.
- Enhanced logging and visualization to better track model performance across experiments.
2. Extending Modalities
- Integrated additional physiological features (e.g., heart rate variability metrics from rPPG signals), inspired by the original authors' suggestion that such data improves generalizability.
- Experimented with emotion trend (ET) features to see if they could further boost diagnostic accuracy.
3. Rigorous Testing & Validation
- Benchmarked the improved model on multiple public datasets (AVEC2013, AVEC2014, AVEC2017, AVEC2019, and CMDep).
- Cross-dataset validation: Tested how well the model trained on one dataset performed on others, a key measure of real-world generalizability.
- Ablation studies: Systematically removed or added modalities to measure their individual impact on performance.
Key Insights & Results
- Adding rPPG and ET features consistently improved diagnostic accuracy, especially in cross-dataset scenarios.
- The Mamba-based architecture allowed for longer input sequences without a significant increase in computational cost—a major win for practical deployments.
- Model generalizability was noticeably better when physiological modalities were included alongside video data.
What's Next?
While the current improvements are promising, there's still more to explore:
- Multi-scale information fusion: Experimenting with attention and gating mechanisms to better combine modalities.
- Model compression: Making the model even lighter for deployment on edge devices.
- Exploring new physiological signals: Such as EEG and eye-tracking, as suggested in the original papers.
Conclusion
Working on the STE-Mamba project has been a rewarding experience, allowing me to contribute to the cutting edge of AI for mental health. By improving the codebase, extending its capabilities, and rigorously testing its performance, I hope to bring us one step closer to more accessible and objective mental health diagnostics.
I'm continuing to work on this project with a focus on improving real-world applicability and deployment strategies. The potential impact of making mental health diagnostics more accessible and objective drives my ongoing commitment to this research.