
\begin{algorithm}[!hbt]
\caption{Backward Pass of Straight-through Estimator}
\label{alg:ste_backward}
\KwInput{Gradient of loss w.r.t. output $\nabla\mathcal{L}$, Input probability $\vec{p}$, Probability $\boldsymbol{\pi}$ from forward pass, Slope $\texttt{slope} = 10$, 
}
\KwOutput{Gradient w.r.t. input vector $\nabla \vec{p}$}
{
    Compute the sigmoid gradient: 
    \[
    \vec{g}  \gets \frac{\texttt{slope} \cdot \exp(-\texttt{slope} \cdot (\vec{p}- \boldsymbol{\pi}))}{(1 + \exp(-\texttt{slope} \cdot (\vec{p}- \boldsymbol{\pi})))^2}
    \] 
}


{$\nabla \vec{p} \gets \vec{g} \cdot \nabla \mathcal{L}$ \tcp*{Compute the gradient with respect to input probability $\vec{p}$} }
\KwRet{$\nabla \vec{p}$}\\
\tcp*[h]{Note: All operations are element-wise}
\end{algorithm}
