documentation/decoder_process.tex

\section{ESM-2 Language Model Head Decoder}
\label{sec:decoder}

Our decoder leverages the pre-trained ESM-2 language model head to convert generated embeddings back to amino acid sequences, avoiding traditional cosine similarity approaches in favor of a principled probabilistic decoding strategy that maintains the semantic structure of the protein embedding space.

\subsection{Decoder Architecture}

The decoding process consists of three main stages: (1) decompression from the flow-generated compressed space back to full ESM-2 embedding space, (2) projection through ESM-2's language model head, and (3) probabilistic sequence sampling.

\subsubsection{Embedding Decompression}
\label{sec:decompression}

The flow matching model generates embeddings in a compressed 80-dimensional space (16× compression from ESM-2's native 1280 dimensions). The decompressor $\mathcal{D}: \mathbb{R}^{L \times 80} \rightarrow \mathbb{R}^{L \times 1280}$ reconstructs full-dimensional embeddings:

\begin{align}
\mathbf{z}^{(dec)} &= \text{LayerNorm}(\mathbf{z}^{(comp)}) \mathbf{W}^{(proj)} \label{eq:proj}\\
\mathbf{x}^{(unpool)} &= \text{Unpool}(\mathbf{z}^{(dec)}) \label{eq:unpool}\\
\mathbf{h}^{(full)} &= \text{TransformerEncoder}(\mathbf{x}^{(unpool)}) \label{eq:decoder_transformer}
\end{align}

where $\mathbf{W}^{(proj)} \in \mathbb{R}^{80 \times 1280}$ is the learned projection matrix, and the unpooling operation restores the original sequence length through interpolation. The transformer encoder consists of 2 layers with 8 attention heads and 5120-dimensional feedforward networks.

\subsubsection{ESM-2 Language Model Head Projection}
\label{sec:lm_head}

Unlike approaches that use cosine similarity between generated embeddings and amino acid token embeddings, our method directly utilizes ESM-2's pre-trained language model head $\text{LM}_{\text{ESM-2}}$. This head was trained to predict amino acids from contextual embeddings during ESM-2's pre-training on evolutionary sequences, ensuring optimal alignment between embedding space and amino acid probabilities.

The language model head applies layer normalization followed by a linear projection:

\begin{align}
\mathbf{h}^{(norm)} &= \text{LayerNorm}_{\text{ESM-2}}(\mathbf{h}^{(full)}) \label{eq:esm_norm}\\
\mathbf{L}^{(full)} &= \mathbf{h}^{(norm)} \mathbf{W}^{(lm)} + \mathbf{b}^{(lm)} \label{eq:lm_projection}\\
\mathbf{L}^{(aa)} &= \mathbf{L}^{(full)}[:, :, \mathcal{I}_{\text{AA}}] \label{eq:aa_selection}
\end{align}

where $\mathbf{W}^{(lm)} \in \mathbb{R}^{1280 \times |\mathcal{V}|}$ and $\mathbf{b}^{(lm)} \in \mathbb{R}^{|\mathcal{V}|}$ are ESM-2's pre-trained language model parameters, $|\mathcal{V}|$ is ESM-2's full vocabulary size, and $\mathcal{I}_{\text{AA}} = \{4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23\}$ are the indices corresponding to the 20 canonical amino acids in ESM-2's vocabulary.

\subsubsection{Probabilistic Sequence Sampling}
\label{sec:sampling}

The logits $\mathbf{L}^{(aa)} \in \mathbb{R}^{L \times 20}$ are converted to amino acid probability distributions and sampled using nucleus sampling (top-p) for controlled stochasticity:

\begin{align}
\mathbf{P}_{i,j} &= \frac{\exp(\mathbf{L}^{(aa)}_{i,j} / \tau)}{\sum_{k=1}^{20} \exp(\mathbf{L}^{(aa)}_{i,k} / \tau)} \label{eq:softmax_temp}\\
\mathbf{P}^{(sorted)}_{i}, \mathbf{I}^{(sorted)}_{i} &= \text{sort}(\mathbf{P}_{i}, \text{descending}=\text{True}) \label{eq:sort_probs}\\
\mathbf{C}_{i} &= \text{cumsum}(\mathbf{P}^{(sorted)}_{i}) \label{eq:cumsum}\\
\mathbf{M}_{i} &= \mathbf{C}_{i} \leq p \label{eq:nucleus_mask}\\
a_i &\sim \text{Categorical}(\mathbf{P}^{(filtered)}_{i}) \label{eq:sample}
\end{align}

where $\tau = 0.8$ is the temperature parameter, $p = 0.9$ is the nucleus sampling threshold, and $\mathbf{P}^{(filtered)}_{i}$ contains only the probability mass within the top-p nucleus, renormalized to sum to 1.

\subsection{Advantages Over Cosine Similarity Approaches}

Our ESM-2 language model head approach offers several key advantages over traditional cosine similarity-based decoders:

\begin{enumerate}
    \item \textbf{Contextual Awareness}: The language model head was trained to predict amino acids from contextual embeddings, incorporating sequence context and evolutionary patterns that pure cosine similarity cannot capture.
    
    \item \textbf{Probability Calibration}: The pre-trained head provides well-calibrated probability distributions over amino acids, enabling principled uncertainty quantification and controlled sampling strategies.
    
    \item \textbf{Evolutionary Consistency}: ESM-2's training on evolutionary sequences ensures that the embedding-to-sequence mapping respects biological constraints and evolutionary relationships.
    
    \item \textbf{Reduced Bias}: Cosine similarity approaches can be biased toward high-frequency amino acids in the embedding space. The language model head learned to balance frequency with contextual appropriateness during pre-training.
\end{enumerate}

\subsection{Decoding Performance}

The decoder successfully converts flow-generated embeddings to valid amino acid sequences with high fidelity. For the 80 sequences generated with different CFG scales, we observe:

\begin{itemize}
    \item \textbf{Sequence Validity}: 100\% of decoded sequences contain only canonical amino acids
    \item \textbf{Length Consistency}: All sequences maintain the target length of 50 residues
    \item \textbf{Diversity}: Strong CFG (scale 7.5) produces the highest diversity while maintaining biological plausibility
    \item \textbf{AMP Classification}: 8.8\% of decoded sequences are classified as antimicrobial peptides by HMD-AMP, with Strong CFG achieving 20\% AMP rate
\end{itemize}

\subsection{Implementation Details}

The decoder is implemented using PyTorch and leverages the pre-trained ESM-2 model (esm2\_t33\_650M\_UR50D) from Facebook's ESM repository. Key implementation considerations include:

\begin{itemize}
    \item \textbf{Memory Efficiency}: Batch processing with automatic chunking to prevent out-of-memory errors
    \item \textbf{Numerical Stability}: Careful handling of temperature scaling and probability renormalization
    \item \textbf{Deterministic Sampling}: Optional seed control for reproducible sequence generation
    \item \textbf{Confidence Estimation}: Per-position maximum probability averaging for sequence confidence scoring
\end{itemize}

The complete decoding pipeline from compressed flow embeddings to amino acid sequences takes approximately 0.1 seconds per sequence on GPU hardware, enabling efficient large-scale sequence generation and analysis.

\begin{algorithm}[h]
\caption{ESM-2 Language Model Head Decoder}
\label{alg:decoder}
\begin{algorithmic}[1]
\REQUIRE Flow-generated compressed embeddings $\mathbf{Z}^{(comp)} \in \mathbb{R}^{B \times L \times 80}$
\REQUIRE Pre-trained ESM-2 model with language model head
\REQUIRE Temperature $\tau = 0.8$, nucleus threshold $p = 0.9$
\ENSURE Amino acid sequences $\mathbf{S} = \{s_1, s_2, \ldots, s_B\}$

\STATE $\mathbf{H}^{(full)} \leftarrow \text{Decompressor}(\mathbf{Z}^{(comp)})$ \COMMENT{Decompress to ESM-2 space}
\STATE $\mathbf{H}^{(norm)} \leftarrow \text{LayerNorm}_{\text{ESM-2}}(\mathbf{H}^{(full)})$ \COMMENT{Apply ESM-2 normalization}
\STATE $\mathbf{L}^{(full)} \leftarrow \mathbf{H}^{(norm)} \mathbf{W}^{(lm)} + \mathbf{b}^{(lm)}$ \COMMENT{Language model head projection}
\STATE $\mathbf{L}^{(aa)} \leftarrow \mathbf{L}^{(full)}[:, :, \mathcal{I}_{\text{AA}}]$ \COMMENT{Extract amino acid logits}
\STATE $\mathbf{P} \leftarrow \text{softmax}(\mathbf{L}^{(aa)} / \tau)$ \COMMENT{Temperature-scaled probabilities}

\FOR{$i = 1$ to $B$}
    \FOR{$j = 1$ to $L$}
        \STATE $\mathbf{p}^{(sorted)}, \mathbf{idx} \leftarrow \text{sort}(\mathbf{P}_{i,j}, \text{desc})$ \COMMENT{Sort probabilities}
        \STATE $\mathbf{c} \leftarrow \text{cumsum}(\mathbf{p}^{(sorted)})$ \COMMENT{Cumulative probabilities}
        \STATE $\mathbf{mask} \leftarrow \mathbf{c} \leq p$ \COMMENT{Nucleus mask}
        \STATE $\mathbf{p}^{(filtered)} \leftarrow \text{renormalize}(\mathbf{p}^{(sorted)}[\mathbf{mask}])$ \COMMENT{Filter and renormalize}
        \STATE $a_{i,j} \leftarrow \text{sample}(\mathbf{p}^{(filtered)})$ \COMMENT{Sample amino acid}
    \ENDFOR
    \STATE $s_i \leftarrow \text{decode\_indices}(\mathbf{a}_i)$ \COMMENT{Convert to sequence string}
\ENDFOR

\RETURN $\mathbf{S}$
\end{algorithmic}
\end{algorithm}