Training Language Models at Scale

7.1 Pretraining: Next-Token Prediction

The pretraining task: given tokens [t₁, t₂, …, t_n], predict [t₂, t₃, …, t_{n+1}]. This self-supervised objective can be applied to any text corpus without labels.

-- Sliding window over tokenised text
createTrainingExamples :: Int -> [Int] -> [([Int], Int)]
createTrainingExamples contextLength tokens =
  [ (take contextLength (drop i tokens), tokens !! (i + contextLength))
  | i <- [0 .. length tokens - contextLength - 1] ]

-- Batch structure
data Batch = Batch
  { inputIds     :: Matrix Int
  , targetIds    :: Vector Int
  , attentionMask :: Matrix Bool
  }

createBatches :: Int -> Int -> [([Int], Int)] -> [Batch]
createBatches batchSize contextLength examples =
  map createBatch (chunksOf batchSize examples)
  where
    createBatch batch =
      Batch { inputIds  = map fst batch
            , targetIds = map snd batch
            , attentionMask = trueMatrix (length batch) contextLength }

7.2 Training Loop

data TrainingState = TrainingState
  { model     :: LanguageModel
  , optimizer :: AdamOptimizer
  , step      :: Int
  , loss      :: Double
  }

-- One training step
trainStep :: TrainingState -> Batch -> TrainingState
trainStep state batch =
  let logits      = forward (model state) (inputIds batch)
      batchLoss   = crossEntropyLoss logits (targetIds batch)
      gradients   = backprop (model state) logits (targetIds batch)
      updatedModel = updateParameters (optimizer state) (model state) gradients
  in state { model = updatedModel
           , step  = step state + 1
           , loss  = batchLoss }

-- Full training run
train :: LanguageModel -> Int -> [Batch] -> LanguageModel
train model numEpochs batches =
  let initialState = TrainingState
        { model = model, optimizer = adamOptimizer, step = 0, loss = 0 }
      finalState   = foldl (\s _ -> trainEpoch s batches) initialState [1..numEpochs]
  in model finalState

trainEpoch :: TrainingState -> [Batch] -> TrainingState
trainEpoch = foldl trainStep

7.3 Learning Rate Scheduling

Learning rates that are too high lead to unstable training; too low leads to slow convergence. The original Transformer paper uses a warmup-then-decay schedule:

learningRateSchedule
  :: Double   -- d_model
  -> Int      -- warmup_steps
  -> Int      -- current_step
  -> Double
learningRateSchedule d_model warmupSteps currentStep =
  let step    = fromIntegral currentStep
      warmup  = fromIntegral warmupSteps
      factor  = d_model ** (-0.5)
      schedule = min (step ** (-0.5)) (step * warmup ** (-1.5))
  in factor * schedule

7.4 Checkpointing

saveCheckpoint :: FilePath -> LanguageModel -> Int -> IO ()
saveCheckpoint path model step = do
  encodeFile (path ++ "/model_" ++ show step ++ ".model") model
  writeFile  (path ++ "/metadata.txt") (show step)

loadCheckpoint :: FilePath -> IO (LanguageModel, Int)
loadCheckpoint path = do
  step  <- read <$> readFile (path ++ "/metadata.txt")
  model <- decodeFile (path ++ "/model_" ++ show step ++ ".model")
  return (model, step)

-- Evaluation
evaluate :: LanguageModel -> [Batch] -> Double
evaluate model batches =
  let losses = map (computeBatchLoss model) batches
  in average losses