| (in-package #:nekod.moe)
|
|
|
|
|
|
|
|
|
| (defun gaussian-noise (mean stddev)
|
| "Box-Muller transform for Gaussian random number."
|
| (let* ((u1 (+ 1e-10 (random 1.0d0)))
|
| (u2 (random 1.0d0))
|
| (z (* (sqrt (* -2.0d0 (log u1))) (cos (* 2.0d0 pi u2)))))
|
| (+ mean (* stddev z))))
|
|
|
| (defun noisy-top-k (router h &optional (noise-stddev 0.1d0) (training t))
|
| "Top-K gating with noise for training. In inference mode: no noise."
|
| (let* ((scores (all-scores router h))
|
| (noisy-scores
|
| (if training
|
| (mapcar (lambda (s)
|
| (cons (car s)
|
| (+ (cdr s) (gaussian-noise 0.0d0 noise-stddev))))
|
| scores)
|
| scores))
|
| (sorted (sort noisy-scores #'> :key #'cdr))
|
| (selected (subseq sorted 0 (min (router-k router) (length sorted)))))
|
| (softmax-over-selected selected)))
|
|
|
|
|
| (defun load-balance-loss (router activations)
|
| "Auxiliary loss to prevent all traffic going to one expert.
|
| Returns scalar loss value — lower is more balanced."
|
| (let* ((n (length (router-expert-names router)))
|
| (counts (make-array n :initial-element 0)))
|
| (dolist (act activations)
|
| (let ((idx (position (car act) (router-expert-names router))))
|
| (when idx (incf (aref counts idx)))))
|
| (let ((total (reduce #'+ counts)))
|
| (if (zerop total) 0.0d0
|
| (loop for c across counts
|
| sum (let ((p (/ c total)))
|
| (if (zerop p) 0.0d0 (* p (log p)))))))))
|
|
|
