Address initial review comments