Update alignment heads

#8
by arda-argmax - opened