## 3 Methodology
### W2 reg Approach for Binary Classification
#### 3.2.2 Gradient Estimation  
We computed the gradient of Equation (3) using the standard backpropagation strategy [34]. For the empirical risk part of Equation (3), this requires computing the derivatives of the losses \(\ell(f_{\theta}(x_{i}),y_{i})\) with respect to the neural network outputs \(f_{\theta}(x_{i})\), something routinely performed by packages such as PyTorch, TensorFlow, or Keras, for all mainstream losses. For the 2-Wasserstein part of Equation (3), the authors of [7] proposed to use a mathematical strategy to compute the pseudo-derivatives of \(\mathcal{W}_{2}^{2}(\mu_{\theta,0},\mu_{\theta,1})\) with respect to the neural network outputs \(f_{\theta}(x_{i})\). Specifically, to compute the pseudo-derivative of a discrete and empirical approximation of \(\mathcal{W}_{2}^{2}(\mu_{\theta,0},\mu_{\theta,1})\) with respect to a mini-batch output \(f_{\theta}(x_{i})\), the following equation was used:  
\[\Delta_{\tau}\left[1_{s_{i}=0}\frac{f_{\theta}(x_{i})-cor_{1}(f_{\theta}(x_{i} ))}{n_{0}\Bigl{(}H_{0}^{j_{i}+1}-H_{0}^{j_{i}}\Bigr{)}}-1_{s_{i}=1}\frac{cor_{ 0}(f_{\theta}(x_{i}))-f_{\theta}(x_{i})}{n_{1}\Bigl{(}H_{1}^{i_{i}+1}-H_{1}^{ j_{i}}\Bigr{)}}\right]\,, \tag{4}\]  
where \(n_{s}\) is the number of observations in class \(S=s\) and \(H_{s}^{j}\) are discrete versions of the cumulative distribution functions \({\mathcal{H}_{\theta,s}}\) defined on a discrete grid of possible output values:  
\[\eta^{j}=\min_{i}(f_{\theta}(x_{i}))+j\Delta_{\eta}\,,\,j=1,\ldots,J_{\eta}\,, \tag{5}\]where \(\Delta_{\eta}=J_{\eta}^{-1}(\max_{i}(f_{\theta}(x_{i}))-\min_{i}(f_{\theta}(x_{i})))\) and \(J_{\eta}\) is the number of discretisation steps. We denote \(H_{s}^{j}=H_{s}(\eta^{j})\), and \(j_{i}\) is defined such that \(\eta^{j_{i}}\leq f_{\theta}(x_{i})<\eta^{j_{i}+1}\). Finally, \(cor_{s}(f_{\theta}(x_{i}))=H_{s}^{-1}(H_{|1-s|}(f_{\theta}(x_{i})))\).  
2.3 Distinction between Mini-Batch Observations and the Observations for \(H_{0}\) and \(H_{1}\)  
As shown in Equation (4), computing the pseudo-derivatives of the 2-Wasserstein distance \(W_{2}^{2}(\mu_{\theta,0}^{n},\mu_{\theta,1}^{n})\) with respect to the model predictions \(f_{\theta}(x_{i})\) requires computing the discrete cumulative distribution functions \(H_{s}\), with \(s\in\{0,1\}\). Computing \(H_{s}\) would ideally require computing \(f_{\theta}(x_{i})\) for all \(n\) observations \(x_{i}\) of the training set, which would be a computational bottleneck. To solve this issue, Reference [7] proposed approximating \(H_{s}\) at each mini-batch iteration, where Equation (4) is computed, using a subset of all training observations. This observation subset is composed of \(m\) randomly drawn observations in group \(S=0\), \(m\) other randomly drawn observations in group \(S=1\), and the mini-batch observations. This guaranties that there are at least \(m\) observations to compute either \(H_{0}\) or \(H_{1}\) and that the impact of each mini-batch observation is represented in \(H_{0}\) and \(H_{1}\). Note that these additional \(2m\) predictions do not require backpropagating any gradient information, so their computational burden is limited in terms of memory resources. Although it is also reasonable in terms of computational resources, the amount of \(2m\) additional predictions should remain relatively small to avoid significantly slowing down the gradient descent. In previous experiences on images, \(m=16\) or \(m=32\) often appeared as reasonable, as this allowed mitigating undesirable algorithmic biases and slowed down the whole training procedure by a factor of less than 2. Finally, preserving the amount of such additional predictions to something reasonable at each gradient descent step is at the heart of our methodological contribution when extending _W2 reg_ to multi-class classification.