{-# LANGUAGE OverloadedStrings #-}

-- | Full backward pass through one transformer block.
--   This is the core of the training backprop chain.
module Circuit.LLM.Backprop
  ( -- * Full backward pass
    gptBackward,
    bertBackward,

    -- * Gradient types
    GptGrads (..),
    BlockGrads (..),
    zeroGptGrads,
    addGptGrads,
    scaleGptGrads,

    -- * Primitives (exported for testing)
    linearBwd,
    layerNormBwd,
    geluBwd,
    softmaxBwd,
    crossEntropyBwd,
    maskedCrossEntropyBwd,
  )
where

import Circuit.LLM.Diff
  ( BlockParams (..),
    GptParams (..),
    bertDiffP,
    gptDiffP,
    gptParamsFromModel,
    primBackward,
    primForward,
    subMatrixW,
    zeroMatrix,
  )
import Circuit.LLM.GPT (Gpt (..), GptConfig (..))
import Numeric.LinearAlgebra
  ( Matrix,
    Vector,
    cmap,
    cols,
    fromList,
    fromRows,
    maxElement,
    reshape,
    rows,
    scale,
    sumElements,
    toList,
    toRows,
    tr,
  )
import Numeric.LinearAlgebra qualified as LA

----------------------------------------------------------------------
-- Gradient accumulators (unchanged)
----------------------------------------------------------------------

data GptGrads = GptGrads
  { GptGrads -> Matrix Double
ggWte :: !(Matrix Double),
    GptGrads -> Matrix Double
ggWpe :: !(Matrix Double),
    GptGrads -> [BlockGrads]
ggBlocks :: ![BlockGrads],
    GptGrads -> Vector Double
ggLnGamma :: !(Vector Double),
    GptGrads -> Vector Double
ggLnBeta :: !(Vector Double),
    GptGrads -> Matrix Double
ggHead :: !(Matrix Double),
    GptGrads -> Vector Double
ggHeadB :: !(Vector Double)
  }

data BlockGrads = BlockGrads
  { BlockGrads -> Matrix Double
bgAttnWq, BlockGrads -> Matrix Double
bgAttnWk, BlockGrads -> Matrix Double
bgAttnWv, BlockGrads -> Matrix Double
bgAttnWo :: !(Matrix Double),
    BlockGrads -> Vector Double
bgAttnLnGamma, BlockGrads -> Vector Double
bgAttnLnBeta :: !(Vector Double),
    BlockGrads -> Matrix Double
bgFfnW1, BlockGrads -> Matrix Double
bgFfnW2 :: !(Matrix Double),
    BlockGrads -> Vector Double
bgFfnB1, BlockGrads -> Vector Double
bgFfnB2 :: !(Vector Double),
    BlockGrads -> Vector Double
bgFfnLnGamma, BlockGrads -> Vector Double
bgFfnLnBeta :: !(Vector Double)
  }

zeroGptGrads :: GptConfig -> GptGrads
zeroGptGrads :: GptConfig -> GptGrads
zeroGptGrads GptConfig
cfg =
  let nEmb :: Int
nEmb = GptConfig -> Int
gptNEmbd GptConfig
cfg
      nLayer :: Int
nLayer = GptConfig -> Int
gptNLayer GptConfig
cfg
      vocab :: Int
vocab = GptConfig -> Int
gptVocabSize GptConfig
cfg
      ffMul :: Int
ffMul = Int
4
      zM :: Int -> Int -> Matrix t
zM Int
r Int
c = Int -> Vector t -> Matrix t
forall t. Storable t => Int -> Vector t -> Matrix t
reshape Int
c ([t] -> Vector t
forall a. Storable a => [a] -> Vector a
fromList (Int -> t -> [t]
forall a. Int -> a -> [a]
replicate (Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c) t
0))
      zV :: Int -> Vector a
zV Int
n = [a] -> Vector a
forall a. Storable a => [a] -> Vector a
fromList (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n a
0)
   in GptGrads
        { ggWte :: Matrix Double
ggWte = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
vocab Int
nEmb,
          ggWpe :: Matrix Double
ggWpe = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
1024 Int
nEmb,
          ggBlocks :: [BlockGrads]
ggBlocks =
            Int -> BlockGrads -> [BlockGrads]
forall a. Int -> a -> [a]
replicate Int
nLayer (BlockGrads -> [BlockGrads]) -> BlockGrads -> [BlockGrads]
forall a b. (a -> b) -> a -> b
$
              BlockGrads
                { bgAttnWq :: Matrix Double
bgAttnWq = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb Int
nEmb,
                  bgAttnWk :: Matrix Double
bgAttnWk = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb Int
nEmb,
                  bgAttnWv :: Matrix Double
bgAttnWv = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb Int
nEmb,
                  bgAttnWo :: Matrix Double
bgAttnWo = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb Int
nEmb,
                  bgAttnLnGamma :: Vector Double
bgAttnLnGamma = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
                  bgAttnLnBeta :: Vector Double
bgAttnLnBeta = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
                  bgFfnW1 :: Matrix Double
bgFfnW1 = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb (Int
ffMul Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nEmb),
                  bgFfnB1 :: Vector Double
bgFfnB1 = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV (Int
ffMul Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nEmb),
                  bgFfnW2 :: Matrix Double
bgFfnW2 = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM (Int
ffMul Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nEmb) Int
nEmb,
                  bgFfnB2 :: Vector Double
bgFfnB2 = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
                  bgFfnLnGamma :: Vector Double
bgFfnLnGamma = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
                  bgFfnLnBeta :: Vector Double
bgFfnLnBeta = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb
                },
          ggLnGamma :: Vector Double
ggLnGamma = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
          ggLnBeta :: Vector Double
ggLnBeta = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
nEmb,
          ggHead :: Matrix Double
ggHead = Int -> Int -> Matrix Double
forall {t}. (Storable t, Num t) => Int -> Int -> Matrix t
zM Int
nEmb Int
vocab,
          ggHeadB :: Vector Double
ggHeadB = Int -> Vector Double
forall {a}. (Storable a, Num a) => Int -> Vector a
zV Int
vocab
        }

addGptGrads :: GptGrads -> GptGrads -> GptGrads
addGptGrads :: GptGrads -> GptGrads -> GptGrads
addGptGrads GptGrads
a GptGrads
b =
  let r1 :: Int
r1 = Matrix Double -> Int
forall t. Matrix t -> Int
rows (GptGrads -> Matrix Double
ggWte GptGrads
a)
      c1 :: Int
c1 = Matrix Double -> Int
forall t. Matrix t -> Int
cols (GptGrads -> Matrix Double
ggWte GptGrads
a)
      r2 :: Int
r2 = Matrix Double -> Int
forall t. Matrix t -> Int
rows (GptGrads -> Matrix Double
ggWte GptGrads
b)
      c2 :: Int
c2 = Matrix Double -> Int
forall t. Matrix t -> Int
cols (GptGrads -> Matrix Double
ggWte GptGrads
b)
   in if Int
r1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r2 Bool -> Bool -> Bool
|| Int
c1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c2
        then [Char] -> GptGrads
forall a. HasCallStack => [Char] -> a
error ([Char] -> GptGrads) -> [Char] -> GptGrads
forall a b. (a -> b) -> a -> b
$ [Char]
"addGptGrads Wte mismatch: (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
r1 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
c1 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") vs (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
r2 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
c2 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
        else
          GptGrads
            { ggWte :: Matrix Double
ggWte = GptGrads -> Matrix Double
ggWte GptGrads
a Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ GptGrads -> Matrix Double
ggWte GptGrads
b,
              ggWpe :: Matrix Double
ggWpe = GptGrads -> Matrix Double
ggWpe GptGrads
a Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ GptGrads -> Matrix Double
ggWpe GptGrads
b,
              ggBlocks :: [BlockGrads]
ggBlocks = (BlockGrads -> BlockGrads -> BlockGrads)
-> [BlockGrads] -> [BlockGrads] -> [BlockGrads]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith BlockGrads -> BlockGrads -> BlockGrads
addBG (GptGrads -> [BlockGrads]
ggBlocks GptGrads
a) (GptGrads -> [BlockGrads]
ggBlocks GptGrads
b),
              ggLnGamma :: Vector Double
ggLnGamma = GptGrads -> Vector Double
ggLnGamma GptGrads
a Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ GptGrads -> Vector Double
ggLnGamma GptGrads
b,
              ggLnBeta :: Vector Double
ggLnBeta = GptGrads -> Vector Double
ggLnBeta GptGrads
a Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ GptGrads -> Vector Double
ggLnBeta GptGrads
b,
              ggHead :: Matrix Double
ggHead = GptGrads -> Matrix Double
ggHead GptGrads
a Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ GptGrads -> Matrix Double
ggHead GptGrads
b,
              ggHeadB :: Vector Double
ggHeadB = GptGrads -> Vector Double
ggHeadB GptGrads
a Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ GptGrads -> Vector Double
ggHeadB GptGrads
b
            }
  where
    addBG :: BlockGrads -> BlockGrads -> BlockGrads
addBG BlockGrads
x BlockGrads
y =
      let r1 :: Int
r1 = Matrix Double -> Int
forall t. Matrix t -> Int
rows (BlockGrads -> Matrix Double
bgAttnWq BlockGrads
x)
          c1 :: Int
c1 = Matrix Double -> Int
forall t. Matrix t -> Int
cols (BlockGrads -> Matrix Double
bgAttnWq BlockGrads
x)
          r2 :: Int
r2 = Matrix Double -> Int
forall t. Matrix t -> Int
rows (BlockGrads -> Matrix Double
bgAttnWq BlockGrads
y)
          c2 :: Int
c2 = Matrix Double -> Int
forall t. Matrix t -> Int
cols (BlockGrads -> Matrix Double
bgAttnWq BlockGrads
y)
       in if Int
r1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
r2 Bool -> Bool -> Bool
|| Int
c1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
c2
            then [Char] -> BlockGrads
forall a. HasCallStack => [Char] -> a
error ([Char] -> BlockGrads) -> [Char] -> BlockGrads
forall a b. (a -> b) -> a -> b
$ [Char]
"addBG Wq mismatch: (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
r1 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
c1 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
") vs (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
r2 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"," [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
c2 [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"
            else
              BlockGrads
                { bgAttnWq :: Matrix Double
bgAttnWq = BlockGrads -> Matrix Double
bgAttnWq BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgAttnWq BlockGrads
y,
                  bgAttnWk :: Matrix Double
bgAttnWk = BlockGrads -> Matrix Double
bgAttnWk BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgAttnWk BlockGrads
y,
                  bgAttnWv :: Matrix Double
bgAttnWv = BlockGrads -> Matrix Double
bgAttnWv BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgAttnWv BlockGrads
y,
                  bgAttnWo :: Matrix Double
bgAttnWo = BlockGrads -> Matrix Double
bgAttnWo BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgAttnWo BlockGrads
y,
                  bgAttnLnGamma :: Vector Double
bgAttnLnGamma = BlockGrads -> Vector Double
bgAttnLnGamma BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgAttnLnGamma BlockGrads
y,
                  bgAttnLnBeta :: Vector Double
bgAttnLnBeta = BlockGrads -> Vector Double
bgAttnLnBeta BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgAttnLnBeta BlockGrads
y,
                  bgFfnW1 :: Matrix Double
bgFfnW1 = BlockGrads -> Matrix Double
bgFfnW1 BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgFfnW1 BlockGrads
y,
                  bgFfnB1 :: Vector Double
bgFfnB1 = BlockGrads -> Vector Double
bgFfnB1 BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgFfnB1 BlockGrads
y,
                  bgFfnW2 :: Matrix Double
bgFfnW2 = BlockGrads -> Matrix Double
bgFfnW2 BlockGrads
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Matrix Double
bgFfnW2 BlockGrads
y,
                  bgFfnB2 :: Vector Double
bgFfnB2 = BlockGrads -> Vector Double
bgFfnB2 BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgFfnB2 BlockGrads
y,
                  bgFfnLnGamma :: Vector Double
bgFfnLnGamma = BlockGrads -> Vector Double
bgFfnLnGamma BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgFfnLnGamma BlockGrads
y,
                  bgFfnLnBeta :: Vector Double
bgFfnLnBeta = BlockGrads -> Vector Double
bgFfnLnBeta BlockGrads
x Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ BlockGrads -> Vector Double
bgFfnLnBeta BlockGrads
y
                }

scaleGptGrads :: Double -> GptGrads -> GptGrads
scaleGptGrads :: Double -> GptGrads -> GptGrads
scaleGptGrads Double
s GptGrads
g =
  GptGrads
    { ggWte :: Matrix Double
ggWte = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Matrix Double
ggWte GptGrads
g),
      ggWpe :: Matrix Double
ggWpe = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Matrix Double
ggWpe GptGrads
g),
      ggBlocks :: [BlockGrads]
ggBlocks = (BlockGrads -> BlockGrads) -> [BlockGrads] -> [BlockGrads]
forall a b. (a -> b) -> [a] -> [b]
map (Double -> BlockGrads -> BlockGrads
scaleBG Double
s) (GptGrads -> [BlockGrads]
ggBlocks GptGrads
g),
      ggLnGamma :: Vector Double
ggLnGamma = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Vector Double
ggLnGamma GptGrads
g),
      ggLnBeta :: Vector Double
ggLnBeta = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Vector Double
ggLnBeta GptGrads
g),
      ggHead :: Matrix Double
ggHead = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Matrix Double
ggHead GptGrads
g),
      ggHeadB :: Vector Double
ggHeadB = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s (GptGrads -> Vector Double
ggHeadB GptGrads
g)
    }
  where
    scaleBG :: Double -> BlockGrads -> BlockGrads
scaleBG Double
s_ BlockGrads
b =
      BlockGrads
        { bgAttnWq :: Matrix Double
bgAttnWq = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgAttnWq BlockGrads
b),
          bgAttnWk :: Matrix Double
bgAttnWk = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgAttnWk BlockGrads
b),
          bgAttnWv :: Matrix Double
bgAttnWv = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgAttnWv BlockGrads
b),
          bgAttnWo :: Matrix Double
bgAttnWo = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgAttnWo BlockGrads
b),
          bgAttnLnGamma :: Vector Double
bgAttnLnGamma = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgAttnLnGamma BlockGrads
b),
          bgAttnLnBeta :: Vector Double
bgAttnLnBeta = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgAttnLnBeta BlockGrads
b),
          bgFfnW1 :: Matrix Double
bgFfnW1 = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgFfnW1 BlockGrads
b),
          bgFfnB1 :: Vector Double
bgFfnB1 = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgFfnB1 BlockGrads
b),
          bgFfnW2 :: Matrix Double
bgFfnW2 = Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Matrix Double
bgFfnW2 BlockGrads
b),
          bgFfnB2 :: Vector Double
bgFfnB2 = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgFfnB2 BlockGrads
b),
          bgFfnLnGamma :: Vector Double
bgFfnLnGamma = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgFfnLnGamma BlockGrads
b),
          bgFfnLnBeta :: Vector Double
bgFfnLnBeta = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
s_ (BlockGrads -> Vector Double
bgFfnLnBeta BlockGrads
b)
        }

----------------------------------------------------------------------
-- Backward primitives (unchanged)
----------------------------------------------------------------------

linearBwd :: Matrix Double -> Matrix Double -> Matrix Double -> (Matrix Double, Matrix Double, Vector Double)
linearBwd :: Matrix Double
-> Matrix Double
-> Matrix Double
-> (Matrix Double, Matrix Double, Vector Double)
linearBwd Matrix Double
x Matrix Double
w Matrix Double
gradY =
  let gradX :: Matrix Double
gradX = Matrix Double
gradY Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
w
      gradW :: Matrix Double
gradW = Matrix Double -> Matrix Double
forall m mt. Transposable m mt => m -> mt
tr Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
LA.<> Matrix Double
gradY
      gradB :: Vector Double
gradB = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
col | Vector Double
col <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
gradY]
   in (Matrix Double
gradX, Matrix Double
gradW, Vector Double
gradB)

geluBwd :: Matrix Double -> Matrix Double -> Matrix Double
geluBwd :: Matrix Double -> Matrix Double -> Matrix Double
geluBwd Matrix Double
x Matrix Double
gradY =
  let xs :: [Double]
xs = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
x)
      gs :: [Double]
gs = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
gradY)
      deriv :: a -> a
deriv a
v =
        let a :: a
a = a
1.59577
            b :: a
b = a
0.044715
            z :: a
z = a
a a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v)
            phi :: a
phi = a
1 a -> a -> a
forall a. Fractional a => a -> a -> a
/ (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a -> a
forall a. Floating a => a -> a
exp (-a
z))
            phi' :: a
phi' = a
phi a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
- a
phi) a -> a -> a
forall a. Num a => a -> a -> a
* a
a a -> a -> a
forall a. Num a => a -> a -> a
* (a
1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
3 a -> a -> a
forall a. Num a => a -> a -> a
* a
b a -> a -> a
forall a. Num a => a -> a -> a
* a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
v)
         in a
phi a -> a -> a
forall a. Num a => a -> a -> a
+ a
v a -> a -> a
forall a. Num a => a -> a -> a
* a
phi'
   in Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x) ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList ((Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) [Double]
gs ((Double -> Double) -> [Double] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map Double -> Double
forall a. Floating a => a -> a
deriv [Double]
xs)))

softmaxBwd :: Matrix Double -> Matrix Double -> Matrix Double
softmaxBwd :: Matrix Double -> Matrix Double -> Matrix Double
softmaxBwd Matrix Double
probs Matrix Double
gradOut =
  [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector Double] -> Matrix Double)
-> [Vector Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Vector Double -> Vector Double -> Vector Double
forall {a}. (Storable a, Num a) => Vector a -> Vector a -> Vector a
softmaxRowBwd (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
probs) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradOut)
  where
    softmaxRowBwd :: Vector a -> Vector a -> Vector a
softmaxRowBwd Vector a
p Vector a
go =
      let pv :: [a]
pv = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
p; gov :: [a]
gov = Vector a -> [a]
forall a. Storable a => Vector a -> [a]
toList Vector a
go; dot :: a
dot = [a] -> a
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Num a => a -> a -> a
(*) [a]
pv [a]
gov)
       in [a] -> Vector a
forall a. Storable a => [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\a
pi_ a
goi -> a
pi_ a -> a -> a
forall a. Num a => a -> a -> a
* (a
goi a -> a -> a
forall a. Num a => a -> a -> a
- a
dot)) [a]
pv [a]
gov

layerNormBwd ::
  Matrix Double ->
  Vector Double ->
  Vector Double ->
  Double ->
  Matrix Double ->
  (Matrix Double, Vector Double, Vector Double)
layerNormBwd :: Matrix Double
-> Vector Double
-> Vector Double
-> Double
-> Matrix Double
-> (Matrix Double, Vector Double, Vector Double)
layerNormBwd Matrix Double
x Vector Double
gamma Vector Double
_beta Double
eps Matrix Double
gradY =
  let mu :: Vector Double
mu = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
LA.fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
row Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
x]
      xm :: Matrix Double
xm = Matrix Double
x Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
- Vector Double -> Matrix Double
forall a. Storable a => Vector a -> Matrix a
LA.asColumn Vector Double
mu
      var :: Vector Double
var = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
LA.fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
row Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
row) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d | Vector Double
row <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm]
      invStd :: Vector Double
invStd = (Double -> Double) -> Vector Double -> Vector Double
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\Double
v -> Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double -> Double
forall a. Floating a => a -> a
sqrt (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps)) Vector Double
var
      xHat :: Matrix Double
xHat = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows ([Vector Double] -> Matrix Double)
-> [Vector Double] -> Matrix Double
forall a b. (a -> b) -> a -> b
$ (Vector Double -> Double -> Vector Double)
-> [Vector Double] -> [Double] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((Double -> Vector Double -> Vector Double)
-> Vector Double -> Double -> Vector Double
forall a b c. (a -> b -> c) -> b -> a -> c
flip Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
invStd)
      gammaList :: [Double]
gammaList = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
gamma
      gradXHat :: Matrix Double
gradXHat = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [[Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
LA.fromList ([Double] -> Vector Double) -> [Double] -> Vector Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double) -> [Double] -> [Double] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Double -> Double -> Double
forall a. Num a => a -> a -> a
(*) (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
go) [Double]
gammaList | Vector Double
go <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradY]
      gGamma :: Vector Double
gGamma = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
go Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
xh) | (Vector Double
go, Vector Double
xh) <- [Vector Double]
-> [Vector Double] -> [(Vector Double, Vector Double)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
gradY) (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
xHat)]
      gBeta :: Vector Double
gBeta = [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
go | Vector Double
go <- Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
LA.toColumns Matrix Double
gradY]
      gradX :: Matrix Double
gradX =
        [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows
          [ Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRow
              (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradXHat [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
xm [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
invStd [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
              (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
var [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
          | Int
i <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
          ]
   in (Matrix Double
gradX, Vector Double
gGamma, Vector Double
gBeta)
  where
    d :: Double
d = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x)
    features :: Int
features = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
x
    gradRow :: Vector Double -> Vector Double -> Double -> Double -> Vector Double
gradRow Vector Double
gxh Vector Double
xi_minus_mu Double
is Double
v =
      let gInvStd :: Double
gInvStd = Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements (Vector Double
gxh Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
* Vector Double
xi_minus_mu)
          gVar :: Double
gVar = Double
gInvStd Double -> Double -> Double
forall a. Num a => a -> a -> a
* (-Double
0.5) Double -> Double -> Double
forall a. Num a => a -> a -> a
* (Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
eps) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** (-Double
1.5)
          gXm :: Vector Double
gXm = Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale Double
is Vector Double
gxh Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ Double -> Vector Double -> Vector Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
gVar Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d) Vector Double
xi_minus_mu
          gMu :: Double
gMu = -Vector Double -> Double
forall (c :: * -> *) e. Container c e => c e -> e
sumElements Vector Double
gXm
       in Vector Double
gXm Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
features (Double
gMu Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
d))

crossEntropyBwd :: Matrix Double -> [Int] -> (Double, Matrix Double)
crossEntropyBwd :: Matrix Double -> [Int] -> (Double, Matrix Double)
crossEntropyBwd Matrix Double
logits [Int]
targetIds =
  let n :: Double
n = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
logits)
      softmaxRows :: [Vector Double]
softmaxRows = (Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
map Vector Double -> Vector Double
forall {c :: * -> *} {t}.
(Container c t, Floating t, Linear t c) =>
c t -> c t
softmaxStable (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
logits)
      losses :: [Double]
losses = (Vector Double -> Int -> Double)
-> [Vector Double] -> [Int] -> [Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Vector Double
p Int
t -> -Double -> Double
forall a. Floating a => a -> a
log (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
p [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
t) Double
1e-12)) [Vector Double]
softmaxRows [Int]
targetIds
      totalLoss :: Double
totalLoss = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
losses Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n
      gradRows :: [Vector Double]
gradRows =
        (Vector Double -> Int -> Vector Double)
-> [Vector Double] -> [Int] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          ( \Vector Double
p Int
t ->
              let pv :: [Double]
pv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
p
               in [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t then [Double]
pv [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1 else [Double]
pv [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
pv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
          )
          [Vector Double]
softmaxRows
          [Int]
targetIds
   in (Double
totalLoss, Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
n) ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
gradRows))
  where
    softmaxStable :: c t -> c t
softmaxStable c t
v =
      let mx :: t
mx = c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
maxElement c t
v; shifted :: c t
shifted = (t -> t) -> c t -> c t
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\t
x -> t -> t
forall a. Floating a => a -> a
exp (t
x t -> t -> t
forall a. Num a => a -> a -> a
- t
mx)) c t
v
       in t -> c t -> c t
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (t
1 t -> t -> t
forall a. Fractional a => a -> a -> a
/ c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
sumElements c t
shifted) c t
shifted

-- | Masked cross-entropy for BERT-style training.
--   The mask indicates which positions are supervised; loss and gradients are
--   averaged only over those positions.
maskedCrossEntropyBwd :: Matrix Double -> [Int] -> [Bool] -> (Double, Matrix Double)
maskedCrossEntropyBwd :: Matrix Double -> [Int] -> [Bool] -> (Double, Matrix Double)
maskedCrossEntropyBwd Matrix Double
logits [Int]
targetIds [Bool]
mask =
  let maskedCount :: Double
maskedCount = Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Bool] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
forall a. a -> a
id [Bool]
mask))
      vocab :: Int
vocab = Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
logits
      softmaxRows :: [Vector Double]
softmaxRows = (Vector Double -> Vector Double)
-> [Vector Double] -> [Vector Double]
forall a b. (a -> b) -> [a] -> [b]
map Vector Double -> Vector Double
forall {c :: * -> *} {t}.
(Container c t, Floating t, Linear t c) =>
c t -> c t
softmaxStable (Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
logits)
      maskedLosses :: [Double]
maskedLosses =
        (Vector Double -> Int -> Bool -> Double)
-> [Vector Double] -> [Int] -> [Bool] -> [Double]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
          ( \Vector Double
p Int
t Bool
m ->
              if Bool
m then -Double -> Double
forall a. Floating a => a -> a
log (Double -> Double -> Double
forall a. Ord a => a -> a -> a
max (Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
p [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
t) Double
1e-12) else Double
0
          )
          [Vector Double]
softmaxRows
          [Int]
targetIds
          [Bool]
mask
      totalLoss :: Double
totalLoss = if Double
maskedCount Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0 then Double
0 else [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Double]
maskedLosses Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
maskedCount
      gradRows :: [Vector Double]
gradRows =
        (Vector Double -> (Int, Bool) -> Vector Double)
-> [Vector Double] -> [(Int, Bool)] -> [Vector Double]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
          ( \Vector Double
p (Int
t, Bool
m) ->
              if Bool
m
                then
                  let pv :: [Double]
pv = Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
p
                   in [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
t then [Double]
pv [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
1 else [Double]
pv [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- [Int
0 .. [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
pv Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
                else [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList (Int -> Double -> [Double]
forall a. Int -> a -> [a]
replicate Int
vocab Double
0)
          )
          [Vector Double]
softmaxRows
          ([Int] -> [Bool] -> [(Int, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
targetIds [Bool]
mask)
   in if Double
maskedCount Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
0
        then (Double
0, Int -> Int -> Matrix Double
zeroMatrix (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
logits) (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
logits))
        else (Double
totalLoss, Double -> Matrix Double -> Matrix Double
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (Double
1 Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
maskedCount) ([Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Vector Double]
gradRows))
  where
    softmaxStable :: c t -> c t
softmaxStable c t
v =
      let mx :: t
mx = c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
maxElement c t
v; shifted :: c t
shifted = (t -> t) -> c t -> c t
forall b (c :: * -> *) e.
(Element b, Container c e) =>
(e -> b) -> c e -> c b
cmap (\t
x -> t -> t
forall a. Floating a => a -> a
exp (t
x t -> t -> t
forall a. Num a => a -> a -> a
- t
mx)) c t
v
       in t -> c t -> c t
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (t
1 t -> t -> t
forall a. Fractional a => a -> a -> a
/ c t -> t
forall (c :: * -> *) e. Container c e => c e -> e
sumElements c t
shifted) c t
shifted

-- | Convert a flat 'BlockParams' (used as the parameter/gradient carrier in
--   'DiffP') back to the existing 'BlockGrads' record.
blockGradsFromParams :: BlockParams -> BlockGrads
blockGradsFromParams :: BlockParams -> BlockGrads
blockGradsFromParams BlockParams
bp =
  BlockGrads
    { bgAttnWq :: Matrix Double
bgAttnWq = BlockParams -> Matrix Double
bpAttnWq BlockParams
bp,
      bgAttnWk :: Matrix Double
bgAttnWk = BlockParams -> Matrix Double
bpAttnWk BlockParams
bp,
      bgAttnWv :: Matrix Double
bgAttnWv = BlockParams -> Matrix Double
bpAttnWv BlockParams
bp,
      bgAttnWo :: Matrix Double
bgAttnWo = BlockParams -> Matrix Double
bpAttnWo BlockParams
bp,
      bgAttnLnGamma :: Vector Double
bgAttnLnGamma = BlockParams -> Vector Double
bpAttnLnGamma BlockParams
bp,
      bgAttnLnBeta :: Vector Double
bgAttnLnBeta = BlockParams -> Vector Double
bpAttnLnBeta BlockParams
bp,
      bgFfnW1 :: Matrix Double
bgFfnW1 = BlockParams -> Matrix Double
bpFfnW1 BlockParams
bp,
      bgFfnB1 :: Vector Double
bgFfnB1 = BlockParams -> Vector Double
bpFfnB1 BlockParams
bp,
      bgFfnW2 :: Matrix Double
bgFfnW2 = BlockParams -> Matrix Double
bpFfnW2 BlockParams
bp,
      bgFfnB2 :: Vector Double
bgFfnB2 = BlockParams -> Vector Double
bpFfnB2 BlockParams
bp,
      bgFfnLnGamma :: Vector Double
bgFfnLnGamma = BlockParams -> Vector Double
bpFfnLnGamma BlockParams
bp,
      bgFfnLnBeta :: Vector Double
bgFfnLnBeta = BlockParams -> Vector Double
bpFfnLnBeta BlockParams
bp
    }

----------------------------------------------------------------------
-- Full model backward passes
----------------------------------------------------------------------

-- | Shared embedding setup used by both GPT and BERT backward passes.
embedInput :: GptConfig -> Gpt -> [Int] -> (Matrix Double, Matrix Double, Matrix Double)
embedInput :: GptConfig
-> Gpt -> [Int] -> (Matrix Double, Matrix Double, Matrix Double)
embedInput GptConfig
cfg Gpt
model [Int]
inputIds =
  let seqLen :: Int
seqLen = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inputIds
      nEmb :: Int
nEmb = GptConfig -> Int
gptNEmbd GptConfig
cfg
      wte :: Matrix Double
wte = Gpt -> Matrix Double
gptWte Gpt
model
      wpe :: Matrix Double
wpe = Gpt -> Matrix Double
gptWpe Gpt
model
      tokEmb :: Matrix Double
tokEmb = [Vector Double] -> Matrix Double
forall t. Element t => [Vector t] -> Matrix t
fromRows [Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
wte [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i | Int
i <- [Int]
inputIds]
      posEmb :: Matrix Double
posEmb = Matrix Double -> Int -> Int -> Int -> Int -> Matrix Double
subMatrixW Matrix Double
wpe Int
0 Int
0 Int
seqLen Int
nEmb
      x0 :: Matrix Double
x0 = Matrix Double
tokEmb Matrix Double -> Matrix Double -> Matrix Double
forall a. Num a => a -> a -> a
+ Matrix Double
posEmb
   in (Matrix Double
x0, Matrix Double
wte, Matrix Double
wpe)

gptBackward :: GptConfig -> Gpt -> [Int] -> [Int] -> (Double, GptGrads)
gptBackward :: GptConfig -> Gpt -> [Int] -> [Int] -> (Double, GptGrads)
gptBackward GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds =
  let seqLen :: Int
seqLen = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inputIds
      nEmb :: Int
nEmb = GptConfig -> Int
gptNEmbd GptConfig
cfg
      eps :: Double
eps = Double
1e-5

      (Matrix Double
x0, Matrix Double
wte, Matrix Double
_) = GptConfig
-> Gpt -> [Int] -> (Matrix Double, Matrix Double, Matrix Double)
embedInput GptConfig
cfg Gpt
model [Int]
inputIds

      -- Run the differentiable GPT body
      gptP :: DiffP GptParams (Matrix Double) (Matrix Double)
gptP = GptConfig
-> Int -> Double -> DiffP GptParams (Matrix Double) (Matrix Double)
gptDiffP GptConfig
cfg Int
seqLen Double
eps
      params :: GptParams
params = Gpt -> GptParams
gptParamsFromModel Gpt
model
      logits :: Matrix Double
logits = DiffP GptParams (Matrix Double) (Matrix Double)
-> GptParams -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP GptParams (Matrix Double) (Matrix Double)
gptP GptParams
params Matrix Double
x0

      -- Loss and output gradient
      (Double
loss, Matrix Double
gradLogits) = Matrix Double -> [Int] -> (Double, Matrix Double)
crossEntropyBwd Matrix Double
logits [Int]
targetIds

      -- Backward through the whole GPT body
      (Matrix Double
gradX0, GptParams
gp) = DiffP GptParams (Matrix Double) (Matrix Double)
-> GptParams
-> Matrix Double
-> Matrix Double
-> (Matrix Double, GptParams)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP GptParams (Matrix Double) (Matrix Double)
gptP GptParams
params Matrix Double
x0 Matrix Double
gradLogits

      -- Embedding gradients
      gWte :: Matrix Double
gWte = Matrix Double -> [Int] -> Matrix Double -> Matrix Double
embedBackward Matrix Double
wte [Int]
inputIds Matrix Double
gradX0
      gWpeFull :: Matrix Double
gWpeFull = Int -> Int -> Matrix Double
zeroMatrix Int
1024 Int
nEmb
      gWpe' :: Matrix Double
gWpe' = Matrix Double -> Int -> Int -> Matrix Double -> Matrix Double
updateSubMatrix Matrix Double
gWpeFull Int
0 Int
0 Matrix Double
gradX0

      grads :: GptGrads
grads =
        (GptConfig -> GptGrads
zeroGptGrads GptConfig
cfg)
          { ggWte = gWte,
            ggWpe = gWpe',
            ggBlocks = map blockGradsFromParams (gpBlocks gp),
            ggLnGamma = gpLnGamma gp,
            ggLnBeta = gpLnBeta gp,
            ggHead = gpHead gp,
            ggHeadB = gpHeadB gp
          }
   in (Double
loss, GptGrads
grads)

-- | BERT-style masked-language-model backward pass.
--
--   * inputIds are the (possibly masked) token IDs fed to the model.
--   * targetIds are the original token IDs to reconstruct.
--   * mask indicates which positions are supervised.
--
-- The model uses bidirectional attention, so every position can attend to
-- every other position.
bertBackward :: GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> (Double, GptGrads)
bertBackward :: GptConfig -> Gpt -> [Int] -> [Int] -> [Bool] -> (Double, GptGrads)
bertBackward GptConfig
cfg Gpt
model [Int]
inputIds [Int]
targetIds [Bool]
mask =
  let seqLen :: Int
seqLen = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
inputIds
      nEmb :: Int
nEmb = GptConfig -> Int
gptNEmbd GptConfig
cfg
      eps :: Double
eps = Double
1e-5

      (Matrix Double
x0, Matrix Double
wte, Matrix Double
_) = GptConfig
-> Gpt -> [Int] -> (Matrix Double, Matrix Double, Matrix Double)
embedInput GptConfig
cfg Gpt
model [Int]
inputIds

      -- Run the differentiable bidirectional body
      bertP :: DiffP GptParams (Matrix Double) (Matrix Double)
bertP = GptConfig
-> Int -> Double -> DiffP GptParams (Matrix Double) (Matrix Double)
bertDiffP GptConfig
cfg Int
seqLen Double
eps
      params :: GptParams
params = Gpt -> GptParams
gptParamsFromModel Gpt
model
      logits :: Matrix Double
logits = DiffP GptParams (Matrix Double) (Matrix Double)
-> GptParams -> Matrix Double -> Matrix Double
forall p a b. TensorPrim p a b -> p -> a -> b
primForward DiffP GptParams (Matrix Double) (Matrix Double)
bertP GptParams
params Matrix Double
x0

      -- Masked loss and output gradient
      (Double
loss, Matrix Double
gradLogits) = Matrix Double -> [Int] -> [Bool] -> (Double, Matrix Double)
maskedCrossEntropyBwd Matrix Double
logits [Int]
targetIds [Bool]
mask

      -- Backward through the whole BERT body
      (Matrix Double
gradX0, GptParams
gp) = DiffP GptParams (Matrix Double) (Matrix Double)
-> GptParams
-> Matrix Double
-> Matrix Double
-> (Matrix Double, GptParams)
forall p a b. TensorPrim p a b -> p -> a -> b -> (a, p)
primBackward DiffP GptParams (Matrix Double) (Matrix Double)
bertP GptParams
params Matrix Double
x0 Matrix Double
gradLogits

      -- Embedding gradients
      gWte :: Matrix Double
gWte = Matrix Double -> [Int] -> Matrix Double -> Matrix Double
embedBackward Matrix Double
wte [Int]
inputIds Matrix Double
gradX0
      gWpeFull :: Matrix Double
gWpeFull = Int -> Int -> Matrix Double
zeroMatrix Int
1024 Int
nEmb
      gWpe' :: Matrix Double
gWpe' = Matrix Double -> Int -> Int -> Matrix Double -> Matrix Double
updateSubMatrix Matrix Double
gWpeFull Int
0 Int
0 Matrix Double
gradX0

      grads :: GptGrads
grads =
        (GptConfig -> GptGrads
zeroGptGrads GptConfig
cfg)
          { ggWte = gWte,
            ggWpe = gWpe',
            ggBlocks = map blockGradsFromParams (gpBlocks gp),
            ggLnGamma = gpLnGamma gp,
            ggLnBeta = gpLnBeta gp,
            ggHead = gpHead gp,
            ggHeadB = gpHeadB gp
          }
   in (Double
loss, GptGrads
grads)

----------------------------------------------------------------------
-- Utilities
----------------------------------------------------------------------

embedBackward :: Matrix Double -> [Int] -> Matrix Double -> Matrix Double
embedBackward :: Matrix Double -> [Int] -> Matrix Double -> Matrix Double
embedBackward Matrix Double
wte [Int]
inputIds Matrix Double
gradX =
  let rowsGradX :: [Vector Double]
rowsGradX = Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
gradX
      updateFn :: Matrix Double -> (Int, Int) -> Matrix Double
updateFn Matrix Double
acc (Int
i, Int
tid) =
        let oldRow :: Vector Double
oldRow = Matrix Double -> [Vector Double]
forall t. Element t => Matrix t -> [Vector t]
toRows Matrix Double
acc [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
tid
            newRow :: Vector Double
newRow = Vector Double
oldRow Vector Double -> Vector Double -> Vector Double
forall a. Num a => a -> a -> a
+ ([Vector Double]
rowsGradX [Vector Double] -> Int -> Vector Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
i)
            vals :: [Double]
vals =
              [ if Int
ri Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
tid
                  then Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList Vector Double
newRow [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! Int
ci
                  else Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
acc) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
ri Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ci)
              | Int
ri <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1],
                Int
ci <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
              ]
         in Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
acc) ([Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList [Double]
vals)
   in (Matrix Double -> (Int, Int) -> Matrix Double)
-> Matrix Double -> [(Int, Int)] -> Matrix Double
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Matrix Double -> (Int, Int) -> Matrix Double
updateFn (Int -> Int -> Matrix Double
zeroMatrix (Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
wte) (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
wte)) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Int]
inputIds)

updateSubMatrix :: Matrix Double -> Int -> Int -> Matrix Double -> Matrix Double
updateSubMatrix :: Matrix Double -> Int -> Int -> Matrix Double -> Matrix Double
updateSubMatrix Matrix Double
target Int
r0 Int
c0 Matrix Double
patch =
  Int -> Vector Double -> Matrix Double
forall t. Storable t => Int -> Vector t -> Matrix t
reshape
    (Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
target)
    ( [Double] -> Vector Double
forall a. Storable a => [a] -> Vector a
fromList
        [ if Int
ri Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
r0 Bool -> Bool -> Bool
&& Int
ri Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
patch Bool -> Bool -> Bool
&& Int
ci Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
c0 Bool -> Bool -> Bool
&& Int
ci Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
c0 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
patch
            then Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
patch) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! ((Int
ri Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
r0) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
patch Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
ci Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c0))
            else Vector Double -> [Double]
forall a. Storable a => Vector a -> [a]
toList (Matrix Double -> Vector Double
forall t. Element t => Matrix t -> Vector t
LA.flatten Matrix Double
target) [Double] -> Int -> Double
forall a. HasCallStack => [a] -> Int -> a
!! (Int
ri Int -> Int -> Int
forall a. Num a => a -> a -> a
* Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
target Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ci)
        | Int
ri <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
rows Matrix Double
target Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1],
          Int
ci <- [Int
0 .. Matrix Double -> Int
forall t. Matrix t -> Int
cols Matrix Double
target Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
        ]
    )