Implement Backpropagtion For Head Class

[damoebe]

The implemented backPropagate method in Head.java currently does no have any content.
To know how to pass the delta gradient through a head, partial derive the output function with respect to the head intern weights. Then implement those functions with static Matrix.java calc methods. Also return the new gradient (
derivation with respect to the input).

[damoebe]

Consider AI generated solution (check for logic errors and simplify)

    /**
     * BACKPROPAGATION IMPLEMENTATION
     * 
     * Mathematical Derivation:
     * 
     * Forward Pass:
     * 1. Q = X * W_q + b_q (queries)
     * 2. K = X * W_k + b_k (keys)
     * 3. V = X * W_v + b_v (values)
     * 4. scores = Q * K^T / sqrt(d_k) (scaled dot product)
     * 5. attention = softmax(scores) (attention weights)
     * 6. output = attention * V (head output)
     * 
     * Backward Pass (Chain Rule):
     * Given dL/dOutput from the next layer:
     * 
     * 1. dL/dAttention = dL/dOutput * V^T
     * 2. dL/dScores = softmax_derivative(scores) * dL/dAttention
     * 3. dL/dQ = dL/dScores * K / sqrt(d_k)
     * 4. dL/dK = dL/dScores^T * Q / sqrt(d_k)
     * 5. dL/dV = dL/dOutput^T * attention
     * 6. dL/dW_q = X^T * dL/dQ
     * 7. dL/dW_k = X^T * dL/dK
     * 8. dL/dW_v = X^T * dL/dV
     * 9. dL/dInput = dL/dQ * W_q^T + dL/dK * W_k^T + dL/dV * W_v^T
     * 
     * @param deltas The delta from the mha of this head (dL/dOutput)
     * @return The new deltas for the next backpropagation step (dL/dInput)
     */
    @Override
    public double[][] backPropagate(double[][] deltas){
        if (deltas == null || inputSequence == null) {
            return null;
        }

        try {
            double[][] inputMatrix = inputSequence.getVerticalMatrix();
            double scaleFactor = Math.sqrt(inputEmbeddingSize);

            // Step 1: Gradient w.r.t attention output
            // dL/dAttention = dL/dOutput * V^T
            double[][] dL_dAttention = Matrix.multiply(deltas, Matrix.transpose(values));

            // Step 2: Gradient w.r.t attention scores (softmax derivative)
            // dL/dScores = softmax_derivative(scores) * dL/dAttention
            double[][] dL_dScores = computeSoftmaxGradient(attention, dL_dAttention);

            // Step 3: Gradient w.r.t queries (scaled by sqrt(d_k))
            // dL/dQ = dL/dScores * K / sqrt(d_k)
            double[][] dL_dQ = Matrix.multiply(dL_dScores, keys);
            scaleMatrix(dL_dQ, 1.0 / scaleFactor);

            // Step 4: Gradient w.r.t keys (scaled by sqrt(d_k))
            // dL/dK = dL/dScores^T * Q / sqrt(d_k)
            double[][] dL_dScoresTransposed = Matrix.transpose(dL_dScores);
            double[][] dL_dK = Matrix.multiply(dL_dScoresTransposed, queries);
            scaleMatrix(dL_dK, 1.0 / scaleFactor);

            // Step 5: Gradient w.r.t values
            // dL/dV = dL/dOutput^T * attention
            double[][] dL_dV = Matrix.multiply(Matrix.transpose(deltas), attention);

            // Step 6: Compute weight gradients
            // dL/dW_q = X^T * dL/dQ
            double[][] dW_q = Matrix.multiply(Matrix.transpose(inputMatrix), dL_dQ);
            // dL/dW_k = X^T * dL/dK
            double[][] dW_k = Matrix.multiply(Matrix.transpose(inputMatrix), dL_dK);
            // dL/dW_v = X^T * dL/dV
            double[][] dW_v = Matrix.multiply(Matrix.transpose(inputMatrix), dL_dV);

            // Step 7: Update weights using gradient descent
            updateWeights(0, dW_q); // Query weights
            updateWeights(1, dW_k); // Key weights
            updateWeights(2, dW_v); // Value weights

            // Step 8: Gradient w.r.t input for previous layer
            // dL/dInput = dL/dQ * W_q^T + dL/dK * W_k^T + dL/dV * W_v^T
            double[][] dL_dInput_q = Matrix.multiply(dL_dQ, Matrix.transpose(weights.get(0)));
            double[][] dL_dInput_k = Matrix.multiply(dL_dK, Matrix.transpose(weights.get(1)));
            double[][] dL_dInput_v = Matrix.multiply(dL_dV, Matrix.transpose(weights.get(2)));

            // Sum all gradients w.r.t input
            double[][] dL_dInput = addMatrices(dL_dInput_q, addMatrices(dL_dInput_k, dL_dInput_v));

            return dL_dInput;

        } catch (Exception e) {
            System.err.println("Error during backpropagation: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }

    /**
     * Computes the softmax gradient
     * softmax_derivative(z) = diag(softmax(z)) - softmax(z) * softmax(z)^T
     * Then multiply by the upstream gradient
     */
    private double[][] computeSoftmaxGradient(double[][] attention, double[][] dL_dAttention) {
        int batchSize = attention.length;
        int seqLength = attention[0].length;
        double[][] dL_dScores = new double[batchSize][seqLength];

        for (int i = 0; i < batchSize; i++) {
            double[] softmax_i = attention[i];
            double[] upstream_i = dL_dAttention[i];

            // Jacobian of softmax: diag(s) - s * s^T
            for (int j = 0; j < seqLength; j++) {
                dL_dScores[i][j] = softmax_i[j] * upstream_i[j];
                for (int k = 0; k < seqLength; k++) {
                    dL_dScores[i][j] -= softmax_i[j] * softmax_i[k] * upstream_i[k];
                }
            }
        }
        return dL_dScores;
    }

    /**
     * Updates weights using gradient descent: W = W - learning_rate * dW
     */
    private void updateWeights(int weightIndex, double[][] weightGradient) {
        double[][] currentWeights = weights.get(weightIndex);
        for (int i = 0; i < currentWeights.length; i++) {
            for (int j = 0; j < currentWeights[i].length; j++) {
                currentWeights[i][j] -= learningRate * weightGradient[i][j];
            }
        }
    }

    /**
     * Scales all elements in a matrix by a scalar value
     */
    private void scaleMatrix(double[][] matrix, double scale) {
        for (int i = 0; i < matrix.length; i++) {
            for (int j = 0; j < matrix[i].length; j++) {
                matrix[i][j] *= scale;
            }
        }
    }