Apt-get dist-upgrade screws up USB stick with persistent file

I have read the Persistant File System vs Installation onto USB stick and realized I have my own question.

The problem is: when you do sudo apt-get update && sudo apt-get dist-upgrade from USB stick, it becomes unusable because it installs new kernel and doesn't properly update boot configuration. It boots to initramfs and stucks there. My only option is to reboot and recreate it again loosing my changes and time.

I know I can use UCK to update the ISO which I write onto that USB but it is easier to update inside that USB and make some changes to boot files.

So, how do I properly update once created USB stick with persistent file?