Vectorized differentiation

Basic operations

It is technically possible to implement a PINN using the basic automatic differentiation previously described. However, it is painfully inefficient to a point that it makes testing simple examples very unpractical. The next simple yet very useful upgrade is, instead of using an auto-differentiation at the scalar level, to vectorize the operations. It means that addition will be allowed between 1d or 2d arrays, element-wise multiplication also, etc. A new operation will be required to represent what it going on between FNN layers: the dot product and the activation. Finally, some helper operations are also required: transpose, split a 1d array into scalar nodes, etc

There is no new concept there, the trick is now to apply the same strategy to vector or matrix operations. Since these operations can be reduced to multiple scalar operations, it should be achievable. Let us slightly modify the Node class to hold an array instead of a scalar.

src
 1@dataclass
 2class NodeVectorized:
 3    """an instance of NodeVectorized can be used in a
 4    unary or binary forward operation. This operation
 5    creates a result node. The result node points to
 6    its parent(s) building a computation graph.
 7
 8    The val attribute is a 1d or 2d array. All the operations
 9    are vectorized.
10
11    For debugging purposes, the name attribute represents
12    the operand(s) and the operation in string. The operation
13    only the operation performed.
14
15    The is_constant indicates if a the node should be considered
16    as a constant or not.
17
18    The index is unique and is used as a hash.
19
20    The backward callbacks which takes a node and derivatives dict
21    computes the derivatives using a reverse sweep and builds
22    the computation graph of these derivatives.
23
24    The derivatives dict[NodeVectorized, NodeVectorized] is
25    a default dict to a zero node, making the cumulate operation
26    valid and returning a correct zero derivative for a node
27    which is not connected as a parent and not present in the dict."""
28
29    val: Array2D | Vector
30    parents: tuple[Self, ...] = field(default_factory=tuple)
31    name: str = ""
32    operation: str = ""
33    is_constant: bool = False
34
35    index: int = field(init=False)
36
37    _counter: ClassVar[int] = 0
38
39    def __post_init__(self) -> None:
40        self.backward: Backward = lambda _, __: None
41        self.index = self._counter
42        self._increment_counter()
43
44    def __hash__(self) -> int:
45        return self.index
46
47    @classmethod
48    def _increment_counter(cls):
49        cls._counter += 1

The differentiate hardly changes too:

src
 1def topological_sort(node: NodeVectorized) -> tuple[NodeVectorized, ...]:
 2    """here the topological order is intended from the leaf node up to the
 3    root nodes.
 4
 5    Each node points towards its parent(s). It means that starting from a leaf
 6    node, a DFS traversal will go all the way up to the last parents (root nodes).
 7
 8    N0 <- N1 <- N2 <- N3
 9           ^          |
10           |          |
11           ------------
12
13    Since the yield statement is located AFTER the recursive call,
14    the nodes would be yielded in the reverse topological order.
15
16    Assuming the first parent of N3 is N1:
17
18    N4, N1, N0
19    yields: N0, N1
20    N2, (N1 already visited)
21    yields: N2, N3
22
23    => N0, N1, N2,  N3
24
25    That is why it is reversed before return.
26    """
27    visited: set[int] = set()
28
29    def _dfs(current: NodeVectorized) -> Generator[NodeVectorized]:
30        if current.index in visited:
31            return
32        visited.add(current.index)
33        for parent in current.parents:
34            yield from _dfs(parent)
35        yield current
36
37    return tuple(reversed(tuple(_dfs(node))))
src
 1def differentiate(node: NodeVectorized) -> Derivatives:
 2    """after getting the topological order,
 3    computes the numerical values of every intermediate
 4    derivatives by calling the backward callbacks of every node.
 5
 6    However, the difference with the first order differentiation is
 7    that the Derivatives dict is [NodeVectorized, NodeVectorized],
 8    building a computation graph during a reverse sweep.
 9
10    If the node is declared constant, all the derivatives are zero.
11    Using a defaultdict derivatives allows to return zero node if a key is not present.
12
13    It serves two purposes:
14    - first the cumulate operation in the backward callbacks is automatically valid
15    - if a node is not key, it means it is not part of the parent tree,
16      so the derivative should be zero node"""
17    order: tuple[NodeVectorized, ...] = topological_sort(node)
18    derivatives: dict[NodeVectorized, NodeVectorized] = defaultdict(lambda: zero((1,)))
19
20    if node.is_constant:
21        return derivatives
22
23    derivatives[node] = one(node.val.shape)
24    for current in order:
25        current.backward(current, derivatives)
26    return derivatives
src
1def one(shape: tuple[int, ...]) -> NodeVectorized:
2    return NodeVectorized(
3        np.ones(shape, dtype=FLOAT),
4        name="1",
5        is_constant=True,
6    )
src
1def zero(shape: tuple[int, ...]) -> NodeVectorized:
2    return NodeVectorized(
3        np.zeros(shape, dtype=FLOAT),
4        name="0",
5        is_constant=True,
6    )

Exercise (easy)

  1. Implement and test the vectorized addition and element-wise multiplication.

  2. Implement and test the vectorized logistic activation. (The vectorized activation is rather straighforward. This is just a unary function, just applied element-wise.)

Correction
src
 1def addition_forward(first: NodeVectorized, second: NodeVectorized) -> NodeVectorized:
 2    result = NodeVectorized(
 3        first.val + second.val,
 4        (first, second),
 5        f"({first.name} + {second.name})",
 6        "+",
 7        first.is_constant and second.is_constant,
 8    )
 9    result.backward = addition_backward
10    return result
src
 1def multiplication_forward(
 2    first_vector: NodeVectorized, second_vector: NodeVectorized
 3) -> NodeVectorized:
 4    if first_vector.val.shape != second_vector.val.shape:
 5        raise IncompatibleDimensionsAutoDiff()
 6
 7    result = NodeVectorized(
 8        first_vector.val * second_vector.val,
 9        (first_vector, second_vector),
10        f"({first_vector.name} * {second_vector.name})",
11        "*",
12        first_vector.is_constant and second_vector.is_constant,
13    )
14    result.backward = multiplication_backward
15    return result
tests
 1def assert_derivatives(
 2    x_z_zv_dzv_d2zv: tuple[NodeVectorized, NodeVectorized, float, float, float],
 3) -> None:
 4    x, z, z_val, dz_val, d2z_val = x_z_zv_dzv_d2zv
 5    print(f"\n x={x.val} f(x)={z.name}")
 6
 7    assert pytest.approx(z.val) == z_val
 8
 9    # ic(z)
10    # breakpoint()
11
12    dz: Derivatives = differentiate(z)
13    dzdx: NodeVectorized = dz[x]
14
15    # ic(dz)
16    # breakpoint()
17
18    assert pytest.approx(dzdx.val) == dz_val
19
20    d_dzdx: Derivatives = differentiate(dzdx)
21    d2zdx2: NodeVectorized = d_dzdx[x]
22
23    assert pytest.approx(d2zdx2.val) == d2z_val
tests
1@pytest.fixture
2def x_val() -> Vector:
3    return np.array([-1.0, 0.1, 2.0], dtype=FLOAT)
tests
 1@pytest.fixture(params=list(range(2, 6)))
 2def x_pow_one_by_one(
 3    request, x_val
 4) -> tuple[NodeVectorized, NodeVectorized, Vector, Vector, Vector]:
 5    max_power: int = request.param
 6    x = NodeVectorized(x_val, name="x")
 7    z = x
 8    for _ in range(max_power - 1):
 9        z = multiplication_forward(z, x)
10    return (
11        x,
12        z,
13        x_val**max_power,
14        max_power * x_val ** (max_power - 1),
15        max_power * (max_power - 1) * x_val ** (max_power - 2),
16    )
tests
 1@pytest.fixture(params=list(range(1, 6)))
 2def x_sum_one_by_one(
 3    request, x_val
 4) -> tuple[NodeVectorized, NodeVectorized, Vector, Vector, Vector]:
 5    max_repet: int = request.param
 6    x = NodeVectorized(x_val, name="x")
 7    z = x
 8    for _ in range(max_repet - 1):
 9        z = addition_forward(z, x)
10    return (
11        x,
12        z,
13        max_repet * x_val,
14        max_repet * np.ones(x_val.shape),
15        np.zeros(1),
16    )
tests
1def test_derivatives_sum_one_by_one(x_sum_one_by_one):
2    assert_derivatives(x_sum_one_by_one)
tests
1def test_derivatives_pow_one_by_one(x_pow_one_by_one):
2    assert_derivatives(x_pow_one_by_one)
src
 1def division_forward(num: NodeVectorized, den: NodeVectorized) -> NodeVectorized:
 2    result = NodeVectorized(
 3        num.val / den.val,
 4        (num, den),
 5        f"({num.name} / {den.name})",
 6        "/",
 7        num.is_constant and den.is_constant,
 8    )
 9    result.backward = division_backward
10    return result
src
 1def exp_forward(node: NodeVectorized) -> NodeVectorized:
 2    result = NodeVectorized(
 3        np.exp(node.val),
 4        (node,),
 5        f"exp({node.name})",
 6        "exp",
 7        node.is_constant,
 8    )
 9    result.backward = exp_backward
10    return result
src
1def minus_one(shape: tuple[int, ...]) -> NodeVectorized:
2    return NodeVectorized(
3        -np.ones(shape, dtype=FLOAT),
4        name="-1",
5        is_constant=True,
6    )
src
 1def logistic_forward(node: NodeVectorized) -> NodeVectorized:
 2    one_ = one(node.val.shape)
 3    minus_one_ = minus_one(node.val.shape)
 4    return division_forward(
 5        one_,
 6        addition_forward(
 7            one_,
 8            exp_forward(multiplication_forward(minus_one_, node)),
 9        ),
10    )
src
1def exp_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
2    if not node.parents[0].is_constant:
3        derivatives[node.parents[0]] = addition_forward(
4            multiplication_forward(exp_forward(node.parents[0]), derivatives[node]),
5            derivatives[node.parents[0]],
6        )
src
 1def logistic_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
 2    if not node.parents[0].is_constant:
 3        one_ = one(node.val.shape)
 4        minus_one_ = minus_one(node.val.shape)
 5        temp: NodeVectorized = logistic_forward(node)
 6        logistic_der: NodeVectorized = multiplication_forward(
 7            temp,
 8            addition_forward(
 9                one_,
10                multiplication_forward(minus_one_, temp),
11            ),
12        )
13        derivatives[node.parents[0]] = addition_forward(
14            multiplication_forward(
15                logistic_der,
16                derivatives[node],
17            ),
18            derivatives[node.parents[0]],
19        )
tests
 1@pytest.fixture(params=list(range(1, 4)))
 2def x_inv_pow_one_by_one(
 3    request, x_val
 4) -> tuple[NodeVectorized, NodeVectorized, Vector, Vector, Vector]:
 5    max_power: int = request.param
 6    x = NodeVectorized(x_val, name="x")
 7    z = NodeVectorized(np.ones(x_val.shape), name="1", is_constant=True)
 8    for _ in range(max_power):
 9        z = division_forward(z, x)
10    return (
11        x,
12        z,
13        x_val ** (-max_power),
14        -max_power / x_val ** (max_power + 1),
15        max_power * (max_power + 1) / x_val ** (max_power + 2),
16    )
tests
1def test_derivatives_inv_one_by_one(x_inv_pow_one_by_one):
2    assert_derivatives(x_inv_pow_one_by_one)
tests
 1@pytest.fixture
 2def x_exp_exp(x_val) -> tuple[NodeVectorized, NodeVectorized, Vector, Vector, Vector]:
 3    x = NodeVectorized(x_val, name="x")
 4    z = exp_forward(exp_forward(x))
 5    return (
 6        x,
 7        z,
 8        np.exp(np.exp(x_val)),
 9        np.exp(np.exp(x_val) + x_val),
10        (np.exp(x_val) + 1) * np.exp(np.exp(x_val) + x_val),
11    )
tests
1def test_derivatives_exp_exp(x_exp_exp):
2    assert_derivatives(x_exp_exp)
tests
 1@pytest.fixture
 2def x_logistic(x_val) -> tuple[NodeVectorized, NodeVectorized, float, float, float]:
 3    x = NodeVectorized(x_val, name="x")
 4    z = logistic_forward(x)
 5    temp: float = 1.0 / (1.0 + np.exp(-x_val))
 6    return (
 7        x,
 8        z,
 9        temp,
10        temp * (1.0 - temp),
11        temp * (1.0 - temp) * (1.0 - 2.0 * temp),
12    )
tests
1def test_derivatives_logistic(x_logistic):
2    assert_derivatives(x_logistic)

Dot product

The dot product between a \((n\times m)\) matrix \(W\) (weights) and a \((m)\) vector \(x\) (values of the previous layer) to obtain the \((n)\) result vector \(y\) is more interesting. The goal is to obtain the vectorized derivatives of the result vector with respect to the vector \(x\) and the weights \(W\). But what does it mean to differentiate with respect to a vector or a matrix? In this context, it just means that any weight or vector coefficient must have its derivative calculated the same way as if only scalar operations were used.

The easiest way is to write what the operation would be on a small example. To simplify things, let us focus on a two step case, first a dot product, then a function \(S\) which sums up all the coefficients of a vector.

dot_product_vectorized

Dot product vectorized

(1)\[\begin{split}\begin{bmatrix} w_{00} & w_{01} & w_{02} \\ w_{10} & w_{11} & w_{12} \end{bmatrix} \cdot \begin{Bmatrix} x_0 \\ x_1 \end{Bmatrix} &= \begin{Bmatrix} y_0 \\ y_1 \end{Bmatrix} = \begin{Bmatrix} w_{00} x_0 + w_{01} x_1 + w_{02} x_2 \\ w_{10} x_0 + w_{11} x_1 + w_{12} x_2 \end{Bmatrix}\\ z&=y_0+y_1\end{split}\]

If a non vectorized auto-differentiation was used:

  • \(x_0\) would see a scalar multiplication with \(w_{00}\) and \(w_{10}\).

  • \(x_1\) would see a scalar multiplication with \(w_{01}\) and \(w_{11}\).

  • \(x_2\) would see a scalar multiplication with \(w_{02}\) and \(w_{12}\).

The backward sweep is relatively straighforward:

(2)\[\begin{split}\nabla_y \, z = \begin{Bmatrix} \cfrac{\partial z}{\partial y_0} \\ \cfrac{\partial z}{\partial y_1} \end{Bmatrix}\end{split}\]

A chain rule using what we know for a product between scalars gives for the derivatives with respect to \(x\) components:

(3)\[\begin{split}\begin{Bmatrix} w_{00} \cfrac{\partial z}{\partial y_0} + w_{10} \cfrac{\partial z}{\partial y_1} \\ w_{01} \cfrac{\partial z}{\partial y_0} + w_{11} \cfrac{\partial z}{\partial y_1} \\ w_{02} \cfrac{\partial z}{\partial y_0} + w_{12} \cfrac{\partial z}{\partial y_1} \end{Bmatrix} = W^\top \cdot \nabla_y \, z\end{split}\]

The second chain rule gives for the derivatives with respect to \(W\) components:

(4)\[\begin{split}\begin{bmatrix} \cfrac{\partial z}{\partial y_0} x_0 & \cfrac{\partial z}{\partial y_0} x_1 & \cfrac{\partial z}{\partial y_0} x_2 \\ \cfrac{\partial z}{\partial y_1} x_0 & \cfrac{\partial z}{\partial y_1} x_1 & \cfrac{\partial z}{\partial y_1} x_2 \\ \end{bmatrix} = \nabla_y \, z \cdot x^\top\end{split}\]

For the specific case of \(z=y_0+y_1=(w_{00}+w_{10}) x_0 + (w_{01}+w_{11}) x_1 + (w_{02}+w_{12}) x_2\), one can check easily all the derivatives because \(\nabla_y \, z=\begin{Bmatrix} 1 \\ 1 \end{Bmatrix}\). However, our result is more general and stands for any \(\mathbb{R}^m \rightarrow \mathbb{R}\) function between \(y\) and \(z\).

An implementation of the dot product with all the helper operations (transpose, etc) is detailed below:

src
 1def dot_forward(matrix: NodeVectorized, vector: NodeVectorized) -> NodeVectorized:
 2    result = NodeVectorized(
 3        matrix.val.dot(vector.val),
 4        (matrix, vector),
 5        f"({matrix.name} . {vector.name})",
 6        ".",
 7        matrix.is_constant and vector.is_constant,
 8    )
 9    result.backward = dot_backward
10    return result
src
 1def dot_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
 2    matrix: NodeVectorized = node.parents[0]
 3    vector: NodeVectorized = node.parents[1]
 4
 5    if not vector.is_constant:
 6        derivatives[vector] = addition_forward(
 7            dot_forward(transpose_matrix_forward(matrix), derivatives[node]),
 8            derivatives[vector],
 9        )
10
11    if not matrix.is_constant:
12        derivatives[matrix] = addition_forward(
13            dot_forward(
14                vert_vector_forward(derivatives[node]), horz_vector_forward(vector)
15            ),
16            derivatives[matrix],
17        )
src
 1def transpose_matrix_forward(node: NodeVectorized) -> NodeVectorized:
 2    result = NodeVectorized(
 3        node.val.T,
 4        (node,),
 5        f"({node.name})^T",
 6        "T",
 7        node.is_constant,
 8    )
 9    result.backward = transpose_matrix_backward
10    return result
src
1def transpose_matrix_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
2    if not node.parents[0].is_constant:
3        derivatives[node.parents[0]] = addition_forward(
4            transpose_matrix_forward(derivatives[node]),
5            derivatives[node.parents[0]],
6        )
src
 1def horz_vector_forward(node: NodeVectorized) -> NodeVectorized:
 2    if node.val.ndim == 2:
 3        temp: Vector = np.squeeze(node.val)
 4    else:
 5        temp: Vector = node.val[np.newaxis, :]
 6    result = NodeVectorized(
 7        temp,
 8        (node,),
 9        f"({node.name})^H",
10        "H",
11        node.is_constant,
12    )
13    result.backward = horz_vector_backward
14    return result
src
1def horz_vector_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
2    if not node.parents[0].is_constant:
3        derivatives[node.parents[0]] = addition_forward(
4            horz_vector_forward(derivatives[node]),
5            derivatives[node.parents[0]],
6        )
src
 1def vert_vector_forward(node: NodeVectorized) -> NodeVectorized:
 2    if node.val.ndim == 2:
 3        temp: Vector = np.squeeze(node.val)
 4    else:
 5        temp: Vector = node.val[:, np.newaxis]
 6    result = NodeVectorized(
 7        temp,
 8        (node,),
 9        f"({node.name})^V",
10        "V",
11        node.is_constant,
12    )
13    result.backward = vert_vector_backward
14    return result
src
1def vert_vector_backward(node: NodeVectorized, derivatives: Derivatives) -> None:
2    if not node.parents[0].is_constant:
3        derivatives[node.parents[0]] = addition_forward(
4            vert_vector_forward(derivatives[node]),
5            derivatives[node.parents[0]],
6        )

Exercise (easy)

  1. Study the implementation of the auto-differentiable dot product above. Comment. Test the dot product.

  2. Using this newly available dot operation, how would you manage to implement a split operation? A split operation consists in splitting a single node containing a vector of size \(n\) into \(n\) nodes of size \(1\). The goal of such an operation is to manipulate independently the different components of a vector typically in the last layer of a FNN.

  3. Same question with the sum operation which consists in summing all the components of a node containing a node.

Correction

  • dot_forward, the value is obtained by a matrix vector product.

  • dot_backward, the derivatives with respect to the vector is \(W^\top \cdot \nabla _y \, z\), the derivatives with respect to the matrix is \(\nabla _y \, z \cdot x^\top\).

  • transpose_matrix_forward, the value is obtained by a matrix transpose.

  • transpose_matrix_backard, the backward transpose is just a transpose.

  • horz_vector_forward, it creates a row vector (of dimension 2). To get the same mechanic of an inverse transpose which is a transpose as previsouly, the vector is reset to a dimension 1 if its already at dimension 2.

  • horz_vector_backard, the backward transpose is just a transpose.

  • vert_vector_forward, same as previously but to create a vertical vector of dimension 2.

  • vert_vector_backard, same as previously.

tests
 1def test_dot_product_vector(x_val) -> None:
 2    x = NodeVectorized(x_val, name="x")
 3    a = NodeVectorized(np.array([[1, 2, 3], [4, 5, 6]], dtype=FLOAT), name="a")
 4    z = dot_forward(a, x)
 5    print(f"\n x={x.val} f(x)={z.name}")
 6
 7    assert pytest.approx(z.val) == a.val.dot(x_val)
 8
 9    # ic(z)
10    # breakpoint()
11
12    dz: Derivatives = differentiate(z)
13    dzdx: NodeVectorized = dz[x]
14    dzda: NodeVectorized = dz[a]
15
16    # ic(dz)
17    # breakpoint()
18
19    assert pytest.approx(dzdx.val) == np.sum(a.val, axis=0)
20    assert pytest.approx(dzda.val) == np.ones((z.val.size, 1)).dot(x_val[np.newaxis, :])
21
22    d_dzdx: Derivatives = differentiate(dzdx)
23    d2zdx2: NodeVectorized = d_dzdx[x]
24
25    # breakpoint()
26
27    assert pytest.approx(d2zdx2.val) == np.zeros(1)

  1. The split operation requires to use a kind of one-hot encoding to create successice matrix with a single one and all zeros, with the one at each location. It can be done using the row of an identity matrix of the correct size and performing all the dot products with the vector. All the result nodes are gathered in a list.

src
1def split_forward(node: NodeVectorized) -> list[NodeVectorized]:
2    base_vectors: Array2D = np.eye(node.val.size)
3    nodes: list[NodeVectorized] = []
4    for ind, base in enumerate(base_vectors):
5        base_node = NodeVectorized(
6            base[np.newaxis, :], name=f"e{ind}", is_constant=True
7        )
8        nodes.append(dot_forward(base_node, node))
9    return nodes
tests
1def test_split():
2    values = np.array([1, 2], dtype=FLOAT)
3    node_two_var = NodeVectorized(values, name="v")
4    node_two_var_splitted: list[NodeVectorized] = split_forward(node_two_var)
5    sum_: NodeVectorized = addition_forward(
6        node_two_var_splitted[0], node_two_var_splitted[1]
7    )
8    d_sum: Derivatives = differentiate(sum_)
9    assert pytest.approx(d_sum[node_two_var].val) == np.ones(values.shape, dtype=FLOAT)

  1. The sum operation is the easiest, one can obtain the sum of every components of a vector by computing the dot product with a one line matrix of ones with the number of columns equal to the size of the vector. Thus, just create a constant matrix of ones and make the dot product.

src
1def sum_forward(node: NodeVectorized) -> NodeVectorized:
2    temp: Array2D = np.ones((1, node.val.size), dtype=FLOAT)
3    matrix = NodeVectorized(temp, name="1", is_constant=True)
4    result: NodeVectorized = dot_forward(matrix, node)
5    result.name = f"sum({node.name})"
6    return result
tests
1def test_sum():
2    values = np.array([1, 2], dtype=FLOAT)
3    node_two_var = NodeVectorized(values, name="v")
4    node_two_var_summed: NodeVectorized = sum_forward(node_two_var)
5    d_sum: Derivatives = differentiate(node_two_var_summed)
6    assert pytest.approx(d_sum[node_two_var].val) == np.ones(values.shape, dtype=FLOAT)

Exercise (easy)

Consider the following differential equation:

\[\begin{split}\begin{cases} \cfrac{\partial u}{\partial x} + u = 0 \\ u(x=0) = u_0 \end{cases}\end{split}\]

  1. For a PINN with two layers and an activation \(\sigma\) for the first and linear for the second, draw the computation graph of the residual loss for a single collocation point. Make sure to remove every superflous operation from the graph.

  2. Same question for the boundary condition loss.

Correction

The dot product with a vector of size 1 could have been omitted, but it is still interesting to make them appear to be perfectly clear about the nature of every node: matrix or vector. It matters for a question of dimension of the result.

The forward sweep of the FNN from \(x\) to \(u\) is in blue. The rest is the backward sweep to build the \(\cfrac{\partial u}{\partial x}\) and the residual and Dirichlet losses.

Obviously, if several colocation points were used, the residual squared of each point would have been summed.

exp_decay_graph

Computation graph for the residual and Dirichlet loss