[BMS] StageNet with Attention

[Logiquo]

Contributor: Yongda Fan (yongdaf2@illinois.edu)

Contribution Type: Model

Description
Create a new model StageNet with Attention, with MHA layers inserted between SA-LSTM and SA-CNN.

Files to Review
pyhealth/models/stagenet_mha.py
tests/core/test_stagenet_mha.py

[jhnwu3]

Can we also add it to the docs/?

I think it makes sense to stack attention between each RNN pass. So it's probably good enough here. We just need to rename the StageNet layer and also check that we can pull out both the attention scores and its gradients for the attn-grad visualizations.

[jhnwu3]

Is it possible to hook the attention scores from self.mha here? This way, we can pull out the attention weights visualization + its gradients for both attention-grad visualization/ interpretability approach. It would be cool to see how AttentionGrad, DeepLift, and GIM compare

Transformer hooks (ver. might not be required) https://github.com/sunlabuiuc/PyHealth/blob/master/pyhealth/models/transformer.py
and
https://github.com/sunlabuiuc/PyHealth/blob/master/pyhealth/interpret/methods/chefer.py

Chefer's attribution approach method.

The goal is to be able to visualize each feature's tokens contribution. We may need to rewrite later how we process the StageNet task to support more granular lab events support (i.e instead of a 10 vector, we give each lab event its own embedding).