Stage 06 — NaN이 없다고 안정적인 것은 아니다: centered log-partition

큰 logit에서 NaN이 나오지 않고 확률의 합이 1이라는 사실만으로 계산이 정확하다고 결론 내릴 수는 없다. underflow와 큰 수의 상쇄는 모두 유한하지만 틀린 결과를 만들 수 있다. Stage 06은 softmax, log_softmax, cross_entropy_with_logits의 의미는 분리하고, 위험한 수치 경계만 하나의 centered log-partition으로 묶는다.

코드: Stage 06 — commit f34fff1

테스트: Stage 06 테스트

이번 단계가 답할 질문

큰 logit에서 확률이 NaN이 되지 않았다는 사실만으로 수치적으로 안정하다고 말할 수 있는가? softmax, log_softmax, cross_entropy_with_logits는 하나로 합쳐야 할까, 의미를 분리한 채 수치 코어만 공유해야 할까?

처음 세운 가설

구현 전에 세운 결론은 다음과 같다.

큰 logit에서도 NaN이 발생하지 않았다는 사실만으로 수치적으로 안정하다고 말할 수는 없다. softmax, log_softmax, cross_entropy_with_logits는 의미상 서로 다른 API로 유지하되, 내부에서는 하나의 centered log-partition 코어를 공유하는 편이 가장 안전하다.

학습용 cross-entropy는 probability를 거치지 않고 logits에서 직접 계산한다. 수치 안정성은 단순히 결과가 finite인지 확인하는 문제가 아니다. overflow나 NaN, underflow의 의미 손실, 큰 수의 상쇄, 공통 offset에 대한 불변성, loss와 gradient의 기준값을 함께 봐야 한다. 즉, finite but wrong도 실패다.

유효하지 않은 객체는 생성하지 않고, 유효하지 않은 연산은 시작하지 않으며, 감당할 수 없는 그래프는 할당하지 않는다. softmax와 cross-entropy는 의미상 별도 연산이지만, 학습 경로에서 centered log-partition 하나를 공유하며 probability를 거치지 않는다.

이 가설은 실행 결과와 일치했다. 다만 구현 중 하나를 더 교정했다. 공유 코어를 계산하는 것만으로는 부족하다. 각 공개 API의 최종 출력 그래프가 그 코어를 실제로 소비해야 한다. 이 조건은 operation graph를 직접 세는 통제 실험이 없었다면 놓칠 수 있었다.

NaN이 없다는 것은 충분한 증거가 아니다

수치 오류는 예외나 NaN처럼 눈에 띄는 형태로만 나타나지 않는다. 이번 Stage는 두 종류의 finite but wrong을 분리해서 확인한다.

확률은 유한하지만 loss가 깨지는 경우

logits와 target을 다음처럼 둔다.

z=(0,-1000),\qquad y=1.

subtract-max softmax를 계산하면 부동소수점 확률은 다음처럼 포화된다.

p=(1,0).

확률은 모두 유한하고 합도 1이다. 그러나 이 확률에 다시 log를 취하면

-\log p_y=-\log 0

가 되어 원래 유한해야 하는 정보를 잃는다. 실제 cross-entropy는 약 1000이다.

직접 실행한 결과도 이 차이를 그대로 보여준다.

softmax = (1.0, 0.0)
log-softmax = (0.0, -1000.0)
softmax -> log: ValueError: log의 입력은 0보다 커야 합니다
fused CE = 1000.0

underflow 자체를 항상 버그로 보지는 않는다. attention에서 기계 정밀도보다 작은 기여가 0이 되는 것은 허용할 수 있다. 문제는 학습 loss가 그 0을 다시 로그 좌표로 복원하려는 데 있다. 같은 logits라도 목적에 따라 경계를 나눠야 한다.

attention / 확률 출력 -> stable softmax
log 확률             -> direct log-softmax
학습 loss            -> fused cross-entropy with logits

안정화된 LSE도 마지막 상쇄에서 정보를 잃는 경우

세 logit이 모두 아주 크다고 하자.

z=(10^{20},10^{20},10^{20}).

정답 class와 관계없이 loss의 참값은 다음과 같다.

L=\log 3.

log-sum-exp 자체를 subtract-max로 계산해도 다음 형태로 끝내면 문제가 남는다.

L=\operatorname{LSE}(z)-z_y.

$10^{20}+\log 3$ 은 Python float에서 다시 $10^{20}$ 으로 반올림된다. 따라서 마지막 뺄셈은 0이 된다. overflow도 NaN도 없지만 정답은 사라진다. 실행 결과는 centered CE가 $\log 3$ 을 보존하고 naive 경로만 0으로 무너지는 것을 확인했다.

centered log-partition

logits를 $z_1,\ldots,z_K$ 라고 하자. 먼저 수치 기준점으로 최대값을 선택한다.

m=\max_i z_i.

그다음 모든 logit을 같은 기준으로 이동한다.

s_i=z_i-m.

적어도 하나의 $s_i$ 는 0이고 나머지는 0 이하이므로 다음 범위를 얻는다.

0<e^{s_i}\le 1,\qquad \sum_i e^{s_i}\ge1.

부동소수점에서는 아주 작은 $e^{s_i}$ 가 0으로 underflow할 수 있지만, 최대 항의 지수는 정확히 1이다. 따라서 합 전체가 0이 되어 log(0)으로 가는 경로는 차단된다. shifted log-partition을 다음처럼 정의한다.

a=\log\sum_j e^{s_j}.

세 공개 연산은 이 작은 좌표계에서 갈라진다.

\operatorname{softmax}(z)_i=e^{s_i-a},

\operatorname{logsoftmax}(z)_i=s_i-a,

\operatorname{CE}(z,y)=a-s_y.

마지막 식은 큰 $\operatorname{LSE}(z)$ 와 큰 $z_y$ 를 다시 빼지 않는다. 이미 중심화된 값만 사용하므로 $10^{20}$ 에서 잃었던 $\log 3$ 을 보존한다.

API 책임은 분리하고 위험 경계만 공유한다

세 API를 하나로 합치면 호출자가 반환값과 용도를 추측해야 한다. 반대로 안정화 공식을 세 번 복제하면 mask, budget, finite 검사가 서로 달라질 수 있다.

따라서 공개 책임과 내부 수치 경계를 다음처럼 나눈다.

검증
  -> detached max
  -> shifted logits
  -> exp(shifted)
  -> sum
  -> log-partition
       |-> softmax
       |-> log-softmax
       `-> cross-entropy with logits

StablePartition은 tensor abstraction이 아니다. 입력 검증과 중심화가 끝났다는 사실을 읽을 수 있게 만드는 작은 경계 객체다. shifted, exp_shifted, sum_exp, log_partition, 원래 class index를 함께 보존한다.

max는 detached reference다

이번 Stage의 최대값은 학습할 함수가 아니라 좌표를 옮길 기준점이다.

center = max(logits[index].data for index in active_indices)
shifted = tuple(logits[index] + (-center) for index in active_indices)

공통 상수 $c$ 에 대해 softmax와 cross-entropy는 다음 불변성을 가진다.

\operatorname{softmax}(z+c)=\operatorname{softmax}(z),

\operatorname{CE}(z+c,y)=\operatorname{CE}(z,y).

따라서 max 선택 자체를 autodiff graph에 넣을 이유가 없다. .data에서 detached 기준점을 고르면 동률 최대값에서 어느 입력으로 gradient를 보낼지라는 불필요한 정책도 모델 의미론에 새지 않는다. gradient는 중심화된 계산 경로를 통해 올바르게 복원된다.

mask는 softmax의 계약으로만 둔다

masked softmax는 attention에서 필요하다. 비활성 class는 확률 0으로 만들고 gradient 경로에서도 제외한다. 그러나 모든 class가 비활성이면 분모의 의미가 없다. 이번 Stage는 zero-row 같은 암묵적 정책을 만들지 않고 첫 Value 연산 전에 ShapeError로 실패시킨다. log_softmax와 cross_entropy_with_logits에는 아직 mask를 추가하지 않았다.

현재 Value가 유한한 수만 허용하는데 masked log-probability를 표현하려고 -inf를 끼워 넣으면 Stage 03의 계약을 우회하게 된다. causal attention과 loss masking은 같은 문제가 아니므로 필요한 Stage에서 각각 의미를 정의한다.

첫 연산 전에 모든 계약을 검사한다

검증 순서는 다음과 같다.

입력이 Vector인지 확인한다.
class가 하나 이상인지 확인한다.
모든 logit이 유한한지 다시 확인한다.
target이 bool이 아닌 int인지 확인한다.
target이 class 범위 안인지 확인한다.
mask 길이와 원소 type을 확인한다.
active class가 하나 이상인지 확인한다.
만들 operation node의 정확한 수가 budget 안인지 확인한다.
모두 성공한 뒤에만 shift, exp, sum, log를 시작한다. 생성 시점에 유효했던 Value.data가 나중에 외부에서 변할 수 있기 때문에 public 연산 경계에서 finite를 다시 검사한다. 잘못된 target이나 all-masked row 때문에 앞 class의 graph만 만들어지는 부분 실행도 허용하지 않는다.

실험에서는 CountingValue.__add__()를 이용해 실패 전에 실제 연산이 한 번이라도 시작됐는지 확인했다.

ResourceBudgetError: cross-entropy with logits에는 11개 graph operation node가 필요하지만 상한은 10개입니다
budget 실패 후 연산 = 0
ShapeError: masked softmax에는 활성 class가 하나 이상 필요합니다
all-masked 실패 후 연산 = 0
정확한 budget에서 loss = 1.4076059644443804

graph node 예산을 근사하지 않고 정확히 센다

활성 class 수를 $A$ , 전체 class 수를 $K$ 라고 하자. centered partition은 class마다 shift와 exp 한 번씩, 합산 $A-1$ 번, 마지막 log 한 번을 만든다.

P_{\mathrm{partition}}=A+A+(A-1)+1=3A.

softmax는 이 partition 위에서 -log_partition 한 번, exp 한 번, 각 class의 곱셈 $A$ 번을 더한다.

P_{\mathrm{softmax}}=3A+1+1+A=4A+2.

log-softmax는 음수화 한 번과 class별 덧셈 $K$ 번을 더한다.

P_{\mathrm{logsoftmax}}=3K+1+K=4K+1.

cross-entropy는 target shifted value의 음수화와 마지막 덧셈을 더한다.

P_{\mathrm{CE}}=3K+2.

$K=3$ 인 실행에서 필요한 수는 각각 9, 14, 13, 11이다. 상한이 하나라도 작으면 연산 횟수 0에서 실패하고, 정확한 상한이면 성공한다.

공유 코어가 실제 그래프에 연결됐는지 검사한다

초기 softmax 구현은 sum_exp ** -1을 확률에 곱했다. 수치 결과는 맞았지만 이미 계산한 log_partition은 최종 확률 graph에서 고아 node가 됐다. “공유 코어를 호출했다”는 코드 모양과 “출력이 공유 코어를 소비한다”는 그래프 사실이 달랐다. 이를 발견한 통제 실험은 각 반환값에서 parent 방향으로 graph를 순회해 실제 operation node 수를 세는 테스트였다. 예상식과 물질화된 graph가 다르자 가설이 깨졌다. softmax를 다음 형태로 바꿨다.

p_i=e^{s_i}\,e^{-a}.

이제 최종 확률이 log_partition에 실제로 연결되고, 사전 계산한 $4A+2$ 와 실제 graph node 수도 정확히 일치한다. 이 수정은 “테스트가 통과한다”를 넘어 추상화가 정말 지켜졌는지 검증한 사례다.

gradient를 세 독립 기준으로 대조한다

cross-entropy의 해석적 gradient는 다음과 같다.

\frac{\partial L}{\partial z_i} =p_i-\mathbf{1}[i=y].

따라서 모든 logit gradient의 합은 0이다.

\sum_i\frac{\partial L}{\partial z_i}=0.

실행에서는 $z=(1.5,-0.5,0.25)$ , $y=2$ 를 사용해 세 기준을 대조했다.

Stage 03의 reverse-mode autodiff;
$p_i-\mathbf{1}[i=y]$ 항등식;
독립 float oracle의 중앙 차분. 중앙 차분은 다음 식을 사용한다.

\frac{\partial L}{\partial z_i} \approx\frac{L(z+h e_i)-L(z-h e_i)}{2h},\qquad h=10^{-6}.

autodiff와 중앙 차분은 $10^{-6}$ 오차 안에서 일치했고 gradient 합은 출력 정밀도에서 0이었다. 극단적인 $10^{20}$ offset에서는 중앙 차분 자체가 반올림에 막히므로 그 경우에는 $\log K$ 항등식을 독립 oracle로 사용했다.

비용과 Blast Radius

실패	가장 작은 차단 위치	늦게 발견했을 때의 영향
empty 또는 non-finite logits	공개 API 진입점	일부 graph 생성 뒤 수치 실패
invalid target	CE 진입점	쓸모없는 partition 전체 생성
mask shape 불일치	softmax 진입점	class 대응 관계 손실
all-masked row	softmax 진입점	0 분모 또는 NaN 정책 누출
graph budget 초과	첫 `Value` 연산 전	CPU·memory를 쓴 뒤 늦은 실패
probability underflow 뒤 log	API 선택 경계	finite loss가 무한대나 예외로 변함
큰 LSE와 target 상쇄	centered 좌표 경계	finite지만 틀린 loss로 학습
orphaned shared core	graph 구조 테스트	비용식과 추상화가 실제 실행과 불일치

이번 Stage도 GPU와 NPU를 사용하지 않는다. 현재 비용은 Python scalar operation과 graph object 수다. 하지만 $O(K)$ 라는 표기만 남기지 않고 정확한 node 수를 계약으로 고정했기 때문에 이후 PyTorch port에서 kernel 수, memory allocation, fused loss와 동일한 의미를 비교할 기준이 생긴다.

직접 실행한 결과

다음 명령을 실행했다.

python .\stages\stage_06_stable_softmax.py
python -m unittest tests.test_stage_06_stable_softmax -v

관찰 결과는 다음과 같다.

$10^{20}$ uniform logits의 centered CE는 $\log 3$ 과 일치했다.
같은 입력의 naive LSE-target은 유한한 0으로 무너졌다.
(0,-1000)의 억제 target 확률은 0이지만 fused CE는 1000이었다.
autodiff와 중앙 차분 gradient가 일치했다.
gradient 합은 0이었다.
budget 실패와 all-masked 실패 뒤 연산 횟수는 모두 0이었다.
Stage 06 테스트 13개가 통과했다.
전체 테스트 70개가 통과했다.
sanity check는 PASS였다.
Stage 06 테스트 13개 통과

이번 Stage가 증명한 것

최초 가설은 맞았다. isfinite()와 확률 합 1은 수치 안정성의 필요조건 일부일뿐 충분조건이 아니다. underflow 뒤의 log와 큰 수끼리의 상쇄는 예외 없이도 학습 의미를 파괴한다. 올바른 추상화는 세 공개 API를 하나로 숨기는 것이 아니다. 호출 목적은 분리하고, overflow,underflow,상쇄가 만나는 centered log-partition만 공유하는 것이다. 학습용 cross-entropy는 probability를 물질화하지 않고 centered logits에서 직접 끝낸다.

Blast Radius도 연산 전에 제한했다. target, mask, finite, 정확한 graph budget을 먼저 검사하고, 실패한 forward와 backward는 기존 gradient를 건드리지 않는다. 마지막으로 실제 graph node를 세어 공유 코어가 코드 이름만이 아니라 실행 그래프에서도 공유된다는 사실까지 확인했다.

다음 Stage.

Stage 07에서는 이 안정적인 확률 경계 앞에 선형 변환을 만든다. 다음에 답할 질문은 다음과 같다.

선형 layer의 weight를 Matrix로 보관하고 출력 shape만 맞추면 충분할까?

initialization scale, bias, parameter identity, graph budget 가운데 무엇을 생성 시점과 forward 시점에 각각 강제해야 여러 layer를 쌓아도 신호와 비용의 Blast Radius가 작게 유지될까?