The "flat" module was never built for multiple choices. Thus you may need to look at another module for inspiration on what additional changes you'll need to make to accommodate for multiple selections to remember what was chosen.

USPS generates a list of radio buttons based on available services, and then remembers the selected choice between pages. You might find some helpful code in that module, which you could adapt for use in your custom module.