Change _theme_entityreference_pageshow() to be aware of when editing an entity, that way the checkboxes can be properly checked for any referenced entity.