0 | module NN.Architectures.Softargmax
 1 |
 2 | import Data.Tensor
 3 | import Data.Para
 4 | import Control.Monad.Distribution
 5 |
 6 |
 7 |
 8 | ||| Numerically stable log-sum-exp operation
 9 | ||| LSE(x) = max(x) + log(Σᵢ exp(xᵢ - max(x)))
10 | ||| See https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/
11 | public export
12 | logSumExp : {i : Axis} -> Exp a => Ord a => Neg a =>
13 |   Foldable (Tensor [i]) =>
14 |   (allAlg : AllAlgebra [i] a) =>
15 |   Tensor [i] a -> Maybe a
16 | logSumExp t = do
17 |   c <- max t
18 |   pure $ c + log (reduce (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 |   (allAlg : AllAlgebra [i] a) =>
27 |   Tensor [i] a -> Tensor [i] a
28 | logSoftargmax t = case logSumExp t of
29 |   Just lse => t <&> (\x => x - lse) -- Non-empty: subtract LSE from each element
30 |   Nothing  => t                     -- t is empty
31 |
32 | ||| Commonly known as 'softmax'
33 | ||| When `temperature=0` it reduces to `argmax`
34 | public export
35 | softargmaxImpl : {i : Axis} -> Fractional a => Exp a => Ord a => Neg a =>
36 |   IsFoldable i .cont =>
37 |   (allAlg : AllAlgebra [i] a) =>
38 |   {default 1 temperature : a} ->
39 |   Tensor [i] a -> Tensor [i] a
40 | softargmaxImpl {temperature} t
41 |   = exp <$> logSoftargmax (t <&> (/ temperature))
42 |
43 | ||| Softargmax as a parametric map, with temperature as a parameter
44 | ||| TODO since distribution is an applicative functor (https://glaive-research.org/2025/02/11/Generalized-Transformers-from-Applicative-Functors.html)
45 | ||| is there a meaningful notion of the "distribution container"?
46 | ||| Is there a sense in which `Dist` is a functor on containers?
47 | public export
48 | softargmax : {i : Axis} ->
49 |   {a : Type} -> Fractional a => Exp a => Ord a => Neg a =>
50 |   IsFoldable i.cont =>
51 |   (allAlg : AllAlgebra [i] a) =>
52 |   Tensor [i] a -\-> Tensor [i] a
53 | softargmax = MkPara 
54 |   (\_ => a) -- temperature is the parameter
55 |   (\(t ** temperature) => softargmaxImpl {temperature} t)
56 |
57 |
58 | -- `Control.Monad.Distribution` and softargmax should probably be merged?
59 | -- todo this is missing beause of a show instance for tensors
60 | -- needs an assert total because it goes through tensors
61 | public export
62 | {i : Nat} -> Show (Dist i) where
63 |   show (MkDist xs) = assert_total $ 
64 |     show @{(?todoTensorShow)} (softargmaxImpl {i="softmaxTemp" ~~> i} (># xs))
65 |
66 | inpp : Tensor ["ieva" ~~> 3] Double
67 | inpp = ># [1000, 999, 998]