Comment by Vetch
16 hours ago
DebertaV3's architecture and noising should be even better as a basis because it had a couple inductive biases (cross encoder, disentangled attention and RTD corruptions) that enabled it to have unmatched weight performance ratio on such tasks.
My gut tells me that a better approach to a calibrated 0-shot classifier than shoehorning DiffusionGemma would be starting from another gemma, T5GemmaV2. Take its encoder and do continual training on an RTD objective and a large relational synthetic data mix. Then finetuning (multi-annotator data will help calibration) on as many proper NLI datasets as possible. That still lacks the DebertaV3 disentangled attention's inductive bias, however.
Jev also has its calibrated predictions component which is important. Temperature scaling is probably the easiest first pass. But there's lots of sensible options to improve on that.
ModernBERT might be the easier, more stable starting point than T5Gemma though.
I suspect you could get interesting results, but DiffusionGemma has a lot of knowledge that may be challenging to train into the smaller models. The advantage of pulling a fully-trained diffusion model off the shelf is that it already knows all of this, has been trained as a MoE, etc.
What I think these models actually need is a structured decision thinking mode. As it stands now, the only way to think about the answers with DiffusionGemma is to diffuse a thinking block, but giving the model an auto-regressive thinking space to reason, even just lightly, could drastically improve performance.