| \section{ESM-2 Contextual Encoder and Compression Pipeline} |
| \label{sec:encoder} |
|
|
| Our encoder transforms raw amino acid sequences into compressed contextual embeddings suitable for flow matching generation. The pipeline consists of four main stages: (1) sequence preprocessing and validation, (2) ESM-2 contextual embedding extraction, (3) statistical normalization, and (4) transformer-based compression with hourglass pooling. |
|
|
| \subsection{Encoder Architecture Overview} |
|
|
| The complete encoding pipeline $\mathcal{E}: \mathcal{S} \rightarrow \mathbb{R}^{L' \times d_{comp}}$ transforms sequences $s \in \mathcal{S}$ from the amino acid alphabet to compressed embeddings, where $L' = L/2$ due to hourglass pooling and $d_{comp} = 80$ is the compressed dimension: |
|
|
| \begin{align} |
| s &\rightarrow \mathbf{H}^{(esm)} \rightarrow \mathbf{H}^{(norm)} \rightarrow \mathbf{Z}^{(comp)} \label{eq:encoding_pipeline} |
| \end{align} |
|
|
| \subsubsection{Sequence Preprocessing and Validation} |
| \label{sec:preprocessing} |
|
|
| Input sequences undergo rigorous preprocessing to ensure compatibility with ESM-2 and biological validity: |
|
|
| \begin{enumerate} |
| \item \textbf{Canonical Amino Acid Filtering}: Only sequences containing the 20 canonical amino acids $\mathcal{A} = \{$A, C, D, E, F, G, H, I, K, L, M, N, P, Q, R, S, T, V, W, Y$\}$ are accepted. |
| |
| \item \textbf{Length Constraints}: Sequences are filtered to $L_{min} \leq |s| \leq L_{max}$ where $L_{min} = 2$ and $L_{max} = 50$ for antimicrobial peptides. |
| |
| \item \textbf{Sequence Standardization}: All sequences are converted to uppercase and stripped of whitespace. |
| |
| \item \textbf{Padding and Truncation}: Sequences are standardized to length $L = 50$ through zero-padding (shorter sequences) or truncation (longer sequences). |
| \end{enumerate} |
|
|
| The preprocessing function $\text{Preprocess}(s)$ ensures uniform input format: |
|
|
| \begin{align} |
| s' = \begin{cases} |
| s \oplus \mathbf{0}^{L-|s|} & \text{if } |s| < L \\ |
| s_{1:L} & \text{if } |s| \geq L |
| \end{cases} \label{eq:padding} |
| \end{align} |
|
|
| where $\oplus$ denotes concatenation and $\mathbf{0}^{k}$ represents $k$ padding tokens. |
|
|
| \subsubsection{ESM-2 Contextual Embedding Extraction} |
| \label{sec:esm_embedding} |
|
|
| We utilize the pre-trained ESM-2 model (esm2\_t33\_650M\_UR50D) to extract contextual per-residue embeddings. ESM-2's 33-layer transformer architecture captures evolutionary relationships and structural constraints learned from 65 million protein sequences. |
|
|
| The embedding extraction process follows ESM-2's standard protocol: |
|
|
| \begin{align} |
| \mathbf{T} &= \text{Tokenize}(s') \in \mathbb{R}^{L+2} \label{eq:tokenization}\\ |
| \mathbf{H}^{(raw)} &= \text{ESM-2}_{33}(\mathbf{T}) \in \mathbb{R}^{(L+2) \times 1280} \label{eq:esm_forward}\\ |
| \mathbf{H}^{(esm)} &= \mathbf{H}^{(raw)}_{2:L+1, :} \in \mathbb{R}^{L \times 1280} \label{eq:cls_eos_removal} |
| \end{align} |
|
|
| where tokenization adds special CLS and EOS tokens, and we extract representations from the 33rd (final) layer, removing the special tokens to obtain per-residue embeddings. |
|
|
| \subsubsection{Statistical Normalization} |
| \label{sec:normalization} |
|
|
| To stabilize training and ensure consistent embedding magnitudes across the dataset, we apply a two-stage normalization scheme computed from dataset statistics: |
|
|
| \begin{align} |
| \boldsymbol{\mu} &= \mathbb{E}[\mathbf{H}^{(esm)}], \quad \boldsymbol{\sigma}^2 = \text{Var}[\mathbf{H}^{(esm)}] \label{eq:dataset_stats}\\ |
| \mathbf{H}^{(z)} &= \text{clamp}\left(\frac{\mathbf{H}^{(esm)} - \boldsymbol{\mu}}{\boldsymbol{\sigma} + \epsilon}, -4, 4\right) \label{eq:z_score}\\ |
| \boldsymbol{\mu}_{min} &= \min(\mathbf{H}^{(z)}), \quad \boldsymbol{\mu}_{max} = \max(\mathbf{H}^{(z)}) \label{eq:minmax_stats}\\ |
| \mathbf{H}^{(norm)} &= \text{clamp}\left(\frac{\mathbf{H}^{(z)} - \boldsymbol{\mu}_{min}}{\boldsymbol{\mu}_{max} - \boldsymbol{\mu}_{min} + \epsilon}, 0, 1\right) \label{eq:minmax_norm} |
| \end{align} |
|
|
| where $\epsilon = 10^{-8}$ prevents division by zero, and clamping operations ensure numerical stability. This normalization scheme combines z-score standardization with min-max scaling to produce embeddings in $[0, 1]^{L \times 1280}$. |
|
|
| \subsubsection{Transformer-Based Compression with Hourglass Pooling} |
| \label{sec:compression} |
|
|
| The compressor $\mathcal{C}: \mathbb{R}^{L \times 1280} \rightarrow \mathbb{R}^{L/2 \times 80}$ employs a hourglass architecture inspired by ProtFlow, combining transformer self-attention with spatial pooling for efficient compression: |
|
|
| \begin{align} |
| \mathbf{H}^{(0)} &= \text{LayerNorm}(\mathbf{H}^{(norm)}) \label{eq:input_norm}\\ |
| \mathbf{H}^{(pre)} &= \text{TransformerEncoder}^{(2)}(\mathbf{H}^{(0)}) \label{eq:pre_transformer}\\ |
| \mathbf{H}^{(pool)} &= \text{HourglassPool}(\mathbf{H}^{(pre)}) \label{eq:hourglass_pool}\\ |
| \mathbf{H}^{(post)} &= \text{TransformerEncoder}^{(2)}(\mathbf{H}^{(pool)}) \label{eq:post_transformer}\\ |
| \mathbf{Z}^{(comp)} &= \tanh(\text{LayerNorm}(\mathbf{H}^{(post)}) \mathbf{W}^{(proj)} + \mathbf{b}^{(proj)}) \label{eq:final_projection} |
| \end{align} |
|
|
| The hourglass pooling operation reduces sequence length while preserving critical information: |
|
|
| \begin{align} |
| \text{HourglassPool}(\mathbf{H}) = \begin{cases} |
| \text{Reshape}(\mathbf{H}_{1:L-1}, [B, (L-1)/2, 2, D]) \text{ if } L \text{ is odd} \\ |
| \text{Reshape}(\mathbf{H}, [B, L/2, 2, D]) |
| \end{cases} \label{eq:reshape_pool} |
| \end{align} |
|
|
| followed by mean pooling across the grouped dimension: |
|
|
| \begin{align} |
| \mathbf{H}^{(pool)} = \text{Mean}(\text{Reshape}(\mathbf{H}), \text{dim}=2) \label{eq:mean_pool} |
| \end{align} |
|
|
| This pooling strategy reduces computational complexity while maintaining spatial relationships between adjacent residues. |
|
|
| \subsection{Transformer Architecture Details} |
|
|
| Both pre-pooling and post-pooling transformer encoders use identical architectures: |
|
|
| \begin{itemize} |
| \item \textbf{Layers}: 2 transformer encoder layers each (4 total) |
| \item \textbf{Attention Heads}: 8 multi-head attention heads |
| \item \textbf{Hidden Dimension}: 1280 (matching ESM-2) |
| \item \textbf{Feedforward Dimension}: 5120 (4× hidden dimension) |
| \item \textbf{Activation}: GELU activation in feedforward layers |
| \item \textbf{Dropout}: 0.1 dropout rate during training |
| \end{itemize} |
|
|
| The final projection layer $\mathbf{W}^{(proj)} \in \mathbb{R}^{1280 \times 80}$ compresses to the target dimension with tanh activation to bound outputs in $[-1, 1]^{L/2 \times 80}$. |
|
|
| \subsection{Training Objective and Optimization} |
|
|
| The encoder-decoder pair is trained end-to-end using reconstruction loss to ensure information preservation: |
|
|
| \begin{align} |
| \mathcal{L}_{\text{recon}} &= \mathbb{E}_{\mathbf{H} \sim \mathcal{D}} \left[ \|\mathbf{H} - \mathcal{D}(\mathcal{C}(\mathbf{H}))\|_2^2 \right] \label{eq:reconstruction_loss} |
| \end{align} |
|
|
| where $\mathcal{D}$ is the decompressor and $\mathcal{D}$ represents the dataset distribution. |
|
|
| Training employs AdamW optimization with cosine annealing: |
|
|
| \begin{align} |
| \text{lr}(t) = \text{lr}_{\min} + \frac{1}{2}(\text{lr}_{\max} - \text{lr}_{\min})(1 + \cos(\pi t / T)) \label{eq:cosine_schedule} |
| \end{align} |
|
|
| with warmup schedule for the first 10,000 steps: |
|
|
| \begin{align} |
| \text{lr}_{\text{warmup}}(t) = \text{lr}_{\max} \cdot \frac{t}{T_{\text{warmup}}} \label{eq:warmup_schedule} |
| \end{align} |
|
|
| \subsection{Computational Efficiency and Scalability} |
|
|
| The encoder pipeline is optimized for large-scale processing: |
|
|
| \begin{itemize} |
| \item \textbf{Batch Processing}: Dynamic batching with GPU memory management |
| \item \textbf{Memory Optimization}: Gradient checkpointing and mixed precision training |
| \item \textbf{Parallel Processing}: Multi-GPU support with data parallelism |
| \item \textbf{Storage Efficiency}: Individual and combined tensor storage formats |
| \end{itemize} |
|
|
| Processing statistics for our dataset: |
| \begin{itemize} |
| \item \textbf{Dataset Size}: 6,983 validated AMP sequences |
| \item \textbf{Processing Speed}: ~100 sequences/second on A100 GPU |
| \item \textbf{Memory Usage}: ~8GB GPU memory for batch size 32 |
| \item \textbf{Storage Requirements}: ~2.1GB for compressed embeddings |
| \end{itemize} |
|
|
| \subsection{Embedding Quality and Validation} |
|
|
| The compressed embeddings maintain high fidelity to the original ESM-2 representations: |
|
|
| \begin{itemize} |
| \item \textbf{Reconstruction MSE}: $< 0.01$ on validation set |
| \item \textbf{Cosine Similarity}: $> 0.95$ between original and reconstructed embeddings |
| \item \textbf{Downstream Performance}: Maintained classification accuracy on AMP prediction tasks |
| \item \textbf{Compression Ratio}: 16× reduction in embedding dimension (1280 → 80) |
| \end{itemize} |
|
|
| \begin{algorithm}[h] |
| \caption{ESM-2 Contextual Encoder Pipeline} |
| \label{alg:encoder} |
| \begin{algorithmic}[1] |
| \REQUIRE Raw amino acid sequences $\mathcal{S} = \{s_1, s_2, \ldots, s_N\}$ |
| \REQUIRE Pre-trained ESM-2 model and compressor weights |
| \REQUIRE Dataset normalization statistics $\{\boldsymbol{\mu}, \boldsymbol{\sigma}, \boldsymbol{\mu}_{min}, \boldsymbol{\mu}_{max}\}$ |
| \ENSURE Compressed embeddings $\mathbf{Z}^{(comp)} \in \mathbb{R}^{N \times L/2 \times 80}$ |
|
|
| \STATE \textbf{// Stage 1: Sequence Preprocessing} |
| \FOR{$i = 1$ to $N$} |
| \STATE $s_i' \leftarrow \text{Preprocess}(s_i)$ \COMMENT{Filter, pad/truncate to length $L$} |
| \STATE \textbf{assert} $|s_i'| = L$ and $s_i' \subset \mathcal{A}^L$ \COMMENT{Validate canonical AAs} |
| \ENDFOR |
|
|
| \STATE \textbf{// Stage 2: ESM-2 Embedding Extraction} |
| \STATE $\mathcal{B} \leftarrow \text{CreateBatches}(\{s_1', \ldots, s_N'\}, \text{batch\_size})$ |
| \FOR{$\mathbf{B} \in \mathcal{B}$} |
| \STATE $\mathbf{T} \leftarrow \text{ESM2Tokenize}(\mathbf{B})$ \COMMENT{Add CLS/EOS tokens} |
| \STATE $\mathbf{H}^{(raw)} \leftarrow \text{ESM-2}_{33}(\mathbf{T})$ \COMMENT{Extract layer 33 representations} |
| \STATE $\mathbf{H}^{(esm)} \leftarrow \mathbf{H}^{(raw)}[:, 1:L+1, :]$ \COMMENT{Remove CLS/EOS tokens} |
| \ENDFOR |
|
|
| \STATE \textbf{// Stage 3: Statistical Normalization} |
| \FOR{$i = 1$ to $N$} |
| \STATE $\mathbf{H}_i^{(z)} \leftarrow \text{clamp}\left(\frac{\mathbf{H}_i^{(esm)} - \boldsymbol{\mu}}{\boldsymbol{\sigma} + \epsilon}, -4, 4\right)$ |
| \STATE $\mathbf{H}_i^{(norm)} \leftarrow \text{clamp}\left(\frac{\mathbf{H}_i^{(z)} - \boldsymbol{\mu}_{min}}{\boldsymbol{\mu}_{max} - \boldsymbol{\mu}_{min} + \epsilon}, 0, 1\right)$ |
| \ENDFOR |
|
|
| \STATE \textbf{// Stage 4: Transformer Compression} |
| \FOR{$i = 1$ to $N$} |
| \STATE $\mathbf{H}_i^{(0)} \leftarrow \text{LayerNorm}(\mathbf{H}_i^{(norm)})$ \COMMENT{Input normalization} |
| \STATE $\mathbf{H}_i^{(pre)} \leftarrow \text{TransformerEncoder}^{(2)}(\mathbf{H}_i^{(0)})$ \COMMENT{Pre-pooling layers} |
| \STATE $\mathbf{H}_i^{(pool)} \leftarrow \text{HourglassPool}(\mathbf{H}_i^{(pre)})$ \COMMENT{Spatial pooling} |
| \STATE $\mathbf{H}_i^{(post)} \leftarrow \text{TransformerEncoder}^{(2)}(\mathbf{H}_i^{(pool)})$ \COMMENT{Post-pooling layers} |
| \STATE $\mathbf{Z}_i^{(comp)} \leftarrow \tanh(\text{LayerNorm}(\mathbf{H}_i^{(post)}) \mathbf{W}^{(proj)} + \mathbf{b}^{(proj)})$ |
| \ENDFOR |
|
|
| \STATE $\mathbf{Z}^{(comp)} \leftarrow \text{Stack}(\{\mathbf{Z}_1^{(comp)}, \ldots, \mathbf{Z}_N^{(comp)}\})$ |
| \RETURN $\mathbf{Z}^{(comp)}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
| \begin{algorithm}[h] |
| \caption{Hourglass Pooling Operation} |
| \label{alg:hourglass_pool} |
| \begin{algorithmic}[1] |
| \REQUIRE Input embeddings $\mathbf{H} \in \mathbb{R}^{B \times L \times D}$ |
| \ENSURE Pooled embeddings $\mathbf{H}^{(pool)} \in \mathbb{R}^{B \times L/2 \times D}$ |
|
|
| \IF{$L \bmod 2 = 1$} \COMMENT{Handle odd sequence lengths} |
| \STATE $\mathbf{H} \leftarrow \mathbf{H}[:, :L-1, :]$ \COMMENT{Remove last position} |
| \STATE $L \leftarrow L - 1$ |
| \ENDIF |
|
|
| \STATE $\mathbf{H}^{(reshaped)} \leftarrow \text{Reshape}(\mathbf{H}, [B, L/2, 2, D])$ \COMMENT{Group adjacent positions} |
| \STATE $\mathbf{H}^{(pool)} \leftarrow \text{Mean}(\mathbf{H}^{(reshaped)}, \text{dim}=2)$ \COMMENT{Average grouped positions} |
|
|
| \RETURN $\mathbf{H}^{(pool)}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
| \begin{algorithm}[h] |
| \caption{Dataset Statistics Computation} |
| \label{alg:dataset_stats} |
| \begin{algorithmic}[1] |
| \REQUIRE ESM-2 embeddings $\{\mathbf{H}_1^{(esm)}, \ldots, \mathbf{H}_N^{(esm)}\}$ |
| \ENSURE Normalization statistics $\{\boldsymbol{\mu}, \boldsymbol{\sigma}, \boldsymbol{\mu}_{min}, \boldsymbol{\mu}_{max}\}$ |
|
|
| \STATE $\mathbf{H}^{(flat)} \leftarrow \text{Concatenate}(\{\text{Flatten}(\mathbf{H}_i^{(esm)})\}_{i=1}^N)$ \COMMENT{Flatten all embeddings} |
|
|
| \STATE \textbf{// Compute z-score statistics} |
| \STATE $\boldsymbol{\mu} \leftarrow \text{Mean}(\mathbf{H}^{(flat)}, \text{dim}=0)$ \COMMENT{Per-dimension mean} |
| \STATE $\boldsymbol{\sigma}^2 \leftarrow \text{Var}(\mathbf{H}^{(flat)}, \text{dim}=0)$ \COMMENT{Per-dimension variance} |
| \STATE $\boldsymbol{\sigma} \leftarrow \sqrt{\boldsymbol{\sigma}^2 + \epsilon}$ \COMMENT{Add epsilon for stability} |
|
|
| \STATE \textbf{// Apply z-score normalization} |
| \STATE $\mathbf{H}^{(z)} \leftarrow \text{clamp}\left(\frac{\mathbf{H}^{(flat)} - \boldsymbol{\mu}}{\boldsymbol{\sigma}}, -4, 4\right)$ |
|
|
| \STATE \textbf{// Compute min-max statistics} |
| \STATE $\boldsymbol{\mu}_{min} \leftarrow \text{Min}(\mathbf{H}^{(z)}, \text{dim}=0)$ \COMMENT{Per-dimension minimum} |
| \STATE $\boldsymbol{\mu}_{max} \leftarrow \text{Max}(\mathbf{H}^{(z)}, \text{dim}=0)$ \COMMENT{Per-dimension maximum} |
|
|
| \STATE \textbf{// Save statistics for inference} |
| \STATE $\text{Save}(\{\boldsymbol{\mu}, \boldsymbol{\sigma}, \boldsymbol{\mu}_{min}, \boldsymbol{\mu}_{max}\}, \text{"normalization\_stats.pt"})$ |
|
|
| \RETURN $\{\boldsymbol{\mu}, \boldsymbol{\sigma}, \boldsymbol{\mu}_{min}, \boldsymbol{\mu}_{max}\}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
