SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Runa Eschenhagen*,1 Alexander Immer2,3 Richard E. Turner
Frank Schneider4,5 Philipp Hennig4,5
<sup>1</sup>University of Cambridge
$^{2}$ Department of Computer Science, ETH Zurich
3Max Planck Institute for Intelligent Systems
<sup>4</sup>University of Tübingen
5Tubingen AI Center
The core components of many modern neural network architectures, such as transformers, convolutional, or graph neural networks, can be expressed as linear layers with weight-sharing. Kronecker-Factored Approximate Curvature (K-FAC), a second-order optimisation method, has shown promise to speed up neural network training and thereby reduce computational costs. However, there is currently no framework to apply it to generic architectures, specifically ones with linear weight-sharing layers. In this work, we identify two different settings of linear weight-sharing layers which motivate two flavours of K-FAC - expand and reduce. We show that they are exact for deep linear networks with weight-sharing in their respective setting. Notably, K-FAC-reduce is generally faster than K-FAC-expand, which we leverage to speed up automatic hyperparameter selection via optimising the marginal likelihood for a Wide ResNet. Finally, we observe little difference between these two K-FAC variations when using them to train both a graph neural network and a vision transformer. However, both variations are able to reach a fixed validation metric target in  $50 - 75\%$  of the number of steps of a first-order reference run, which translates into a comparable improvement in wall-clock time. This highlights the potential of applying K-FAC to modern neural network architectures.
One of the key driving forces behind the success of deep learning is arguably the development and scaling of novel deep neural network (DNN) architectures like transformers (Vaswani et al., 2017) or graph neural networks (Battaglia et al., 2018; Scarselli et al., 2009). While the landscape of neural network architectures seems vast, their core building blocks like attention, graph, recurrent, and convolutional layers can be expressed as simple linear layers with weight-sharing. For example, the weight-sharing can happen over a sequence of tokens, e.g. in language modelling, over image patches, or by tying weights in the forward pass. The perspective of expressing the core neural network operations as linear operations has been formalised in a different context with the Tensor Programs framework (Yang, 2019), which develops a language to express arbitrary neural network computations as a composition of simple matrix multiplications and coordinate-wise nonlinearities. One downside of relying on increasing the scale of the models to achieve better performance is the increased compute cost. To decrease these costs, we can try to improve the efficiency of the optimisation algorithms, which motivates exploring second-order methods like Newton's method and natural gradient descent (Amari, 1998), which uses the Fisher information matrix, or short, the Fisher.
*Correspondence to: re393@cam.ac.uk.
Kronecker-Factored Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015, K-FAC) is an approximation to the Fisher, or for common loss functions equivalently, the generalised Gauss-Newton matrix (GGN). It has first been derived in the context of optimisation for linear layers and later for convolutional (Grosse & Martens, 2016) and recurrent (Martens et al., 2018) layers. In deep learning, using the Fisher/GGN instead of the Hessian has been the de-facto standard; in turn, K-FAC has arguably been one of the most popular approximations of the Fisher/GGN, probably due to its relative efficiency as it approximates each layer's Fisher/GGN independently with a Kronecker product. While K-FAC has also been used for transformers (Grosse et al., 2023; Ottawa et al., 2022; Pauloski et al., 2021; Zhang et al., 2019a) and a graph neural network (Izadi et al., 2020), there is no theoretical framework for these cases and for applying K-FAC to new architectures.
Contributions. We propose such a framework by leveraging the idea of linear layers with weight-sharing. This concept reveals the additional structure in the Fisher/GGN due to the weight-sharing that has to be explicitly considered when applying K-FAC. We identify two settings of such layers that motivate two different flavours of the K-FAC approximation: K-FAC-expand and K-FAC-reduce. The former corresponds to the Kronecker Factors for Convolution (KFC) approximation in Grosse & Martens (2016), to the simplest approximation for recurrent layers proposed in Martens et al. (2018), and has been used for transformers (Grosse et al., 2023; Osawa et al., 2022; Pauloski et al., 2021; Zhang et al., 2019a) – though without discussion or motivation. Notably, K-FAC-reduce is exact for certain settings, like deep linear networks with convolutions and average pooling, under the same conditions that K-FAC is exact for deep linear networks without weight-sharing, whereas the currently used K-FAC-expand is not. In practice, both approximations can be used in each setting and K-FAC-reduce generally has a lower computational and memory complexity than K-FAC-expand. We empirically verify this speed difference with a Wide ResNet on CIFAR-10. Moreover, we show that the two K-FAC variations applied to a graph neural network and vision transformer can reach a fixed validation metric target in  $50 - 75\%$  of the steps of a first-order reference method, which translates into an almost equivalent decrease in wall-clock time. Although this does not demonstrate that K-FAC is a superior training algorithm, it indicates the potential of extending K-FAC to modern deep learning architectures. In general, K-FAC can be used as a drop-in Hessian approximation, e.g., in Laplace approximations for Bayesian deep learning (Daxberger et al., 2021; MacKay, 1992; Ritter et al., 2018) or in natural gradient variational inference (Khan et al., 2018; Zhang et al., 2018). To demonstrate this, we show that K-FAC-reduce can speed up automatic weight decay selection via approximate marginal likelihood optimisation, compared to the previously used K-FAC-expand.
We consider the supervised learning setup with a dataset  $\mathcal{D}$  of  $N$  i.i.d. samples  $\{\pmb{x}_n,\pmb{y}_n\}_{n = 1}^N$ , with  $\pmb{x}_n\in \mathbb{R}^D$  and  $\pmb{y}_n\in \mathbb{R}^C$ , a neural net  $f_{\theta}:\mathbb{R}^{D}\to \mathbb{R}^{C}$ , parameterised with  $\pmb {\theta}\in \mathbb{R}^{P}$ , and a loss function  $\ell :\mathbb{R}^C\times \mathbb{R}^C\rightarrow \mathbb{R}$ , which is often equivalent to a negative log likelihood, i.e.  $\ell (\pmb {y},f_{\pmb{\theta}}(\pmb {x})) = -\log p(\pmb {y}|f_{\pmb{\theta}}(\pmb {x}))$ . The DNN is usually trained with a (stochastic) gradient-based iterative update rule, in its simplest form  $\pmb{\theta}_{t + 1} = \pmb{\theta}_t - \alpha \nabla_{\pmb{\theta}_t}\ell (\pmb {y},f_{\pmb{\theta}_t}(\pmb {x}))$ , where  $t$  indicates the iteration and  $\alpha$  is the learning rate. If we assume a local quadratic approximation of the loss around the current parameter iterate  $\pmb{\theta}_t$ , we get a preconditioned gradient update step, i.e.
$$
\boldsymbol {\theta} _ {t + 1} = \boldsymbol {\theta} _ {t} - \alpha \boldsymbol {C} _ {t} ^ {- 1} \nabla_ {\boldsymbol {\theta} _ {t}} \ell (\boldsymbol {y}, f _ {\boldsymbol {\theta} _ {t}} (\boldsymbol {x})), \tag {1}
$$
where  $C_t \in \mathbb{R}^{P \times P}$  is the symmetric positive definite (p.d.) preconditioner. Setting  $C_t$  to the Hessian of the loss w.r.t. the model parameters  $\pmb{H}_t \coloneqq \nabla_{\pmb{\theta}_t}^2 \ell(\pmb{y}, f_{\pmb{\theta}_t}(\pmb{x}))$  yields a simplified version of the classic Newton's method; however, since  $\pmb{H}_t$  is expensive to compute, store, and not guaranteed to be p.d., it is not commonly used in deep learning.
2.1 Second-order optimisation in deep learning
The term "second-order" is used ambiguously in the literature: it can refer to using second-order derivatives like the Hessian of the loss w.r.t. the parameters of the model, or approximations thereof, like the GGN, which however only contains second-order derivatives of the loss w.r.t. the model outputs and not the parameters. It can also refer to using some variation of the second moment of the (average) mini-batch gradient, which blurs the line between commonly used methods that are usually considered "first-order" methods like Adam (Kingma & Ba, 2015).
There have been numerous works on "second-order" methods in deep learning. For example, Martens (2010) proposes to use Hessian-free optimisation, Gupta et al. (2018) introduce an architecture-agnostic structured version of AdaGrad (Duchi et al., 2011) called Shampoo, Goldfarb et al. (2020) develop K-BFGS, a Kronecker-factored stochastic quasi-Newton method, and Ren & Goldfarb (2021) introduce Tensor Normal Training, which is similar to Shampoo, but approximates the Fisher; just to name a few among many more methods (Yang et al., 2022, 2023; Yao et al., 2021). While these methods have been shown to require fewer steps than commonly used first-order methods to reach similar performance in many settings, the mechanism behind their benefits is not fully understood. It could be orthogonal to unique properties of the true GGN or Fisher but might be more related to being able to adapt to the stochastic gradient noise (Kunstner et al., 2019) or to the fact that K-FAC with a specific but commonly used form of damping can be viewed as an approximation to gradient descent on neurons instead of weights (Benzing, 2022). In any case, generalising K-FAC to modern architectures is a promising direction because it is used in many training algorithms that can potentially speed up training and train a wider range of architectures. Moreover, it can be used to study the mechanisms behind the benefits of second-order methods in deep learning.
2.2 Linear weight-sharing layers
Here, we present three examples of how the core building blocks of modern DNN architectures can be expressed as linear layers with weight-sharing. DNNs have a layered structure, i.e. they can be written as  $f_{\theta} = f_{\theta_L} \circ \ldots \circ f_{\theta_\ell} \circ \ldots \circ f_{\theta_1}$ , with  $\pmb{\theta} = \mathrm{concat}\big(\pmb{\theta}_1, \dots, \pmb{\theta}_\ell, \dots, \pmb{\theta}_L\big)$  and  $L$  the number of layers;  $\mathrm{concat}(\cdot, \dots, \cdot)$  concatenates vector inputs to a larger vector. For a linear layer, we have  $f_{\theta_\ell}(\pmb{x}) = \phi(\pmb{W}_\ell \pmb{x} + \pmb{b}_\ell)$ , where  $\pmb{x} \in \mathbb{R}^{P_{\ell,\mathrm{in}}}, \pmb{W}_\ell \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times P_{\ell,\mathrm{in}}}, \pmb{b} \in \mathbb{R}^{P_{\ell,\mathrm{out}}}, \pmb{\theta}_\ell = \mathrm{concat}\big(\mathrm{vec}(\pmb{W}_\ell), \pmb{b}_\ell\big) \in \mathbb{R}^{P_\ell}$ , and  $P_\ell = P_{\ell,\mathrm{out}} P_{\ell,\mathrm{in}} + P_{\ell,\mathrm{out}}$ ;  $\mathrm{vec}(\cdot)$  vectorises a matrix by concatenating its column vectors and  $\phi$  is an element-wise nonlinearity. For simplicity, we subsume the bias into the weights. We are interested in linear layers with weights shared across an additional input dimension of size  $R$ , i.e. with inputs  $X \in \mathbb{R}^{R \times D}$ . The linear layer is applied to this input as  $X W^\top$  and the weight matrix  $W$  is shared across the additional first dimension, i.e. every element in  $W$  contributes to each row of the output.
Example 1: Attention. The transformer architecture is popular in vision (Dosovitskiy et al., 2021), language (Brown et al., 2020), and spatio-temporal modelling (Bi et al., 2022). The inputs and outputs of the attention operation (Bahdanau et al., 2015) in transformers are exactly shaped as described above, where  $R$  is the sequence length, e.g. the number of tokens in language data or the number of image patches in vision. For the  $i$ -th row of the output matrix  $O$  and the  $r$ -th row of the inputs  $X$ , the attention operation is generally defined as
$$
\boldsymbol {o} _ {i} = \sum_ {r = 1} ^ {R} A _ {i, r} \boldsymbol {x} _ {r}, \tag {2}
$$
where  $\mathbf{A} \in \mathbb{R}^{R \times R}$  is the attention matrix, which is normalised such that  $\sum_{r=1}^{R} A_{i,r} = 1$ ; alternatively, we can write it using matrix operations as  $\mathbf{O} = \mathbf{A}\mathbf{X}$ . Intuitively, the outputs are a weighted sum of the inputs, where each weighting factor indicates the importance of the  $r$ -th sequence element to the element with index  $i$ . The most common choice for  $\mathbf{A}$  is a variation of so-called self-attention, defined as  $\mathbf{A} = \text{softmax}_{\text{row}}(\mathbf{X}\mathbf{W}_Q^\top\mathbf{W}_K\mathbf{X}^\top)$ , where  $\mathbf{W}_Q \in \mathbb{R}^{P_{\text{out}} \times D}$  and  $\mathbf{W}_K \in \mathbb{R}^{P_{\text{out}} \times D}$  are weight matrices and  $\text{softmax}_{\text{row}}(\cdot)$  applies the softmax function  $\text{softmax}(\mathbf{x})_i := \exp(x_i) / \sum_{j=1}^{J} \exp(x_j)$  for an input  $\mathbf{x} \in \mathbb{R}^J$  to each row of the input matrix. This corresponds to a dot-product similarity measure of two independent linear projections of the input, allowing for asymmetric relationships between sequence elements. The two weight matrices  $\mathbf{W}_Q, \mathbf{W}_K$ , and any linear projection of the attention output are applied to all  $R$  sequence elements - they are linear weight-sharing layers.
Example 2: Convolution. A typical input to a 2d convolution is an image  $X^I \in \mathbb{R}^{C_{\mathrm{in}} \times H \times W}$ , where  $C_{\mathrm{in}}$  is the number of input channels,  $H$  is the height, and  $W$  the width of the image. We choose a kernel of size  $K \times K$ , such that we have a rank-4 weight tensor of shape  $C_{\mathrm{out}} \times C_{\mathrm{in}} \times K \times K$ , where  $C_{\mathrm{out}}$  is the number of output channels. For simplicity, following Grosse & Martens (2016), we assume that the convolution is performed with a stride of 1 and that the padding is chosen such that the height and width of the inputs are maintained, i.e. the output image will be  $O^I \in \mathbb{R}^{C_{\mathrm{out}} \times H \times W}$ .
<sup>1</sup>Our framework also applies to weight-sharing where the weights are used multiple times during the forward pass; however, we focus the presentation on an input with an explicit weight-sharing dimension.
The elements of the output tensor of the convolution are then defined as
$$
o _ {d, i, j} ^ {I} = \sum_ {c = 1} ^ {C _ {\mathrm {i n}}} \sum_ {k = 1} ^ {K} \sum_ {k ^ {\prime} = 1} ^ {K} X _ {c, i + k - 1, j + k ^ {\prime} - 1} ^ {I} W _ {d, c, k, k ^ {\prime}}; \tag {3}
$$
the input elements  $X_{c,i + k - 1,j + k^{\prime} - 1}$  are defined as 0 outside of the image's boundaries. Now, we can unfold the input tensor to a matrix  $\mathbf{X}\in \mathbb{R}^{HW\times C_{\mathrm{in}}K^2}$ , where we have extracted and vectorised the  $K\times K$  sized patches for each of the  $H\cdot W$  spatial locations. The weight tensor can be reshaped to the matrix  $\mathbf{W}\in \mathbb{R}^{C_{\mathrm{out}}\times C_{\mathrm{in}}K^2}$ . Using these quantities, we can write the convolution as  $\mathbf{O} = \mathbf{X}\mathbf{W}^{\intercal}\in \mathbb{R}^{HW\times C_{\mathrm{out}}}$ . Again, we can recognise the canonical form of a linear weight-sharing layer – here, with  $R = HW$ , i.e. the weights are shared across all spatial locations.
Example 3: Graph neural network layer. We follow Battaglia et al. (2018) since their formulation encompasses many graph neural network (GNN) layers. For the full example and notation, please refer to Appendix A. Once again, the core operations in these kinds of models can be expressed as linear weight-sharing layers. Here, the weights are shared across all  $R = N_V$  nodes or  $R = N_E$  edges of a graph, depending on whether node or edge features are being updated.
2.3 The generalised Gauss-Newton and the Fisher matrix
To avoid computing the Hessian, approximations like the generalised Gauss-Newton matrix (GGN) are usually used in deep learning (Botev et al., 2017; Martens, 2010; Schraudolph, 2002). Defining  $\Lambda(f_{\theta}(\boldsymbol{x}_n)) \coloneqq \nabla_{f_{\theta}}^2 \ell(\boldsymbol{y}_n, f_{\theta}(\boldsymbol{x}_n)) \in \mathbb{R}^{C \times C}$ , we have
$$
\mathbf {G G N} (\boldsymbol {\theta}) = \sum_ {n = 1} ^ {N} \boldsymbol {J} _ {\boldsymbol {\theta}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right) ^ {\top} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \boldsymbol {J} _ {\boldsymbol {\theta}} f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right), \tag {4}
$$
which is equivalent to the Hessian when  $f_{\theta}$  is a linear model or the residual  $g^{-1}(f_{\theta}(\pmb{x}_n)) - \pmb{y}_n$  is zero for all data points, where  $g^{-1}$  is the inverse link function, e.g. the softmax function for classification. For all likelihoods of the exponential family with natural parameterisation the GGN coincides with the Fisher information matrix (Martens, 2014; Wang, 2010). This is, for example, the case for the common cross-entropy loss corresponding to a categorical likelihood and the mean-square error loss corresponding to a Gaussian likelihood. Setting  $C_t$  in Equation (1) to the Fisher leads to natural gradient descent (Amari, 1998). We will only state our derivations for the GGN explicitly, but they follow analogously for the Fisher.
2.4 Kronecker-Factored Approximate Curvature
The GGN is a semi-p.d. approximation to the Hessian, of size  $P \times P$ , which is still prohibitively large for DNNs. K-FAC was proposed as an efficient approximation to a neural network's Fisher (Heskes, 2000; Martens & Grosse, 2015) and GGN (Botev et al., 2017) matrix. First, we only approximate the GGN for each layer, ignoring cross-layer interactions, which results in a block-diagonal matrix. We define  $J_{\theta_{\ell}}(\boldsymbol{x}_n) \coloneqq J_{\theta_{\ell}}f_{\theta}(\boldsymbol{x}_n) = J_{\boldsymbol{s}_{\ell,n}}f_{\theta}(\boldsymbol{x}_n)J_{\theta_{\ell}}\boldsymbol{s}_{\ell,n} \in \mathbb{R}^{C \times P_{\ell}}$  as the Jacobian of the model outputs w.r.t. the parameters of the  $\ell$ -th layer. Now we can write  $s_{\ell,n} = W_{\ell}\boldsymbol{a}_{\ell,n} = (\boldsymbol{a}_{\ell,n}^{\top} \otimes \mathbf{I}_{P_{\ell,\mathrm{out}}})\mathrm{vec}(W_{\ell})$  and with this we have  $J_{\theta_{\ell}}s_{\ell,n} = a_{\ell,n}^{\top} \otimes \mathbf{I}_{P_{\ell,\mathrm{out}}}$ . Additionally, by defining  $\boldsymbol{b}_{\ell,n} \coloneqq J_{\boldsymbol{s}_{\ell,n}}f_{\theta}(\boldsymbol{x}_n)^{\top} \in \mathbb{R}^{P_{\ell,\mathrm{out}} \times C}$  as the transposed Jacobian of the model outputs w.r.t. the pre-activations of the  $\ell$ -th layer, we note that  $J_{\theta_{\ell}}(\boldsymbol{x}_n)^{\top} = (\boldsymbol{a}_{\ell,n}^{\top} \otimes \mathbf{I}_{P_{\ell,\mathrm{out}}})^{\top}\boldsymbol{b}_{\ell,n} = \boldsymbol{a}_{\ell,n} \otimes \boldsymbol{b}_{\ell,n}$ . K-FAC approximates the GGN of the  $\ell$ -th layer as
$$
\mathbf {G G N} \left(\boldsymbol {\theta} _ {\ell}\right) \approx \underbrace {\left[ \frac {1}{N} \sum_ {n = 1} ^ {N} \boldsymbol {a} _ {\ell , n} \boldsymbol {a} _ {\ell , n} ^ {\intercal} \right]} _ {=: \boldsymbol {A} _ {\ell}} \otimes \underbrace {\left[ \sum_ {n = 1} ^ {N} \boldsymbol {b} _ {\ell , n} \boldsymbol {\Lambda} \left(f _ {\boldsymbol {\theta}} \left(\boldsymbol {x} _ {n}\right)\right) \boldsymbol {b} _ {\ell , n} ^ {\intercal} \right]} _ {=: \boldsymbol {B} _ {\ell}}, \tag {5}
$$
where we have replaced the (transposed) Jacobians in the definition of the GGN in Equation (4) by  $\pmb{a}_{\ell,n} \otimes \pmb{b}_{\ell,n}$  and approximated the sum of Kronecker products with a Kronecker product of sums.
2In the literature, the name "K-FAC" is used to refer to the approximation of the Fisher/GNN and to a specific algorithm using this approximation (Martens & Grosse, 2015); here, we use it to refer to just the approximation.
Note that the Kronecker product  $A \otimes B$  is in reverse order compared to in Figure 1, because we concatenate the columns in the derivation and the rows in the implementation when flattening a matrix with  $\mathrm{vec}(\cdot)$ .
Figure 1: Visualisation of K-FAC-expand and K-FAC-reduce in the expand setting. The shown quantities are for a single layer within a deep linear network. We have  $N = 4$ ,  $R = 2$ ,  $P_{\ell, \mathrm{in}} = 8$ ,  $P_{\ell, \mathrm{out}} = 8$ , and  $P_{\ell} = P_{\ell, \mathrm{in}} \cdot P_{\ell, \mathrm{out}} = 64$ . As we have seen in Section 3.2, K-FAC-expand is exact for the expand case in this setting and K-FAC-reduce is not. For better visibility, the color scale is not the same for all quantities, except for the approximation error (right) where black represents zero.
3 K-FAC for linear weight-sharing layers
In this section, we propose to distinguish between two different linear weight-sharing settings when applying K-FAC. The two settings motivate two corresponding approximations,  $K$ -FAC-expand and  $K$ -FAC-reduce. However, both approximations can be applied in each setting, see Listing 1 for an illustration with code. More details on the derivations and the proofs for Proposition 1 and Proposition 2 can be found in Appendix B.
3.1 The expand and reduce settings
Within a network, linear layers with weight-sharing can be classified based on the point where the weight-sharing dimension is aggregated. Specifically, there are three possible aggregation points, which will define the setting for the  $\ell$ -th layer:
(i) Before the  $\ell$ -th layer, i.e.  $A_{\ell,n} \in \mathbb{R}^{R \times P_{\ell,\mathrm{in}}}$  is reduced to  $\tilde{\pmb{a}}_{\ell,n} \in \mathbb{R}^{P_{\ell,\mathrm{in}}}$  before being multiplied with the weight matrix of the layer  $\pmb{W}_{\ell}$ .  $\rightarrow$  We are in the setting of a regular linear layer and no additional considerations are necessary when using K-FAC.
(ii) After the per-example loss, i.e. there will be  $N \cdot R$  labels and outputs of the model. The per-example loss is applied to each of the  $N \cdot R$  output-label pairs.  $\rightarrow$  This is the expand setting, as the loss for each of the  $N$  data points is expanded with  $R$  terms.
(iii) In between the pre-activation of the  $\ell$ -th layer  $S_{\ell,n} \in \mathbb{R}^{R \times P_{\ell,\mathrm{out}}}$  and the model output  $f_{\theta}(x_n)$ , i.e. before the final aggregation over the per-example losses.  $\rightarrow$  This is the reduce setting, as all weight-sharing dimensions have been reduced during the forward pass.
If we assume a single weight-sharing dimension which will be reduced once, the number of loss terms determines which setting applies to all linear weight-sharing layers within the model. So for our purposes, we can identify the setting we are in simply by looking at the form of the loss.
3.2 K-FAC-expand
The expand setting can be characterised by a loss with  $N\cdot R$  terms, similar to  $N\cdot R$  i.i.d. examples,  $\mathcal{L}_{\mathrm{expand}}(f_{\pmb{\theta}},\mathcal{D}):= -\sum_{n = 1}^{N}\sum_{r = 1}^{R}\log p(\pmb{y}_{n,r}|f_{\pmb{\theta}}(\pmb{x}_n)_r)$ , where  $f_{\pmb{\theta}}(\pmb{x}_n)_r$  is the  $r$ -th row of the model output  $f_{\pmb{\theta}}(\pmb{x}_n)\in \mathbb{R}^{R\times C}$  and  $\pmb{y}_{n,r}$  is the  $r$ -th row of the label  $\pmb{Y}_n\in \mathbb{R}^{R\times C}$ . A typical example of this type of loss function is language translation, where  $N$  is the dataset size and  $R$  is the sequence length.