Revert "Revert "Improvements to model serialization""

This reverts commit 3348870ac39ff04a415fe6dc9b15181cfe6b6495.