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.Architectures.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 |   Tensor [i] a -> Maybe a
16 | logSumExp t = do
17 |   c <- max t
18 |   pure $ c + log (reduce @{algebraFoldable} (t <&> (\x => exp $ x - c)))
19 |
20 | ||| Log(softargmax(x)), but computationally efficient and numerically stable
21 | ||| Used for computing cross-entropy loss
22 | ||| Returns empty tensor for empty input
23 | public export
24 | logSoftargmax : {i : Axis} -> Exp a => Ord a => Neg a =>
25 |   Foldable (Tensor [i]) =>
26 |   Tensor [i] a -> Tensor [i] a
27 | logSoftargmax t = case logSumExp t of
28 |   Just lse => t <&> (\x => x - lse) -- Non-empty: subtract LSE from each element
29 |   Nothing  => t                     -- t is empty
30 |
31 | ||| Commonly known as 'softmax'
32 | ||| As `temperature->0` it reduces to `argmax`
33 | ||| TODO fix numerics when temperature=0
34 | public export
35 | softargmaxImpl : {i : Axis} -> Fractional a => Exp a => Ord a => Neg a =>
36 |   IsFoldable i .cont =>
37 |   {default 1 temperature : a} ->
38 |   Tensor [i] a -> Tensor [i] a
39 | softargmaxImpl {temperature} t
40 |   = exp <$> logSoftargmax (t <&> (/ temperature))
41 |