Improve persistence implementation