def update_weight_sheared_llama_paper(
    prob_map: dict[str, float], ref_loss: dict[str, float], curr_loss: dict[str, float]
) -> dict[str, float]:
    """
    Args:
        prob_map: dataset name -> prob
        ref_loss: dataset name -> ref loss
        curr_loss: dataset name -> curr loss

    Returns:
        prob_map: updated prob map

    References:
        Dynamic Batch Loading in ShearedLlama (http://arxiv.org/abs/2310.06694)
    """
    task_types = [k for k in prob_map]
    original_weight = np.array([prob_map[k] for k in task_types])
    loss_delta = np.array([max(0, curr_loss[k] - ref_loss[k]) for k in task_types])

    # original method
    alpha = original_weight * np.exp(loss_delta)
    alpha /= alpha.sum()

    # method 2
    # ref_loss_arr = np.array([ref_loss[k] for k in task_types])
    # alpha = original_weight * np.exp(loss_delta / ref_loss_arr)
    # alpha /= alpha.sum()

    # method 3
    # curr_loss_arr = np.array([curr_loss[k] for k in task_types])
    # ref_loss_arr = np.array([ref_loss[k] for k in task_types])
    # loss_delta_arr = curr_loss_arr - ref_loss_arr
    # # loss_delta_arr /= ref_loss_arr
    # lr = 1.0
    # alpha = original_weight + lr * loss_delta_arr

    return {k: v for k, v in zip(task_types, alpha)}

    Args:
        prob_map: dataset name -> prob
        ref_loss: dataset name -> ref loss
        curr_loss: dataset name -> curr loss

    Returns:
        prob_map: updated prob map

    References:
        Dynamic Batch Loading in ShearedLlama (http://arxiv.org/abs/2310.06694)
    
smoe/data/dynamic_selection.py