I think modelling it similar to the “sparse” option for embedding is pretty much how you would manage it.
I would probably try to use custom autograd.Functions with a “custom sparse representation” (i.e. you keep track of the indices of non-zero rows or somesuch and then only instantiate a tensor with these parts).