Add keep probability argument to program