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 | ||| A loss is a parametric map whose parameter is the label
 18 | ||| It is not a `Model` because there is no concept of initialisation
 19 | public export
 20 | Loss : (y, l : AddCont) -> Type
 21 | Loss y l = y =\\=> l
 22 |
 23 | ||| The "parameter" object of the loss function is the type of supervised
 24 | ||| learning labels
 25 | public export
 26 | Label : Loss y l -> Type
 27 | Label loss = (Param loss).Shp
 28 |
 29 | namespace Combinators
 30 |   ||| Run two losses in parallel, add their results
 31 |   public export
 32 |   pairLossFunctions : {y, z : AddCont} -> {l : Type} -> Num l =>
 33 |     Loss y (Const l) -> Loss z (Const l) -> Loss (y >*< z) (Const l)
 34 |   pairLossFunctions f g = postcomposeLens (composeParallel f g) sum
 35 |
 36 |   ||| The loss of a section against a labelled branch
 37 |   ||| We always choose the branch the label chooses
 38 |   public export
 39 |   chosenBranchLoss : {c : Cont} -> {branches : c.filledWith AddCont} ->
 40 |     {0 l : AddCont} ->
 41 |     (losses : (i : c.Pos (shapeExt branches)) -> Loss (index branches i) l) ->
 42 |     Loss (Section (index branches)) l
 43 |   chosenBranchLoss losses = MkPara (AddContDPair (\i => Param (losses i)))
 44 |     (evalSection {a = index branches} {p = \i => Param (losses i)}
 45 |       %+>> copair (\i => Run (losses i)))
 46 |
 47 |   ||| Same as above, except we take an already chosen branch, and fail if 
 48 |   ||| label disagrees with it
 49 |   public export
 50 |   matchedBranchLoss : {c : Cont} -> (dec : IsDecidable c) =>
 51 |     {branches : c.filledWith AddCont} -> {0 l : AddCont} ->
 52 |     (losses : (i : c.Pos (shapeExt branches)) -> Loss (index branches i) l) ->
 53 |     Loss (Coproduct branches) (Maybe l)
 54 |   matchedBranchLoss losses = MkPara (AddContDPair (\i => Param (losses i)))
 55 |     (matchIndex @{GetInterface dec (shapeExt branches)}
 56 |       {b = index branches} {l = \i => Param (losses i)}
 57 |       %+>> Maybe (copair (\i => Run (losses i))))
 58 |
 59 |
 60 |   ||| Given a made choice, pair with with a dirac delta distribution
 61 |   public export
 62 |   madeChoice : {ax : Axis} -> IsNaperian ax => DecEq (Log ax) =>
 63 |     {f : Log ax -> AddCont} ->
 64 |     AddContDPair f =%+> Simplex ax >*< AddContDPair f
 65 |   madeChoice = !%+ \ix => ((diracDelta (fst ix), ix) ** snd)
 66 |
 67 | namespace Instances
 68 |   public export
 69 |   SquaredError : {a : Type} -> Num a => Neg a => Loss (Const a) (Const a)
 70 |   SquaredError = MkPara (Const a) SquaredDifference
 71 |
 72 |   public export
 73 |   MeanSquaredError : {n : Axis} -> IsCubical n => TensorMonoid n.cont =>
 74 |     {a : Type} -> Num a => Neg a => Fractional a => Cast Nat a =>
 75 |     Loss (Const (Tensor [n] a)) (Const (Tensor [] a))
 76 |   MeanSquaredError = MkPara (Const (Tensor [n] a)) meanSquaredDifference
 77 |
 78 |   ||| The payoff object is the rank-0 tensor, not `Double`
 79 |   ||| Defined for Naperian axes, where the prediction and the label have one shape
 80 |   ||| TODO the prediction and the label can be distributions of different shapes!
 81 |   public export
 82 |   softargmaxCrossEntropyLogits : {a : Axis} ->
 83 |     (nap : IsNaperian a) => IsFoldable a.cont =>
 84 |     Simplex a >*< Simplex a =%+> Const (Tensor [] Double)
 85 |   softargmaxCrossEntropyLogits @{MkIsNaperian _ pos} = !%+ \(predicted, labels) =>
 86 |     let logSoftargmaxLogits = logSoftargmax predicted.logits
 87 |         targetProbs = softargmaxImpl labels.logits
 88 |         out = - dot logSoftargmaxLogits targetProbs
 89 |     in (out ** \l' =>
 90 |       ((extract l' *) <$> (Prelude.exp <$> logSoftargmaxLogits) - targetProbs,
 91 |        -- the derivative in the label's logits, though training discards it
 92 |        (extract l' *) <$> negate (targetProbs * ((+ extract out) <$> logSoftargmaxLogits))))
 93 |
 94 |   public export
 95 |   SoftargmaxCrossEntropyLogits : {a : Axis} ->
 96 |     IsNaperian a => IsFoldable a.cont =>
 97 |     Loss (Simplex a) (Const (Tensor [] Double))
 98 |   SoftargmaxCrossEntropyLogits
 99 |     = MkPara (Simplex a) softargmaxCrossEntropyLogits
100 |