| \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} | |