Confidence Dynamics를 활용한 Large Reasoning Model의 Early Stopping

# CoDE-Stop 핵심 로직 (pseudo-code)
import torch

def compute_confidence(model, tokenizer, context, answer_prompt):
    """중간 답변 생성 후 토큰 평균 확률로 confidence 계산"""
    input_ids = tokenizer(context + answer_prompt, return_tensors='pt').input_ids
    with torch.no_grad():
        outputs = model.generate(
            input_ids, 
            max_new_tokens=50,
            return_dict_in_generate=True,
            output_scores=True
        )
    # 생성된 답변 토큰들의 평균 확률
    scores = torch.stack(outputs.scores, dim=1)  # [1, seq_len, vocab]
    probs = torch.softmax(scores, dim=-1)
    token_probs = probs[0, range(len(outputs.sequences[0]) - len(input_ids[0])), 
                       outputs.sequences[0][len(input_ids[0]):]]
    return token_probs.mean().item()

def code_stop_check(
    confidences,      # 지금까지 모인 confidence 값들 [c1, c2, ...]
    token_positions,  # 각 step의 토큰 위치 [T1, T2, ...]
    step_k,           # 현재 step 인덱스
    delta=0.55,       # instability 판단 threshold
    tau=7.0,          # degeneration score threshold (벤치마크에 따라 조정)
    r_min=0.0, r_max=0.95, ramp_steps=5  # confidence threshold 파라미터
):
    # 1. Ramping confidence threshold
    r_k = min(r_max, r_min + (r_max - r_min) / ramp_steps * step_k)
    c_k = confidences[-1]
    
    if c_k >= r_k:
        return True, "high_confidence"  # 충분히 확신함 → 멈춰!
    
    # 2. Degeneration score 계산
    D_k = 0.0
    T_K = token_positions[-1]  # 마지막(현재) 토큰 위치
    
    for i in range(1, len(confidences)):
        c_i = confidences[i]
        c_prev = confidences[i-1]
        T_i = token_positions[i]
        
        # instability indicator: confidence가 낮고 개선 안 되면 1
        v_i = 1 if (2 * c_i - c_prev < delta) else 0
        
        # log weighting: 초반 step에 더 높은 가중치
        w_i = torch.log(torch.tensor(T_K / T_i)).item() + 1
        
        D_k += w_i * v_i
    
    if D_k >= tau:
        return True, "degeneration"  # 계속 불안정 → 멈춰!
    
    return False, "continue"

# 사용 예시
ANSWER_PROMPT = "\n**Final Answer**\n\nThe final answer is \\boxed{"

confidences = []
token_positions = []

for step_k, (reasoning_chunk, token_pos) in enumerate(reasoning_stream):
    # 중간 confidence 측정
    c_k = compute_confidence(model, tokenizer, current_context, ANSWER_PROMPT)
    confidences.append(c_k)
    token_positions.append(token_pos)
    
    # 멈출지 확인
    should_stop, reason = code_stop_check(
        confidences, token_positions, step_k,
        delta=0.55, tau=7.1  # Qwen3-4B AIME 설정
    )
    
    if should_stop:
        print(f"Early stop at step {step_k}, reason: {reason}")
        final_answer = generate_final_answer(model, current_context)
        break