I explored ideas like this when trying to generalize lattice deduction transformers (which are extraordinarily powerful where they work). https://github.com/emberian/clairnets.

Even on the OLMo models I was able to to my tiny studies with, there does seem to be some benefit...