nn.CrossEntropyLoss() 에 softmax 값을 넣지 않는 이유

Yezun Chung·2023년 10월 4일

Deep Learning

목록 보기

1/16

파이토치에서는 nn.CrossEntropyLoss가 log softmax와 negative log likelihood를 결합하여 구현되어 있습니다. 따라서, 일반적으로 softmax를 거친 결과를 cross entropy에 직접 입력하지 않습니다.

그러나 softmax를 거친 결과를 cross entropy에 입력하려면, nn.NLLLoss를 사용하여 negative log likelihood를 계산할 수 있습니다. 이 경우 softmax를 거친 결과에 log를 취한 후, nn.NLLLoss에 입력해야 합니다.

하지만 이 방법은 불필요한 계산이 추가되기 때문에, 일반적으로는 nn.CrossEntropyLoss를 사용하는 것이 좋습니다. 이 함수는 softmax를 거치지 않은 logit 값을 입력으로 받습니다.

결론적으로, softmax를 거친 결과를 사용하려면 nn.NLLLoss를 사용할 수 있지만, 권장되지 않으며 nn.CrossEntropyLoss를 사용하는 것이 더 효율적입니다.

문제 : Transformer 의 마지막 단에 괜히 softmax layer 를 넣었다가 계속 output 값이 똑같은 값으로 나오는 문제가 발생했습니다.

해결 : 괜히 softmax layer 를 추가하지 말고 (batch size, sequence_length, K-dimension), 즉 K차원의 output dimension 을 그대로 nn.CrossEntropyLoss() 에 넣으면 됩니다~

CS undergraduate student, interested in ML, NLP, CV, RL, multi-modal AI