Understanding Transformer Architecture

// Self-Attention implementation in JavaScript
class SelfAttention {
  constructor(embedSize, heads) {
    this.embedSize = embedSize;
    this.heads = heads;
    this.headDim = Math.floor(embedSize / heads);

    if (this.headDim * heads !== embedSize) {
      throw new Error("Embed size must be divisible by heads");
    }

    // Initialize weight matrices (simplified for demonstration)
    this.wQ = this.initializeWeights(embedSize, embedSize);
    this.wK = this.initializeWeights(embedSize, embedSize);
    this.wV = this.initializeWeights(embedSize, embedSize);
    this.wO = this.initializeWeights(embedSize, embedSize);
  }

  // Helper to initialize weights (random for demonstration)
  initializeWeights(inputDim, outputDim) {
    const weights = [];
    for (let i = 0; i < inputDim; i++) {
      const row = [];
      for (let j = 0; j < outputDim; j++) {
        row.push((Math.random() - 0.5) * 0.1); // Small random values
      }
      weights.push(row);
    }
    return weights;
  }

  // Matrix multiplication
  matMul(a, b) {
    const result = [];
    for (let i = 0; i < a.length; i++) {
      const row = [];
      for (let j = 0; j < b[0].length; j++) {
        let sum = 0;
        for (let k = 0; k < a[0].length; k++) {
          sum += a[i][k] * b[k][j];
        }
        row.push(sum);
      }
      result.push(row);
    }
    return result;
  }

  // Softmax function
  softmax(arr) {
    const maxVal = Math.max(...arr);
    const expValues = arr.map(val => Math.exp(val - maxVal));
    const sumExp = expValues.reduce((acc, val) => acc + val, 0);
    return expValues.map(val => val / sumExp);
  }

  // Self-attention calculation
  forward(inputs, mask = null) {
    const batchSize = inputs.length;
    const seqLength = inputs[0].length;

    // Linear projections
    const Q = this.matMul(inputs, this.wQ);
    const K = this.matMul(inputs, this.wK);
    const V = this.matMul(inputs, this.wV);

    // Split into heads (simplified, not showing tensor reshaping)

    // Calculate attention scores
    const scores = this.matMul(Q, this.transpose(K));

    // Scale scores
    const scaledScores = scores.map(row =>
      row.map(val => val / Math.sqrt(this.headDim))
    );

    // Apply mask if provided
    if (mask) {
      // Apply mask logic here
    }

    // Apply softmax to get attention weights
    const attention = scaledScores.map(row => this.softmax(row));

    // Calculate weighted values
    const output = this.matMul(attention, V);

    // Final linear projection
    return this.matMul(output, this.wO);
  }

  // Helper to transpose 2D array
  transpose(matrix) {
    return matrix[0].map((_, colIndex) =>
      matrix.map(row => row[colIndex])
    );
  }
}