@dniskav/neuron 0.1.6 → 0.2.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -21,15 +21,22 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
21
21
  var index_exports = {};
22
22
  __export(index_exports, {
23
23
  Adam: () => Adam,
24
+ AttentionHead: () => AttentionHead,
25
+ EmbeddingMatrix: () => EmbeddingMatrix,
24
26
  LSTMLayer: () => LSTMLayer,
25
27
  Layer: () => Layer,
28
+ LayerNorm: () => LayerNorm,
26
29
  Momentum: () => Momentum,
30
+ MultiHeadAttention: () => MultiHeadAttention,
27
31
  Network: () => Network,
28
32
  NetworkLSTM: () => NetworkLSTM,
29
33
  NetworkN: () => NetworkN,
34
+ NetworkTransformer: () => NetworkTransformer,
30
35
  Neuron: () => Neuron,
31
36
  NeuronN: () => NeuronN,
32
37
  SGD: () => SGD,
38
+ TransformerBlock: () => TransformerBlock,
39
+ WeightMatrix: () => WeightMatrix,
33
40
  crossEntropy: () => crossEntropy,
34
41
  crossEntropyDelta: () => crossEntropyDelta,
35
42
  crossEntropyDeltaRaw: () => crossEntropyDeltaRaw,
@@ -38,11 +45,15 @@ __export(index_exports, {
38
45
  linear: () => linear,
39
46
  makeElu: () => makeElu,
40
47
  makeLeakyRelu: () => makeLeakyRelu,
48
+ matMul: () => matMul,
41
49
  mse: () => mse,
42
50
  mseDelta: () => mseDelta,
43
51
  relu: () => relu,
44
52
  sigmoid: () => sigmoid2,
45
- tanh: () => tanh
53
+ softmax: () => softmax,
54
+ softmaxBackward: () => softmaxBackward,
55
+ tanh: () => tanh,
56
+ transpose: () => transpose
46
57
  });
47
58
  module.exports = __toCommonJS(index_exports);
48
59
 
@@ -523,6 +534,565 @@ var NetworkLSTM = class {
523
534
  }
524
535
  };
525
536
 
537
+ // src/MatMul.ts
538
+ function matMul(A, B) {
539
+ const rows = A.length;
540
+ const inner = B.length;
541
+ const cols = B[0].length;
542
+ const C = Array.from({ length: rows }, () => new Array(cols).fill(0));
543
+ for (let i = 0; i < rows; i++)
544
+ for (let k = 0; k < inner; k++) {
545
+ const aik = A[i][k];
546
+ for (let j = 0; j < cols; j++)
547
+ C[i][j] += aik * B[k][j];
548
+ }
549
+ return C;
550
+ }
551
+ function transpose(A) {
552
+ const rows = A.length, cols = A[0].length;
553
+ const T = Array.from({ length: cols }, () => new Array(rows).fill(0));
554
+ for (let i = 0; i < rows; i++)
555
+ for (let j = 0; j < cols; j++)
556
+ T[j][i] = A[i][j];
557
+ return T;
558
+ }
559
+ function softmax(row) {
560
+ const max = Math.max(...row);
561
+ const exps = row.map((v) => Math.exp(v - max));
562
+ const sum = exps.reduce((a, b) => a + b, 0);
563
+ return exps.map((e) => e / sum);
564
+ }
565
+ function softmaxBackward(dS, s) {
566
+ const dot = s.reduce((acc, si, i) => acc + dS[i] * si, 0);
567
+ return s.map((si, i) => si * (dS[i] - dot));
568
+ }
569
+ var WeightMatrix = class {
570
+ constructor(rows, cols) {
571
+ const limit = Math.sqrt(2 / (rows + cols));
572
+ this.W = Array.from(
573
+ { length: rows },
574
+ () => Array.from({ length: cols }, () => (Math.random() * 2 - 1) * limit)
575
+ );
576
+ this.opts = Array.from(
577
+ { length: rows },
578
+ () => Array.from({ length: cols }, () => new Adam())
579
+ );
580
+ }
581
+ // Apply pre-computed gradient (same shape as W).
582
+ update(dW, lr) {
583
+ for (let i = 0; i < this.W.length; i++)
584
+ for (let j = 0; j < this.W[0].length; j++)
585
+ this.W[i][j] = this.opts[i][j].step(this.W[i][j], dW[i][j], lr);
586
+ }
587
+ };
588
+ var EmbeddingMatrix = class {
589
+ constructor(vocabSize, d_model) {
590
+ const limit = Math.sqrt(1 / d_model);
591
+ this.W = Array.from(
592
+ { length: vocabSize },
593
+ () => Array.from({ length: d_model }, () => (Math.random() * 2 - 1) * limit)
594
+ );
595
+ }
596
+ get(idx) {
597
+ return [...this.W[idx]];
598
+ }
599
+ update(idx, grad, lr) {
600
+ for (let m = 0; m < this.W[idx].length; m++)
601
+ this.W[idx][m] += lr * grad[m];
602
+ }
603
+ };
604
+
605
+ // src/AttentionHead.ts
606
+ var AttentionHead = class {
607
+ constructor(d_model, d_k, d_v) {
608
+ // d_v × d_model
609
+ this.cache = null;
610
+ this.d_k = d_k;
611
+ this.d_v = d_v;
612
+ this.Wq = new WeightMatrix(d_k, d_model);
613
+ this.Wk = new WeightMatrix(d_k, d_model);
614
+ this.Wv = new WeightMatrix(d_v, d_model);
615
+ }
616
+ // ── Forward ───────────────────────────────────────────────────────────────
617
+ // X: seqLen × d_model → out: seqLen × d_v
618
+ predict(X) {
619
+ const seqLen = X.length;
620
+ const scale = 1 / Math.sqrt(this.d_k);
621
+ const Q = X.map(
622
+ (x) => this.Wq.W.map((wq) => wq.reduce((s, w, m) => s + w * x[m], 0))
623
+ );
624
+ const K = X.map(
625
+ (x) => this.Wk.W.map((wk) => wk.reduce((s, w, m) => s + w * x[m], 0))
626
+ );
627
+ const V = X.map(
628
+ (x) => this.Wv.W.map((wv) => wv.reduce((s, w, m) => s + w * x[m], 0))
629
+ );
630
+ const scores = Array.from(
631
+ { length: seqLen },
632
+ (_, i) => Array.from(
633
+ { length: seqLen },
634
+ (_2, j) => Q[i].reduce((s, q, k) => s + q * K[j][k], 0) * scale
635
+ )
636
+ );
637
+ const attn = scores.map((row) => softmax(row));
638
+ const out = Array.from(
639
+ { length: seqLen },
640
+ (_, i) => Array.from(
641
+ { length: this.d_v },
642
+ (_2, d) => attn[i].reduce((s, a, j) => s + a * V[j][d], 0)
643
+ )
644
+ );
645
+ this.cache = { X, Q, K, V, scores, attn };
646
+ return out;
647
+ }
648
+ // ── Backward ──────────────────────────────────────────────────────────────
649
+ // dOut: seqLen × d_v → dX: seqLen × d_model
650
+ //
651
+ // Steps:
652
+ // 1. dV = attn^T @ dOut
653
+ // 2. dAttn = dOut @ V^T (attention weight gradients)
654
+ // 3. dScores = softmaxBwd(dAttn) / √d_k
655
+ // 4. dQ = dScores @ K, dK = dScores^T @ Q
656
+ // 5. dWq = dQ^T @ X, dWk = dK^T @ X, dWv = dV^T @ X
657
+ // 6. dX = dQ @ Wq + dK @ Wk + dV @ Wv
658
+ backward(dOut, lr) {
659
+ const { X, Q, K, V, attn } = this.cache;
660
+ const seqLen = X.length;
661
+ const d_model = X[0].length;
662
+ const scale = 1 / Math.sqrt(this.d_k);
663
+ const dV = Array.from(
664
+ { length: seqLen },
665
+ (_, j) => Array.from(
666
+ { length: this.d_v },
667
+ (_2, d) => attn.reduce((s, a, i) => s + a[j] * dOut[i][d], 0)
668
+ )
669
+ );
670
+ const dAttn = Array.from(
671
+ { length: seqLen },
672
+ (_, i) => Array.from(
673
+ { length: seqLen },
674
+ (_2, j) => dOut[i].reduce((s, d, k) => s + d * V[j][k], 0)
675
+ )
676
+ );
677
+ const dScores = dAttn.map(
678
+ (da, i) => softmaxBackward(da, attn[i]).map((v) => v * scale)
679
+ );
680
+ const dQ = Array.from(
681
+ { length: seqLen },
682
+ (_, i) => Array.from(
683
+ { length: this.d_k },
684
+ (_2, k) => dScores[i].reduce((s, ds, j) => s + ds * K[j][k], 0)
685
+ )
686
+ );
687
+ const dK = Array.from(
688
+ { length: seqLen },
689
+ (_, j) => Array.from(
690
+ { length: this.d_k },
691
+ (_2, k) => dScores.reduce((s, ds, i) => s + ds[j] * Q[i][k], 0)
692
+ )
693
+ );
694
+ const dWq = Array.from(
695
+ { length: this.d_k },
696
+ (_, k) => Array.from(
697
+ { length: d_model },
698
+ (_2, m) => dQ.reduce((s, dq, i) => s + dq[k] * X[i][m], 0)
699
+ )
700
+ );
701
+ const dWk = Array.from(
702
+ { length: this.d_k },
703
+ (_, k) => Array.from(
704
+ { length: d_model },
705
+ (_2, m) => dK.reduce((s, dk, i) => s + dk[k] * X[i][m], 0)
706
+ )
707
+ );
708
+ const dWv = Array.from(
709
+ { length: this.d_v },
710
+ (_, k) => Array.from(
711
+ { length: d_model },
712
+ (_2, m) => dV.reduce((s, dv, i) => s + dv[k] * X[i][m], 0)
713
+ )
714
+ );
715
+ this.Wq.update(dWq, lr);
716
+ this.Wk.update(dWk, lr);
717
+ this.Wv.update(dWv, lr);
718
+ const dX = Array.from(
719
+ { length: seqLen },
720
+ (_, i) => Array.from(
721
+ { length: d_model },
722
+ (_2, m) => dQ[i].reduce((s, dq, k) => s + dq * this.Wq.W[k][m], 0) + dK[i].reduce((s, dk, k) => s + dk * this.Wk.W[k][m], 0) + dV[i].reduce((s, dv, k) => s + dv * this.Wv.W[k][m], 0)
723
+ )
724
+ );
725
+ return dX;
726
+ }
727
+ // Attention weights from the last predict() call — useful for visualization.
728
+ getAttentionWeights() {
729
+ return this.cache ? this.cache.attn : null;
730
+ }
731
+ };
732
+
733
+ // src/MultiHeadAttention.ts
734
+ var MultiHeadAttention = class {
735
+ // seqLen × (nHeads * d_k)
736
+ constructor(d_model, nHeads) {
737
+ // d_model × (nHeads * d_k)
738
+ // Cached for backward
739
+ this._concat = null;
740
+ this.nHeads = nHeads;
741
+ this.d_model = d_model;
742
+ this.d_k = Math.floor(d_model / nHeads);
743
+ this.heads = Array.from(
744
+ { length: nHeads },
745
+ () => new AttentionHead(d_model, this.d_k, this.d_k)
746
+ );
747
+ this.Wo = new WeightMatrix(d_model, nHeads * this.d_k);
748
+ }
749
+ // ── Forward ───────────────────────────────────────────────────────────────
750
+ // X: seqLen × d_model → out: seqLen × d_model
751
+ predict(X) {
752
+ const seqLen = X.length;
753
+ const headOuts = this.heads.map((h) => h.predict(X));
754
+ const concat = Array.from(
755
+ { length: seqLen },
756
+ (_, i) => headOuts.flatMap((ho) => ho[i])
757
+ );
758
+ const out = concat.map(
759
+ (c) => this.Wo.W.map((row) => row.reduce((s, w, j) => s + w * c[j], 0))
760
+ );
761
+ this._concat = concat;
762
+ return out;
763
+ }
764
+ // ── Backward ──────────────────────────────────────────────────────────────
765
+ // dOut: seqLen × d_model → dX: seqLen × d_model
766
+ backward(dOut, lr) {
767
+ const seqLen = dOut.length;
768
+ const concatD = this.nHeads * this.d_k;
769
+ const d_model = this.d_model;
770
+ const concat = this._concat;
771
+ const dConcat = dOut.map(
772
+ (do_) => Array.from(
773
+ { length: concatD },
774
+ (_, j) => this.Wo.W.reduce((s, row, k) => s + do_[k] * row[j], 0)
775
+ )
776
+ );
777
+ const dWo = Array.from(
778
+ { length: d_model },
779
+ (_, k) => Array.from(
780
+ { length: concatD },
781
+ (_2, j) => dOut.reduce((s, row, i) => s + row[k] * concat[i][j], 0)
782
+ )
783
+ );
784
+ this.Wo.update(dWo, lr);
785
+ const dX = Array.from(
786
+ { length: seqLen },
787
+ () => new Array(d_model).fill(0)
788
+ );
789
+ for (let h = 0; h < this.nHeads; h++) {
790
+ const start = h * this.d_k;
791
+ const dHeadOut = dConcat.map((dc) => dc.slice(start, start + this.d_k));
792
+ const dXh = this.heads[h].backward(dHeadOut, lr);
793
+ for (let i = 0; i < seqLen; i++)
794
+ for (let m = 0; m < d_model; m++)
795
+ dX[i][m] += dXh[i][m];
796
+ }
797
+ return dX;
798
+ }
799
+ // Attention weights per head from the last predict() — for visualization.
800
+ // Returns: nHeads × seqLen × seqLen
801
+ getAttentionWeights() {
802
+ return this.heads.map((h) => h.getAttentionWeights());
803
+ }
804
+ };
805
+
806
+ // src/LayerNorm.ts
807
+ var LayerNorm = class {
808
+ constructor(dim) {
809
+ this.eps = 1e-5;
810
+ // Per-position cache populated during the forward pass.
811
+ // resetCache() must be called before each sequence forward pass.
812
+ this._cache = [];
813
+ this.gamma = new Array(dim).fill(1);
814
+ this.beta = new Array(dim).fill(0);
815
+ }
816
+ // Call once before forward-passing a new sequence.
817
+ resetCache(seqLen) {
818
+ this._cache = new Array(seqLen);
819
+ }
820
+ // Normalize a single position's feature vector.
821
+ // pos must match the position index used in the corresponding backwardOne call.
822
+ predictOne(x, pos) {
823
+ const N = x.length;
824
+ const mean = x.reduce((s, v) => s + v, 0) / N;
825
+ const vari = x.reduce((s, v) => s + (v - mean) ** 2, 0) / N;
826
+ const std = Math.sqrt(vari + this.eps);
827
+ const x_norm = x.map((v) => (v - mean) / std);
828
+ this._cache[pos] = { x_norm, std };
829
+ return x_norm.map((xn, i) => this.gamma[i] * xn + this.beta[i]);
830
+ }
831
+ // Backprop through layer norm for one position.
832
+ //
833
+ // Given dL/dy (dOut), computes dL/dx:
834
+ // Let D = dOut ⊙ γ
835
+ // dL/dx_i = (1/std) * (D_i − mean(D) − x_norm_i * mean(D ⊙ x_norm))
836
+ //
837
+ // Also updates γ and β via SGD:
838
+ // γ_i += lr * dOut_i * x_norm_i
839
+ // β_i += lr * dOut_i
840
+ //
841
+ // SGD (not Adam) for γ/β: they are aggregated across all positions in the
842
+ // sequence (de-facto mini-batch update), so the gradient is already smoothed.
843
+ backwardOne(dOut, pos, lr) {
844
+ const { x_norm, std } = this._cache[pos];
845
+ const N = dOut.length;
846
+ for (let i = 0; i < N; i++) {
847
+ this.gamma[i] += lr * dOut[i] * x_norm[i];
848
+ this.beta[i] += lr * dOut[i];
849
+ }
850
+ const D = dOut.map((d, i) => d * this.gamma[i]);
851
+ const mD = D.reduce((s, v) => s + v, 0) / N;
852
+ const mDxn = D.reduce((s, d, i) => s + d * x_norm[i], 0) / N;
853
+ return D.map((d, i) => (d - mD - x_norm[i] * mDxn) / std);
854
+ }
855
+ };
856
+
857
+ // src/TransformerBlock.ts
858
+ var TransformerBlock = class {
859
+ constructor({ d_model, nHeads, d_ff }) {
860
+ // Forward caches (needed for backprop)
861
+ this._X = null;
862
+ this._attnOut = null;
863
+ this._h1 = null;
864
+ this._ff1Pre = null;
865
+ // pre-ReLU
866
+ this._ff1Out = null;
867
+ // post-ReLU
868
+ this._ff2Out = null;
869
+ this.d_model = d_model;
870
+ this.d_ff = d_ff;
871
+ this.attn = new MultiHeadAttention(d_model, nHeads);
872
+ this.norm1 = new LayerNorm(d_model);
873
+ this.norm2 = new LayerNorm(d_model);
874
+ this.ff1 = new WeightMatrix(d_ff, d_model);
875
+ this.ff2 = new WeightMatrix(d_model, d_ff);
876
+ this.b1 = new Array(d_ff).fill(0);
877
+ this.b2 = new Array(d_model).fill(0);
878
+ this.b1Opts = Array.from({ length: d_ff }, () => new Adam());
879
+ this.b2Opts = Array.from({ length: d_model }, () => new Adam());
880
+ }
881
+ // ── Forward ───────────────────────────────────────────────────────────────
882
+ // X: seqLen × d_model → out: seqLen × d_model
883
+ predict(X) {
884
+ const seqLen = X.length;
885
+ const attnOut = this.attn.predict(X);
886
+ this.norm1.resetCache(seqLen);
887
+ const h1 = X.map((x, i) => {
888
+ const added = x.map((v, k) => v + attnOut[i][k]);
889
+ return this.norm1.predictOne(added, i);
890
+ });
891
+ const ff1Pre = h1.map(
892
+ (h) => this.ff1.W.map((row, k) => row.reduce((s, w, m) => s + w * h[m], this.b1[k]))
893
+ );
894
+ const ff1Out = ff1Pre.map((pre) => pre.map((v) => Math.max(0, v)));
895
+ const ff2Out = ff1Out.map(
896
+ (h) => this.ff2.W.map((row, k) => row.reduce((s, w, m) => s + w * h[m], this.b2[k]))
897
+ );
898
+ this.norm2.resetCache(seqLen);
899
+ const out = h1.map((h, i) => {
900
+ const added = h.map((v, k) => v + ff2Out[i][k]);
901
+ return this.norm2.predictOne(added, i);
902
+ });
903
+ this._X = X;
904
+ this._attnOut = attnOut;
905
+ this._h1 = h1;
906
+ this._ff1Pre = ff1Pre;
907
+ this._ff1Out = ff1Out;
908
+ this._ff2Out = ff2Out;
909
+ return out;
910
+ }
911
+ // ── Backward ──────────────────────────────────────────────────────────────
912
+ // dOut: seqLen × d_model → dX: seqLen × d_model
913
+ backward(dOut, lr) {
914
+ const seqLen = dOut.length;
915
+ const d_model = this.d_model;
916
+ const h1 = this._h1;
917
+ const ff1Out = this._ff1Out;
918
+ const ff1Pre = this._ff1Pre;
919
+ const dAdded2 = dOut.map((do_, i) => this.norm2.backwardOne(do_, i, lr));
920
+ const dFf1Out = dAdded2.map(
921
+ (da) => Array.from(
922
+ { length: this.d_ff },
923
+ (_, k) => this.ff2.W.reduce((s, row, m) => s + row[k] * da[m], 0)
924
+ )
925
+ );
926
+ const dW2 = Array.from(
927
+ { length: d_model },
928
+ (_, m) => Array.from(
929
+ { length: this.d_ff },
930
+ (_2, k) => dAdded2.reduce((s, da, i) => s + da[m] * ff1Out[i][k], 0)
931
+ )
932
+ );
933
+ const db2 = Array.from(
934
+ { length: d_model },
935
+ (_, m) => dAdded2.reduce((s, da) => s + da[m], 0)
936
+ );
937
+ this.ff2.update(dW2, lr);
938
+ for (let m = 0; m < d_model; m++)
939
+ this.b2[m] = this.b2Opts[m].step(this.b2[m], db2[m], lr);
940
+ const dFf1Pre = dFf1Out.map(
941
+ (d, i) => d.map((v, k) => ff1Pre[i][k] > 0 ? v : 0)
942
+ );
943
+ const dH1_fromFf = dFf1Pre.map(
944
+ (dp) => Array.from(
945
+ { length: d_model },
946
+ (_, m) => this.ff1.W.reduce((s, row, k) => s + dp[k] * row[m], 0)
947
+ )
948
+ );
949
+ const dW1 = Array.from(
950
+ { length: this.d_ff },
951
+ (_, k) => Array.from(
952
+ { length: d_model },
953
+ (_2, m) => dFf1Pre.reduce((s, dp, i) => s + dp[k] * h1[i][m], 0)
954
+ )
955
+ );
956
+ const db1 = Array.from(
957
+ { length: this.d_ff },
958
+ (_, k) => dFf1Pre.reduce((s, dp) => s + dp[k], 0)
959
+ );
960
+ this.ff1.update(dW1, lr);
961
+ for (let k = 0; k < this.d_ff; k++)
962
+ this.b1[k] = this.b1Opts[k].step(this.b1[k], db1[k], lr);
963
+ const dH1 = Array.from(
964
+ { length: seqLen },
965
+ (_, i) => dH1_fromFf[i].map((v, m) => v + dAdded2[i][m])
966
+ );
967
+ const dAdded1 = dH1.map((d, i) => this.norm1.backwardOne(d, i, lr));
968
+ const dAttnOut = dAdded1;
969
+ const dX_skip = dAdded1;
970
+ const dX_fromAttn = this.attn.backward(dAttnOut, lr);
971
+ const dX = Array.from(
972
+ { length: seqLen },
973
+ (_, i) => Array.from(
974
+ { length: d_model },
975
+ (_2, m) => dX_fromAttn[i][m] + dX_skip[i][m]
976
+ )
977
+ );
978
+ return dX;
979
+ }
980
+ // Attention weights from the last predict() — for visualization.
981
+ getAttentionWeights() {
982
+ return this.attn.getAttentionWeights();
983
+ }
984
+ };
985
+
986
+ // src/NetworkTransformer.ts
987
+ var NetworkTransformer = class {
988
+ constructor(seqLen, options = {}) {
989
+ const {
990
+ vocabSize = 10,
991
+ d_model = 64,
992
+ nHeads = 4,
993
+ d_ff = 128,
994
+ nBlocks = 4,
995
+ nClasses = 9
996
+ } = options;
997
+ this.seqLen = seqLen;
998
+ this.vocabSize = vocabSize;
999
+ this.d_model = d_model;
1000
+ this.nClasses = nClasses;
1001
+ this.tokenEmb = new EmbeddingMatrix(vocabSize, d_model);
1002
+ this.posEmb = new EmbeddingMatrix(seqLen, d_model);
1003
+ this.blocks = Array.from(
1004
+ { length: nBlocks },
1005
+ () => new TransformerBlock({ d_model, nHeads, d_ff })
1006
+ );
1007
+ this.outputProj = new WeightMatrix(nClasses, d_model);
1008
+ this.outputBias = new Array(nClasses).fill(0);
1009
+ this.outBiasOpts = Array.from({ length: nClasses }, () => new Adam());
1010
+ }
1011
+ // ── Forward pass ──────────────────────────────────────────────────────────
1012
+ // tokens: seqLen integer ids → seqLen * nClasses logits (flattened)
1013
+ predict(tokens) {
1014
+ const h = this._forward(tokens);
1015
+ return h.flatMap(
1016
+ (hi) => this.outputProj.W.map(
1017
+ (row, c) => row.reduce((s, w, m) => s + w * hi[m], this.outputBias[c])
1018
+ )
1019
+ );
1020
+ }
1021
+ // ── Training step (online, one sample at a time) ───────────────────────────
1022
+ // tokens: seqLen integer ids
1023
+ // targets: seqLen * nClasses values (e.g. one-hot per cell)
1024
+ // mask: optional boolean[seqLen] — only compute loss/gradients for
1025
+ // positions where mask[i] = true (e.g. empty cells in Sudoku)
1026
+ // Returns: MSE loss over the masked positions.
1027
+ train(tokens, targets, lr, mask) {
1028
+ const h = this._forward(tokens);
1029
+ const logits = h.map(
1030
+ (hi) => this.outputProj.W.map(
1031
+ (row, c) => row.reduce((s, w, m) => s + w * hi[m], this.outputBias[c])
1032
+ )
1033
+ );
1034
+ let loss = 0;
1035
+ let count = 0;
1036
+ const dLogits = Array.from({ length: this.seqLen }, (_, i) => {
1037
+ if (mask && !mask[i]) return new Array(this.nClasses).fill(0);
1038
+ count++;
1039
+ return Array.from({ length: this.nClasses }, (_2, c) => {
1040
+ const t = targets[i * this.nClasses + c];
1041
+ const p = logits[i][c];
1042
+ loss += (p - t) ** 2;
1043
+ return 2 * (p - t);
1044
+ });
1045
+ });
1046
+ if (count > 0) loss /= count * this.nClasses;
1047
+ const dH = Array.from(
1048
+ { length: this.seqLen },
1049
+ (_, i) => Array.from(
1050
+ { length: this.d_model },
1051
+ (_2, m) => dLogits[i].reduce((s, dl, c) => s + dl * this.outputProj.W[c][m], 0)
1052
+ )
1053
+ );
1054
+ const dWout = Array.from(
1055
+ { length: this.nClasses },
1056
+ (_, c) => Array.from(
1057
+ { length: this.d_model },
1058
+ (_2, m) => dLogits.reduce((s, dl, i) => s + dl[c] * h[i][m], 0)
1059
+ )
1060
+ );
1061
+ const dBout = Array.from(
1062
+ { length: this.nClasses },
1063
+ (_, c) => dLogits.reduce((s, dl) => s + dl[c], 0)
1064
+ );
1065
+ this.outputProj.update(dWout, lr);
1066
+ for (let c = 0; c < this.nClasses; c++)
1067
+ this.outputBias[c] = this.outBiasOpts[c].step(this.outputBias[c], dBout[c], lr);
1068
+ let dX = dH;
1069
+ for (let b = this.blocks.length - 1; b >= 0; b--)
1070
+ dX = this.blocks[b].backward(dX, lr);
1071
+ for (let i = 0; i < this.seqLen; i++) {
1072
+ this.tokenEmb.update(tokens[i], dX[i], lr);
1073
+ this.posEmb.update(i, dX[i], lr);
1074
+ }
1075
+ return loss;
1076
+ }
1077
+ // Attention weights from every block for visualization.
1078
+ // Returns: nBlocks × nHeads × seqLen × seqLen (nulls if not yet run)
1079
+ getAttentionWeights() {
1080
+ return this.blocks.map((b) => b.getAttentionWeights());
1081
+ }
1082
+ // ── Internal ──────────────────────────────────────────────────────────────
1083
+ // Shared embedding + block forward pass.
1084
+ _forward(tokens) {
1085
+ let h = tokens.map((id, i) => {
1086
+ const te = this.tokenEmb.get(id);
1087
+ const pe = this.posEmb.get(i);
1088
+ return te.map((v, m) => v + pe[m]);
1089
+ });
1090
+ for (const block of this.blocks)
1091
+ h = block.predict(h);
1092
+ return h;
1093
+ }
1094
+ };
1095
+
526
1096
  // src/losses.ts
527
1097
  function mse(predicted, actual) {
528
1098
  return predicted.reduce((sum, p, i) => sum + (actual[i] - p) ** 2, 0) / predicted.length;
@@ -548,15 +1118,22 @@ function crossEntropyDeltaRaw(predicted, actual) {
548
1118
  // Annotate the CommonJS export names for ESM import in node:
549
1119
  0 && (module.exports = {
550
1120
  Adam,
1121
+ AttentionHead,
1122
+ EmbeddingMatrix,
551
1123
  LSTMLayer,
552
1124
  Layer,
1125
+ LayerNorm,
553
1126
  Momentum,
1127
+ MultiHeadAttention,
554
1128
  Network,
555
1129
  NetworkLSTM,
556
1130
  NetworkN,
1131
+ NetworkTransformer,
557
1132
  Neuron,
558
1133
  NeuronN,
559
1134
  SGD,
1135
+ TransformerBlock,
1136
+ WeightMatrix,
560
1137
  crossEntropy,
561
1138
  crossEntropyDelta,
562
1139
  crossEntropyDeltaRaw,
@@ -565,9 +1142,13 @@ function crossEntropyDeltaRaw(predicted, actual) {
565
1142
  linear,
566
1143
  makeElu,
567
1144
  makeLeakyRelu,
1145
+ matMul,
568
1146
  mse,
569
1147
  mseDelta,
570
1148
  relu,
571
1149
  sigmoid,
572
- tanh
1150
+ softmax,
1151
+ softmaxBackward,
1152
+ tanh,
1153
+ transpose
573
1154
  });