catniff 0.6.4 → 0.6.5

package/README.md

@@ -76,9 +76,9 @@ optim.step();
 console.log("Updated weight:", w.data);  // Should move toward 3.0
 ```
 
-## Neural networks
+## Neural networks & Deep learning
 
-There are built-in neural network constructs in Catniff as well:
+There are built-in neural network constructs in Catniff as well, from simple prebuilt nn layers:
 ```js
 const { Tensor, nn } = require("catniff");
 
@@ -102,6 +102,24 @@ gruCell.forward(b, c);
 lstmCell.forward(b, c, c);
 ```
 
+to more advanced constructs like normalization, embedding, and attention:
+```js
+// 1. Embedding: tokens -> vectors
+const embedding = new nn.Embedding(100, 64);
+const tokens = new Tensor([[1, 5, 23], [8, 2, 15]]);
+const embedded = embedding.forward(tokens);
+
+// 2. Self-Attention
+const attention = new nn.MultiheadAttention(64, 8, 0.1);
+const [output, weights] = attention.forward(embedded, embedded, embedded);
+
+// 3. Layer Normalization
+const layerNorm = new nn.LayerNorm(64);
+const normalized = layerNorm.forward(output);
+
+console.log(normalized.val());
+```
+
 And it can still do much more, check out the docs and examples below for more information.
 
 ## Documentation

package/dist/core.d.ts

@@ -184,6 +184,7 @@ export declare class Tensor {
     dropout(rate: number): Tensor;
     triu(diagonal?: number): Tensor;
     tril(diagonal?: number): Tensor;
+    maskedFill(mask: Tensor | TensorValue, value: number): Tensor;
     static full(shape: readonly number[], num: number, options?: TensorOptions): Tensor;
     static fullLike(tensor: Tensor, num: number, options?: TensorOptions): Tensor;
     static ones(shape?: readonly number[], options?: TensorOptions): Tensor;

package/dist/core.js

@@ -1465,6 +1465,11 @@ class Tensor {
         });
         return this.mul(mask);
     }
+    // Fill specific positions of this tensor with a value through a mask
+    maskedFill(mask, value) {
+        mask = this.handleOther(mask);
+        return this.mul(mask.logicalNot()).add(mask.mul(value));
+    }
     // Utility to create a new tensor filled with a number
     static full(shape, num, options = {}) {
         if (shape.length === 0)

package/dist/nn.d.ts

@@ -62,6 +62,18 @@ declare class Embedding {
     constructor(numEmbeddings: number, embeddingDim: number, device: string);
     forward(input: Tensor | TensorValue): Tensor;
 }
+declare class MultiheadAttention {
+    qProjection: Linear;
+    kProjection: Linear;
+    vProjection: Linear;
+    oProjection: Linear;
+    embedDim: number;
+    numHeads: number;
+    headDim: number;
+    dropout: number;
+    constructor(embedDim: number, numHeads: number, dropout?: number, bias?: boolean, device?: string);
+    forward(query: Tensor, key: Tensor, value: Tensor, needWeights?: boolean, attnMask?: Tensor, averageAttnWeights?: boolean): [Tensor, Tensor | undefined];
+}
 export interface StateDict {
     [key: string]: any;
 }
@@ -72,6 +84,7 @@ export declare const nn: {
     LSTMCell: typeof LSTMCell;
     LayerNorm: typeof LayerNorm;
     Embedding: typeof Embedding;
+    MultiheadAttention: typeof MultiheadAttention;
     state: {
         getParameters(model: any, visited?: WeakSet<object>): Tensor[];
         moveParameters(model: any, device: string): void;

package/dist/nn.js

@@ -197,6 +197,57 @@ class Embedding {
         return this.weight.index(input);
     }
 }
+class MultiheadAttention {
+    qProjection;
+    kProjection;
+    vProjection;
+    oProjection;
+    embedDim;
+    numHeads;
+    headDim;
+    dropout;
+    constructor(embedDim, numHeads, dropout = 0, bias = true, device) {
+        this.qProjection = new exports.nn.Linear(embedDim, embedDim, bias, device);
+        this.kProjection = new exports.nn.Linear(embedDim, embedDim, bias, device);
+        this.vProjection = new exports.nn.Linear(embedDim, embedDim, bias, device);
+        this.oProjection = new exports.nn.Linear(embedDim, embedDim, bias, device);
+        this.embedDim = embedDim;
+        this.numHeads = numHeads;
+        this.headDim = Math.floor(embedDim / numHeads);
+        this.dropout = dropout;
+    }
+    forward(query, key, value, needWeights = true, attnMask, averageAttnWeights = true) {
+        // Batch-first
+        const [batchSize, targetLen, embedDim] = query.shape;
+        const sourceLen = key.shape[1];
+        let Q = this.qProjection.forward(query); // (batchSize, targetLen, embedDim)
+        let K = this.kProjection.forward(key); // (batchSize, sourceLen, embedDim)
+        let V = this.vProjection.forward(value); // (batchSize, sourceLen, embedDim)
+        // (batchSize, numHeads, targetLen/sourceLen, headDim)
+        Q = Q.reshape([batchSize, targetLen, this.numHeads, this.headDim]).transpose(1, 2);
+        K = K.reshape([batchSize, sourceLen, this.numHeads, this.headDim]).transpose(1, 2);
+        V = V.reshape([batchSize, sourceLen, this.numHeads, this.headDim]).transpose(1, 2);
+        // Attention scores
+        let scores = Q.matmul(K.transpose(-2, -1)).div(Math.sqrt(this.headDim));
+        // Apply attention mask if specified
+        if (attnMask) {
+            scores = scores.maskedFill(attnMask, -Infinity);
+        }
+        // Calculate attention weights
+        let attnWeights = scores.softmax().dropout(this.dropout);
+        // Apply attention to values
+        let attnOutput = attnWeights.matmul(V); // (batchSize, numHeads, targetLen, headDim)
+        // (batchSize, targetLen, embedDim)
+        attnOutput = attnOutput.transpose(1, 2).reshape([batchSize, targetLen, embedDim]);
+        // Output
+        const output = this.oProjection.forward(attnOutput);
+        // Average weights if needed
+        if (averageAttnWeights) {
+            attnWeights = attnWeights.mean(1);
+        }
+        return [output, needWeights ? attnWeights : undefined];
+    }
+}
 const state = {
     getParameters(model, visited = new WeakSet()) {
         if (visited.has(model))
@@ -266,5 +317,6 @@ exports.nn = {
     LSTMCell,
     LayerNorm,
     Embedding,
+    MultiheadAttention,
     state
 };

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "catniff",
-  "version": "0.6.4",
+  "version": "0.6.5",
   "description": "A small Torch-like deep learning framework for Javascript",
   "main": "index.js",
   "scripts": {