class_docstr: <a id="SquaredErrorBayesRisk"></a>
  
  ## Bayes Risk with Squared Error Loss
  
  Here the cost function is squared error,
  $$\sum_{k=1}^K (y_k - p_k)^2 = \Vert \mathbf{y} - \mathbf{p} \Vert_2^2$$
  
  We integrate this cost over all $\mathbf{p}$
  
  egin{align}
  \mathcal{L}(\Theta)
  &= -\log \Bigg(
   \int
    \Big[ \sum_{k=1}^K (y_k - p_k)^2 \Big]
    rac{1}{B(        extcolor{orange}{\mathbf{lpha}})}
    \prod_{k=1}^K p_k^{       extcolor{orange}{lpha_k} - 1}
   d\mathbf{p}
   \Bigg ) \
  &= \sum_{k=1}^K \mathbb{E} \Big[ y_k^2 -2 y_k p_k + p_k^2 \Big] \
  &= \sum_{k=1}^K \Big( y_k^2 -2 y_k \mathbb{E}[p_k] + \mathbb{E}[p_k^2] \Big)
  \end{align}
  
  Where $$\mathbb{E}[p_k] = \hat{p}_k = rac{ extcolor{orange}{lpha_k}}{S}$$
  is the expected probability when sampled from the Dirichlet distribution
  and $$\mathbb{E}[p_k^2] = \mathbb{E}[p_k]^2 +       ext{Var}(p_k)$$
   where
  $$  ext{Var}(p_k) = rac{   extcolor{orange}{lpha_k}(S -   extcolor{orange}{lpha_k})}{S^2 (S + 1)}
  = rac{\hat{p}_k(1 - \hat{p}_k)}{S + 1}$$
   is the variance.
  
  This gives,
  
  egin{align}
  \mathcal{L}(\Theta)
  &= \sum_{k=1}^K \Big( y_k^2 -2 y_k \mathbb{E}[p_k] + \mathbb{E}[p_k^2] \Big) \
  &= \sum_{k=1}^K \Big( y_k^2 -2 y_k \mathbb{E}[p_k] +  \mathbb{E}[p_k]^2 +   ext{Var}(p_k) \Big) \
  &= \sum_{k=1}^K \Big( ig( y_k -\mathbb{E}[p_k] ig)^2 +    ext{Var}(p_k) \Big) \
  &= \sum_{k=1}^K \Big( ( y_k -\hat{p}_k)^2 + rac{\hat{p}_k(1 - \hat{p}_k)}{S + 1} \Big)
  \end{align}
  
  This first part of the equation $ig(y_k -\mathbb{E}[p_k]ig)^2$ is the error term and
  the second part is the variance.
class_docstr_tok: ["<","a","id=","''","SquaredErrorBayesRisk","''",">","<","\/a",">","#","#","Bayes","Risk","with","Squared","Error","Loss","Here","the","cost","function","is","squared","error",",","$","$","\\sum_","{","k=1","}","^K","(","y_k","-","p_k",")","^2","=","\\Vert","\\mathbf","{","y","}","-","\\mathbf","{","p","}","\\Vert_2^2","$","$","We","integrate","this","cost","over","all","$","\\mathbf","{","p","}","$","\begin","{","align","}","\\mathcal","{","L","}","(","\\Theta",")","&","=","-\\log","\\Bigg","(","\\int","\\Big","[","\\sum_","{","k=1","}","^K","(","y_k","-","p_k",")","^2","\\Big","]","rac","{","1","}","{","B","(","extcolor","{","orange","}","{","\\mathbf","{","\u0007lpha","}","}",")","}","\\prod_","{","k=1","}","^K","p_k^","{","extcolor","{","orange","}","{","\u0007lpha_k","}","-","1","}","d\\mathbf","{","p","}","\\Bigg",")","\\","&","=","\\sum_","{","k=1","}","^K","\\mathbb","{","E","}","\\Big","[","y_k^2","-2","y_k","p_k","+","p_k^2","\\Big","]","\\","&","=","\\sum_","{","k=1","}","^K","\\Big","(","y_k^2","-2","y_k","\\mathbb","{","E","}","[","p_k","]","+","\\mathbb","{","E","}","[","p_k^2","]","\\Big",")","\\end","{","align","}","Where","$","$","\\mathbb","{","E","}","[","p_k","]","=","\\hat","{","p","}","_k","=","rac","{","extcolor","{","orange","}","{","\u0007lpha_k","}","}","{","S","}","$","$","is","the","expected","probability","when","sampled","from","the","Dirichlet","distribution","and","$","$","\\mathbb","{","E","}","[","p_k^2","]","=","\\mathbb","{","E","}","[","p_k","]","^2","+","ext","{","Var","}","(","p_k",")","$","$","where","$","$","ext","{","Var","}","(","p_k",")","=","rac","{","extcolor","{","orange","}","{","\u0007lpha_k","}","(","S","-","extcolor","{","orange","}","{","\u0007lpha_k","}",")","}","{","S^2","(","S","+","1",")","}","=","rac","{","\\hat","{","p","}","_k","(","1","-","\\hat","{","p","}","_k",")","}","{","S","+","1","}","$","$","is","the","variance",".","This","gives",",","\begin","{","align","}","\\mathcal","{","L","}","(","\\Theta",")","&","=","\\sum_","{","k=1","}","^K","\\Big","(","y_k^2","-2","y_k","\\mathbb","{","E","}","[","p_k","]","+","\\mathbb","{","E","}","[","p_k^2","]","\\Big",")","\\","&","=","\\sum_","{","k=1","}","^K","\\Big","(","y_k^2","-2","y_k","\\mathbb","{","E","}","[","p_k","]","+","\\mathbb","{","E","}","[","p_k","]","^2","+","ext","{","Var","}","(","p_k",")","\\Big",")","\\","&","=","\\sum_","{","k=1","}","^K","\\Big","(","\big","(","y_k","-\\mathbb","{","E","}","[","p_k","]","\big",")","^2","+","ext","{","Var","}","(","p_k",")","\\Big",")","\\","&","=","\\sum_","{","k=1","}","^K","\\Big","(","(","y_k","-\\hat","{","p","}","_k",")","^2","+","rac","{","\\hat","{","p","}","_k","(","1","-","\\hat","{","p","}","_k",")","}","{","S","+","1","}","\\Big",")","\\end","{","align","}","This","first","part","of","the","equation","$","\big","(","y_k","-\\mathbb","{","E","}","[","p_k","]","\big",")","^2","$","is","the","error","term","and","the","second","part","is","the","variance","."]
docstr: * `evidence` is $\mathbf{e} \ge 0$ with shape `[batch_size, n_classes]`
  * `target` is $\mathbf{y}$ with shape `[batch_size, n_classes]`
docstr_tok: ["*","`","evidence","`","is","$","\\mathbf","{","e","}","\\ge","0","$","with","shape","`","[","batch_size",",","n_classes","]","`","*","`","target","`","is","$","\\mathbf","{","y","}","$","with","shape","`","[","batch_size",",","n_classes","]","`"]
code: def forward(self, evidence: torch.Tensor, target: torch.Tensor):
      """
      * `evidence` is $\mathbf{e} \ge 0$ with shape `[batch_size, n_classes]`
      * `target` is $\mathbf{y}$ with shape `[batch_size, n_classes]`
      """
      # $\textcolor{orange}{\alpha_k} = e_k + 1$
      alpha = evidence + 1.0
      # $S = \sum_{k=1}^K \textcolor{orange}{\alpha_k}$
      strength = alpha.sum(dim=-1)
      # $\hat{p}_k = \frac{\textcolor{orange}{\alpha_k}}{S}$
      p = alpha / strength[:, None]
  
      # Error $(y_k -\hat{p}_k)^2$
      err = (target - p) ** 2
      # Variance $\text{Var}(p_k) = \frac{\hat{p}_k(1 - \hat{p}_k)}{S + 1}$
      var = p * (1 - p) / (strength[:, None] + 1)
  
      # Sum of them
      loss = (err + var).sum(dim=-1)
  
      # Mean loss over the batch
      return loss.mean()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/uncertainty/evidence/__init__.py
imports: import torch
  from labml import tracker
  from labml_helpers.module import Module
label_desc: Use image node_id 1 for calling the SquaredErrorBayesRisk obj's underlying member method code with example usage: obj.forward(evidence, target) and returns: loss