Introduction to Functional Automatic Differentiation

Automatic differentiation is a core technique in neural network training that enables efficient computation of gradients for optimization. MindSpore implements a functional approach to automatic differentiation through its grad and value_and_grad interfaces, which provide mathematical semantics for gradient computation.

Core Components and Setup

To begin working with automatic differentiation in MindSpore, we first import the necessary modules:

import numpy as np
import mindspore
from mindspore import nn, ops, Tensor, Parameter

These imports provide access to numerical operations, neural network components, tensor data structures, and trainable parameters.

Defining Inputs and Parameters

We create sample input data and model parameters for a simple linear transformation:

input_data = ops.ones(5, mindspore.float32)
target_output = ops.zeros(3, mindspore.float32)
weight_matrix = Parameter(Tensor(np.random.randn(5, 3), mindspore.float32), name='weights')
bias_vector = Parameter(Tensor(np.random.randn(3,), mindspore.float32), name='bias')

Computational Function Definition

We define a function that performs the forward computation and loss calculation:

def compute_loss(input_tensor, target, weights, bias):
    linear_output = ops.matmul(input_tensor, weights) + bias
    loss_value = ops.binary_cross_entropy_with_logits(
        linear_output, target, 
        ops.ones_like(linear_output), 
        ops.ones_like(linear_output)
    )
    return loss_value

This function computes a linear transformation followed by a binary cross-entropy loss.

Gradient Computation

To compute gradients with respect to model parameters, we use MindSpore's automatic differentiation:

gradient_function = mindspore.grad(compute_loss, (2, 3))
parameter_gradients = gradient_function(input_data, target_output, weight_matrix, bias_vector)
print(parameter_gradients)

Handling Multiple Outputs

When functions return multiple outputs, we can control gradient computation using stop gradient operations:

def compute_with_intermediate(input_tensor, target, weights, bias):
    intermediate = ops.matmul(input_tensor, weights) + bias
    loss_val = ops.binary_cross_entropy_with_logits(
        intermediate, target, 
        ops.ones_like(intermediate), 
        ops.ones_like(intermediate)
    )
    return loss_val, ops.stop_gradient(intermediate)

grad_func = mindspore.grad(compute_with_intermediate, (2, 3))
grad_results = grad_func(input_data, target_output, weight_matrix, bias_vector)

Neural Network Implementation

For a more structured approach, we can implement the computation as a neural network cell:

class LinearModel(nn.Cell):
    def __init__(self, weight_param, bias_param):
        super().__init__()
        self.weights = weight_param
        self.bias = bias_param
        
    def construct(self, x):
        return ops.matmul(x, self.weights) + self.bias

model = LinearModel(weight_matrix, bias_vector)
loss_function = nn.BCEWithLogitsLoss()

def forward_computation(x, y):
    predictions = model(x)
    return loss_function(predictions, y)

grad_compute = mindspore.value_and_grad(
    forward_computation, 
    None, 
    weights=model.trainable_params()
)
loss_result, gradients = grad_compute(input_data, target_output)