0 | module NN.Architectures.LossFunctions
 1 |
 2 | import Data.List
 3 | import Data.Fin
 4 | import Data.Vect
 5 | import Data.Zippable
 6 |
 7 | import Data.Tensor
 8 | import Data.Tensor.Utils
 9 | import Data.Container.Additive
10 | import Data.Autodiff.Ops
11 | import Control.Monad.Distribution
12 |
13 | import Data.Container.Additive.Quantifiers
14 |
15 | import Data.Para
16 |
17 | %hide Data.Container.Base.Morphism.Definition.DependentLenses.(=%>)
18 |
19 | ||| A loss is a parametric map whose parameter is the label
20 | ||| It is not a `Model` because there is no concept of initialisation
21 | public export
22 | Loss : (y, l : AddCont) -> Type
23 | Loss y l = y =\\=> l
24 |
25 | namespace Combinators
26 |   ||| Run two losses in parallel, add their results
27 |   public export
28 |   pairLossFunctions : {y, z : AddCont} -> {l : Type} -> Num l =>
29 |     Loss y (Const l) -> Loss z (Const l) -> Loss (y >*< z) (Const l)
30 |   pairLossFunctions f g = postcomposeLens (composeParallel f g) sum
31 |
32 |   ||| The loss of a coproduct of choices. When the types don't match, gradient
33 |   ||| is infinite. In our examples we don't expect this to happen; but loss type
34 |   ||| should be refined to exclude it eventually
35 |   public export
36 |   chosenBranchLoss : {n : Nat} -> {branches : Vect n AddCont} ->
37 |     {default branches labels : Vect n AddCont} ->
38 |     {0 lc : AddCont} -> Fractional lc.Shp =>
39 |     (losses : (i : Fin n) -> index i branches >*< index i labels =%+> lc) ->
40 |     Loss (Coproduct branches) lc
41 |   chosenBranchLoss losses = MkPara (Coproduct labels) $
42 |     !%+ \((i ** x), (j ** y)) => case decEq i j of
43 |       Yes Refl => (%!+) (losses i) (x, y)
44 |       No _ => (1 / 0 ** \_ => ((index i branches).Zero x, (index j labels).Zero y))
45 |
46 |   ||| The loss variant of `resolveByLabel`. Given a loss on resolved choices we
47 |   ||| can produce a loss on an effectful output, where the ground-truth label
48 |   ||| selects the effect
49 |   ||| This means that the training loop does not need to handle any effects 
50 |   ||| anymore
51 |   public export
52 |   resolveLoss : {distName : AxisName} -> {n : Nat} ->
53 |     {branches : Vect n AddCont} ->
54 |     {0 l : AddCont} ->
55 |     (ChoiceMade distName branches >*< ChoiceMade distName branches =%+> l) ->
56 |     Loss (ProbabilisticChoice distName branches) l
57 |   resolveLoss loss = MkPara
58 |     (ChoiceMade distName branches)
59 |     (resolveByLabel %+>> loss)
60 |
61 | namespace Instances
62 |   public export
63 |   SquaredError : {a : Type} -> Num a => Neg a => Loss (Const a) (Const a)
64 |   SquaredError = MkPara (Const a) SquaredDifference
65 |
66 |   public export
67 |   MeanSquaredError : {n : Axis} -> IsCubical n => TensorMonoid n.cont =>
68 |     {a : Type} -> Num a => Neg a => Fractional a => Cast Nat a =>
69 |     Loss (Const (Tensor [n] a)) (Const (Tensor [] a))
70 |   MeanSquaredError = MkPara (Const (Tensor [n] a)) meanSquaredDifference
71 |
72 |   ||| The payoff object is the rank-0 tensor, not `Double`
73 |   public export
74 |   softargmaxCrossEntropyLogits : {name : AxisName} -> {n : Nat} ->
75 |     Simplex name n >*< Simplex name n =%+> Const (Tensor [] Double)
76 |   softargmaxCrossEntropyLogits = !%+ \(predicted, labels) =>
77 |     let logSoftargmaxLogits = logSoftargmax predicted.logits
78 |         targetProbs = softargmaxImpl labels.logits
79 |         out = - dot logSoftargmaxLogits targetProbs
80 |     in (out ** \l' =>
81 |       ((extract l' *) <$> (Prelude.exp <$> logSoftargmaxLogits) - targetProbs,
82 |         fill 0)) -- zeros for now
83 |
84 |   public export
85 |   SoftargmaxCrossEntropyLogits : {name : AxisName} -> {n : Nat} ->
86 |     Loss (Simplex name n) (Const (Tensor [] Double))
87 |   SoftargmaxCrossEntropyLogits
88 |     = MkPara (Simplex name n) softargmaxCrossEntropyLogits
89 |