.. role:: python(code) :language: python .. index:: vectorized auto-differentiation PINN; vectorized auto-differentiation .. `vectorized operations 1 ` `vectorized operations 2 ` `PINN loss graph 2 ` .. _`sec:vectorized`: 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 :python:`Node` class to hold an array instead of a scalar. .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: NodeVectorized :name: Node_vectorized :linenos: :caption: src The :python:`differentiate` hardly changes too: .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: topological_sort :name: topological_sort_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: differentiate :name: differentiate_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: one :name: one_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: zero :name: zero_vectorized :linenos: :caption: src .. admonition:: 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.) .. _`ex:vectorized_mult_add_logistic`: .. collapse:: Correction 1. .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: addition_forward :name: addition_forward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: multiplication_forward :name: multiplication_forward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: assert_derivatives :name: assert_derivatives_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_val :name: x_val_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_pow_one_by_one :name: x_pow_one_by_one_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_sum_one_by_one :name: x_sum_one_by_one_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_derivatives_sum_one_by_one :name: test_derivatives_sum_one_by_one_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_derivatives_pow_one_by_one :name: test_derivatives_pow_one_by_one_vectorized :linenos: :caption: tests 2. .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: division_forward :name: division_forward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: exp_forward :name: exp_forward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: minus_one :name: minus_one_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: logistic_forward :name: logistic_forward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: exp_backward :name: exp_backward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: logistic_backward :name: logistic_backward_vectorized :linenos: :caption: src .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_inv_pow_one_by_one :name: x_inv_pow_one_by_one_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_derivatives_inv_one_by_one :name: test_derivatives_inv_one_by_one_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_exp_exp :name: x_exp_exp_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_derivatives_exp_exp :name: test_derivatives_exp_exp_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: x_logistic :name: x_logistic_vectorized :linenos: :caption: tests .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_derivatives_logistic :name: test_derivatives_logistic_vectorized :linenos: :caption: tests Dot product """"""""""" The dot product between a :math:`(n\times m)` matrix :math:`W` (weights) and a :math:`(m)` vector :math:`x` (values of the previous layer) to obtain the :math:`(n)` result vector :math:`y` is more interesting. The goal is to obtain the vectorized derivatives of the result vector with respect to the vector :math:`x` and the weights :math:`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 :math:`S` which sums up all the coefficients of a vector. .. figure:: dot_product_vectorized.png :width: 120 :alt: dot_product_vectorized Dot product vectorized .. math:: :label: eq:vectorized_dot_example \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 If a non vectorized auto-differentiation was used: + :math:`x_0` would see a scalar multiplication with :math:`w_{00}` and :math:`w_{10}`. + :math:`x_1` would see a scalar multiplication with :math:`w_{01}` and :math:`w_{11}`. + :math:`x_2` would see a scalar multiplication with :math:`w_{02}` and :math:`w_{12}`. The backward sweep is relatively straighforward: .. math:: :label: eq:vectorized_dot_backward_chain \nabla_y \, z = \begin{Bmatrix} \cfrac{\partial z}{\partial y_0} \\ \cfrac{\partial z}{\partial y_1} \end{Bmatrix} A chain rule using what we know for a product between scalars gives for the derivatives with respect to :math:`x` components: .. math:: :label: eq:vectorized_dot_backward_x_example \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 The second chain rule gives for the derivatives with respect to :math:`W` components: .. math:: :label: eq:vectorized_dot_backward_w_example \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 For the specific case of :math:`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 :math:`\nabla_y \, z=\begin{Bmatrix} 1 \\ 1 \end{Bmatrix}`. However, our result is more general and stands for any :math:`\mathbb{R}^m \rightarrow \mathbb{R}` function between :math:`y` and :math:`z`. An implementation of the dot product with all the helper operations (transpose, etc) is detailed below: .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: dot_forward :name: dot_forward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: dot_backward :name: dot_backward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: transpose_matrix_forward :name: transpose_matrix_forward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: transpose_matrix_backward :name: transpose_matrix_backward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: horz_vector_forward :name: horz_vector_forward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: horz_vector_backward :name: horz_vector_backward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: vert_vector_forward :name: vert_vector_forward :linenos: :caption: src .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: vert_vector_backward :name: vert_vector_backward :linenos: :caption: src .. admonition:: 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 :math:`n` into :math:`n` nodes of size :math:`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. .. _`ex:vectorized_dot_split_add`: .. collapse:: Correction 1. + :python:`dot_forward`, the value is obtained by a matrix vector product. + :python:`dot_backward`, the derivatives with respect to the vector is :math:`W^\top \cdot \nabla _y \, z`, the derivatives with respect to the matrix is :math:`\nabla _y \, z \cdot x^\top`. + :python:`transpose_matrix_forward`, the value is obtained by a matrix transpose. + :python:`transpose_matrix_backard`, the backward transpose is just a transpose. + :python:`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. + :python:`horz_vector_backard`, the backward transpose is just a transpose. + :python:`vert_vector_forward`, same as previously but to create a vertical vector of dimension 2. + :python:`vert_vector_backard`, same as previously. .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_dot_product_vector :name: test_dot_product_vector :linenos: :caption: tests 2. 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. .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: split_forward :name: split_forward :linenos: :caption: src .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_split :name: test_split :linenos: :caption: tests 3. 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. .. literalinclude:: ../../../../../../src/ai/fnn/autodiff_higher_vectorized.py :pyobject: sum_forward :name: sum_forward :linenos: :caption: src .. literalinclude:: ../../../../../../tests/ai/fnn/test_autodiff_higher_vectorized.py :pyobject: test_sum :name: test_sum :linenos: :caption: tests .. admonition:: Exercise (easy) Consider the following differential equation: .. math:: \begin{cases} \cfrac{\partial u}{\partial x} + u = 0 \\ u(x=0) = u_0 \end{cases} 1. For a PINN with two layers and an activation :math:`\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. .. _`ex:pinn_loss_graph`: .. collapse:: 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 :math:`x` to :math:`u` is in blue. The rest is the backward sweep to build the :math:`\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. .. figure:: exp_decay_graph.png :width: 400 :alt: exp_decay_graph Computation graph for the residual and Dirichlet loss .. The residual loss is: .. .. math:: .. J_{\textrm{residual}} &= \cfrac{1}{2 n_c} \sum_i (r_i)^2 \\ .. J_{\textrm{res}\,i} &= (r_i)^2 .. The goal is now to be able to obtain the gradient at every colocation with respect to the network parameters: .. :math:`\nabla_{\theta_j} J_{\textrm{res}\,i}`.