Add support to dropout.

animated-transformer/src/app/animated-transformer/model-selector/model-selector.component.ts

@@ -80,14 +80,14 @@ export class ModelSpecAndData {
 const layerSpec: TransformerParamLayerSpec = {
   nHeads: 4,
   hasPosEncoding: true,
-  computeSpec: { residuals: true },
+  computeSpec: { residuals: true, dropoutRate: 0.0 },
   layerNormFF: false,
   layerNormHeadsProjection: false,
   addLayerNormBias: false,
 };
 
 const defaultConfig: ModelConfig = {
-  name: 'd=8 l=1 h=4, !layerN',
+  name: 'd=8 l=1 h=4, !layerN !dropout',
   transformer: {
     spec: {
       inputRep: 8,
@@ -105,14 +105,14 @@ const defaultConfig: ModelConfig = {
 const layerSpecWithNorm: TransformerParamLayerSpec = {
   nHeads: 4,
   hasPosEncoding: true,
-  computeSpec: { residuals: true },
+  computeSpec: { residuals: true, dropoutRate: 0.0 },
   layerNormFF: true,
   layerNormHeadsProjection: true,
   addLayerNormBias: false,
 };
 
 const transWithLayerNormed: ModelConfig = {
-  name: 'd=8 l=1 h=4 +layerN',
+  name: 'd=8 l=1 h=4 +layerN !dropout',
   transformer: {
     spec: {
       inputRep: 8,
@@ -127,15 +127,42 @@ const transWithLayerNormed: ModelConfig = {
   },
 };
 
+const layerSpecWithNormAndDropout: TransformerParamLayerSpec = {
+  nHeads: 4,
+  hasPosEncoding: true,
+  computeSpec: { residuals: true, dropoutRate: 0.1 },
+  layerNormFF: true,
+  layerNormHeadsProjection: true,
+  addLayerNormBias: false,
+};
+
+const transWithLayerNormedAndDropout: ModelConfig = {
+  name: 'd=8 l=1 h=4 +layerN +dropout',
+  transformer: {
+    spec: {
+      inputRep: 8,
+      kqvRep: 8,
+      layers: [layerSpecWithNormAndDropout],
+    },
+    init: {
+      stddev: 0.5,
+      mean: 0,
+      seed: 96,
+    },
+  },
+};
+
 const simpleTransformer = new ModelSpecAndData('transformer', defaultConfig);
 
 const simpleTransformerWithLayerNorm = new ModelSpecAndData('transformer', transWithLayerNormed);
 
+const simpleTransformerWithLayerNormAndDropout = new ModelSpecAndData('transformer', transWithLayerNormedAndDropout)
+
 export interface ModelUpdate {
   model: ModelSpecAndData | null;
 }
 
-const initModels: ModelSpecAndData[] = [simpleTransformer, simpleTransformerWithLayerNorm];
+const initModels: ModelSpecAndData[] = [simpleTransformer, simpleTransformerWithLayerNorm, simpleTransformerWithLayerNormAndDropout];
 const initModelsMap: { [name: string]: ModelSpecAndData } = {};
 initModels.forEach((m) => (initModelsMap[m.config.name] = m));
 

animated-transformer/src/lib/trainer/basic_transformer_trainer.spec.ts

@@ -31,7 +31,70 @@ describe('basic_transformer_trainer', () => {
     const layerSpec: transformer.TransformerParamLayerSpec = {
       nHeads: 1,
       hasPosEncoding: false,
-      computeSpec: { residuals: true },
+      computeSpec: { residuals: true, dropoutRate: 0},
+      // TODO: investigate: these make 0 gradients?
+      layerNormFF: false,
+      layerNormHeadsProjection: false,
+      addLayerNormBias: false,
+    };
+    const decoderConfig: transformer.TransformerConfig = {
+      spec: {
+        inputRep: 4,
+        kqvRep: 3,
+        layers: [layerSpec, layerSpec],
+      },
+      init: {
+        stddev: 0.5,
+        mean: 0,
+        seed: 1,
+      },
+    };
+    const taskConfig: BasicRandSeededTaskConfig = {
+      name: 'AorBisMaxTask',
+      maxInputLen: 4,
+      maxOutputLen: 4,
+      seed: 0,
+    };
+    const trainStateConfig: TrainStateConfig = {
+      learningRate: 0.5,
+      batchSize: 64,
+      maxInputlength: taskConfig.maxInputLen,
+      testSetSize: 0,
+      trainSetSize: 64,
+    };
+    const task = new abtask.AorBisMaxTask(taskConfig);
+    const tokenRep = prepareBasicTaskTokenRep(task.baseVocab);
+    const initParams = transformer.initDecoderParamsTree(tokenRep, decoderConfig);
+    console.log('initTransformerTrainState...');
+    const trainState = initTransformerTrainState(
+      task,
+      tokenRep,
+      strSeqPrepFn,
+      singleNextTokenIdxOutputPrepFn,
+      decoderConfig,
+      initParams,
+      trainStateConfig
+    );
+    // Taking a couple of steps...
+    const initLoss = trainState.batchMeanLoss;
+    expect(trainState.nSteps).toBe(0);
+    expect(trainState.nExamples).toBe(0);
+    const stillTraining = trySgdTrainStep(trainState);
+    expect(stillTraining).toBe(true);
+    const newLoss = trainState.batchMeanLoss;
+    expect(trainState.nSteps).toBe(1);
+    expect(trainState.nExamples).toBe(trainState.batchExamples.length);
+    expect(newLoss).toBeLessThan(initLoss);
+
+    // Memory cleanup
+    jstree.forEach((g: GTensor<any>) => g.dispose(), initParams);
+    trainState.dispose();
+  });
+  it('AorBisMaxTaskWithDropout training', async () => {
+    const layerSpec: transformer.TransformerParamLayerSpec = {
+      nHeads: 1,
+      hasPosEncoding: true,
+      computeSpec: { residuals: true, dropoutRate: 0.1 },
       // TODO: investigate: these make 0 gradients?
       layerNormFF: false,
       layerNormHeadsProjection: false,

animated-transformer/src/lib/trainer/basic_transformer_trainer.ts

@@ -109,6 +109,8 @@ export function initTransformerTrainState(
 }
 
 export function computeMetrics(state: TransformerTrainState): TrainMetrics {
+  // Eval mode is passed inside computeStateBatchAccuracy and computeLossAndAccuracy to
+  // computeTransformer.
   const trainBatchAcc: number = computeStateBatchAccuracy(state);
   const testLossAndAcc = computeLossAndAccuracy(state, state.taskSplit.testSetExamples);
   return {
@@ -128,7 +130,8 @@ export function computeStateBatchAccuracy(state: TransformerTrainState): number
     const decoderComputation = transformer.computeTransformer(
       state.spec,
       state.params,
-      state.inputsVar
+      state.inputsVar,
+      true
     );
     meanAcc = transformerAccuracy(
       decoderComputation,
@@ -156,7 +159,8 @@ export function computeLossAndAccuracy(
       const decoderComputation = transformer.computeTransformer(
         state.spec,
         state.params,
-        state.inputsVar
+        state.inputsVar,
+        true
       );
       const batchAcc = transformerAccuracy(
         decoderComputation,

animated-transformer/src/lib/transformer/dropout.spec.ts

@@ -0,0 +1,74 @@
+/* Copyright 2023 Google LLC. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+/* Dropout */
+import * as tf from '@tensorflow/tfjs';
+import { gtensor } from '..';
+import { dropout } from './dropout';
+
+describe('dropout', () => {
+
+  it('Basic dropout', () => {
+    const beforeDropout =
+      gtensor.makeRange('input', 1, 5, 1, 'float32');
+    const afterDropout = dropout(0.5, beforeDropout, false, 0);
+    afterDropout.tensor.print();
+
+    tf.test_util.expectArraysClose(afterDropout.tensor.dataSync(),
+      [2, 4, 0, 8]);
+    expect(afterDropout.dimNames).toEqual(
+      ['input']);
+  });
+
+  it('Deterministic output', () => {
+    const beforeDropout =
+      gtensor.makeRange('input', 1, 5, 1, 'float32');
+    const afterDropout = dropout(0.5, beforeDropout, true, 0);
+    afterDropout.tensor.print();
+
+    tf.test_util.expectArraysClose(afterDropout.tensor.dataSync(),
+      [1, 2, 3, 4]);
+    expect(afterDropout.dimNames).toEqual(
+      ['input']);
+  });
+
+  it('Dropout with noise shape ', () => {
+    const beforeDropout = new gtensor.GTensor(
+        tf.tensor([
+          [
+            [1, 2, 3, 4],
+            [5, 6, 7, 8],
+            [9, 10, 11, 12],
+          ]
+        ]),
+        ['batch', 'pos', 'inputRep']
+      );
+    const afterDropout = dropout(0.5, beforeDropout, false, 1, ['pos']);
+    afterDropout.tensor.print();
+
+    tf.test_util.expectArraysClose(afterDropout.tensor.dataSync(),
+    [
+        [
+          [2, 4, 6, 8],
+          [0, 0, 0, 0],
+          [0, 0, 0, 0],
+        ]
+      ]);
+    expect(afterDropout.dimNames).toEqual(
+      ['batch', 'pos', 'inputRep']);
+  });
+
+  // TODO(@aliciafmachado): Test that grads are not applied on de-activated neurons.
+});

animated-transformer/src/lib/transformer/dropout.ts

@@ -0,0 +1,48 @@
+/* Copyright 2023 Google LLC. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+import {
+    GTensor,
+  } from '../gtensor/gtensor';
+
+  import {dropout as tf_dropout} from '@tensorflow/tfjs';
+
+  // Wrapper for tf ts dropout.
+  export function dropout<G extends string, D extends G>(
+    dropoutRate: number,
+    g: GTensor<G>,
+    deterministic: boolean,
+    seed?: number,
+    dimNames?: string[],
+  ): GTensor<G> {
+    if (deterministic) {
+      return g;
+    }
+
+    let dimensions: number[] = g.tensor.shape;
+    if (dimNames) {
+      dimensions = [];
+      for (const d of g.dimNames) {
+        if (dimNames.includes(d)) {
+          dimensions = dimensions.concat(g.dim[d].size);
+        }
+        else {
+          dimensions = dimensions.concat(1);
+        }
+      }
+    }
+    return new GTensor(tf_dropout(g.tensor, dropoutRate, dimensions, seed), g.dimNames);
+  }
+

animated-transformer/src/lib/transformer/transformer_gtensor.spec.ts

@@ -24,6 +24,7 @@ describe('GTensor Transformers', () => {
   it('basic transformer shapes', () => {
     const spec: AttnHeadComputeSpec = {
       residuals: true,
+      dropoutRate: 0.0
     };
     const paramSizes: AttnHeadParamSpec = {
       inputRep: 2,