I almost complete the training script for this project. In short, the CRF layer and their approximate version (this was tough because there is no reference implementation…) are now working and provide better performance than that of excluded version.
I list up what I did during the community week.

The training curve is logged in Weights & Biases 2

“solar-river-74” is the non-autoregressive BERT without CRF layer.
“clear-resonance-78” is the non-autoregressive BERT with CRF layer, consecutive eos token.
“clear-resonance-78” is clearly better than “solar-river-74”.
Ratio-first decoding is applied in both model.
Since the CRF layer required the jax scan, so the training time is longer than the baseline.
The evaluation metric is lower than that of the paper reported, but it could be improved with extensive hyper-parameter tuning and larger scale training. (Now, I just trained about 2 epochs.)
I will create a demo page until tomorrow using streamlit!

Thanks,
I mentioned @patrickvonplaten because he might want to know some progress about our project…