Updates to registration flow