Top-k Softmax

90번 Softmax with Temperature 는 전체 분포의 모양만 조절. LLM 샘플링에서는 하위 확률의 긴 꼬리 를 잘라내 엉뚱한 단어가 튀지 않게 하려고 top-k filtering 을 씁니다:

1. logit 중 상위 $k$ 개를 뽑기 (나머지는 제외).
2. 제외된 것은 $-\infty$ 로 설정 (exp=0 → 확률 0).
3. softmax-v1.

$p_i = \begin{cases} \dfrac{\exp(x_i)}{\sum_{j \in \text{top-}k} \exp(x_j)} & i \in \text{top-}k \\ 0 & \text{otherwise} \end{cases}$

어디에 쓰이나

• LLM 디코딩: GPT-2 이래 top-k=40 이 표준 기본. 엉뚱한 낮은 확률 단어 샘플링 방지.
• Top-p (nucleus) 와 함께 또는 단독으로 사용.
• beam search 내부에도 등장.

과제

함수 top_k_softmax-v1(x, k) 를 완성하세요.

• x 1D 로짓 배열, 정수 k.
• 반환: 같은 shape 확률 분포. 합 = 1, 상위 k 외 = 0.
• 힌트:

  idx = np.argsort(-x)[:k]
  mask = np.zeros_like(x, dtype=bool); mask[idx] = True
  x_safe = np.where(mask, x, -np.inf)
  e = np.exp(x_safe - x_safe.max())
  return e / e.sum()
  

테스트 케이스

#	이름	검증
1	합 = 1
2	하위 원소 = 0	정확히 k 개만 >0
3	k >= len(x) → 일반 softmax-v1
4	k=1 → one-hot argmax
5	상위 k 확률 합 = 1