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