0 | module Data.Tensor.Softargmax
 1 |
 2 | import Data.Tensor.Tensor
 3 | import Data.Tensor.Utils
 4 |
 5 | {-------------------------------------------------------------------------------
 6 | Used by both `Control.Monad.Distribution` and `NN.Architetures.Softargmax`
 7 | -------------------------------------------------------------------------------}
 8 |
 9 | ||| Numerically stable log-sum-exp operation
10 | ||| LSE(x) = max(x) + log(Σᵢ exp(xᵢ - max(x)))
11 | ||| See https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/
12 | public export
13 | logSumExp : {i : Axis} -> Exp a => Ord a => Neg a =>
14 |   Foldable (Tensor [i]) =>
15 |   (allAlg : AllAlgebra [i] a) =>
16 |   Tensor [i] a -> Maybe a
17 | logSumExp t = do
18 |   c <- max t
19 |   pure $ c + log (reduce (t <&> (\x => exp $ x - c)))
20 |
21 | ||| Log(softargmax(x)), but computationally efficient and numerically stable
22 | ||| Used for computing cross-entropy loss
23 | ||| Returns empty tensor for empty input
24 | public export
25 | logSoftargmax : {i : Axis} -> Exp a => Ord a => Neg a =>
26 |   Foldable (Tensor [i]) =>
27 |   (allAlg : AllAlgebra [i] a) =>
28 |   Tensor [i] a -> Tensor [i] a
29 | logSoftargmax t = case logSumExp t of
30 |   Just lse => t <&> (\x => x - lse) -- Non-empty: subtract LSE from each element
31 |   Nothing  => t                     -- t is empty
32 |
33 | ||| Commonly known as 'softmax'
34 | ||| When `temperature=0` it reduces to `argmax`
35 | public export
36 | softargmaxImpl : {i : Axis} -> Fractional a => Exp a => Ord a => Neg a =>
37 |   IsFoldable i .cont =>
38 |   (allAlg : AllAlgebra [i] a) =>
39 |   {default 1 temperature : a} ->
40 |   Tensor [i] a -> Tensor [i] a
41 | softargmaxImpl {temperature} t
42 |   = exp <$> logSoftargmax (t <&> (/ temperature))
43 |