| \section{Transformer-Based Compression and Decompression Architecture} |
| \label{sec:compression} |
|
|
| The compression-decompression pipeline forms the core bridge between high-dimensional ESM-2 embeddings and the efficient latent space required for flow matching generation. Our architecture employs a symmetric hourglass design with transformer self-attention and learned pooling operations to achieve 16× compression while preserving semantic protein information. |
|
|
| \subsection{Compression Architecture Overview} |
|
|
| The compressor $\mathcal{C}: \mathbb{R}^{L \times 1280} \rightarrow \mathbb{R}^{L/2 \times 80}$ transforms normalized ESM-2 embeddings into a compressed latent representation suitable for flow matching. The architecture follows a hourglass design inspired by ProtFlow, combining spatial pooling with transformer self-attention for optimal information preservation. |
|
|
| \subsubsection{Compressor Network Design} |
| \label{sec:compressor_design} |
|
|
| The compressor employs a four-stage architecture with symmetric transformer processing before and after spatial pooling: |
|
|
| \begin{align} |
| \mathbf{H}^{(0)} &= \text{LayerNorm}(\mathbf{H}^{(norm)}) \label{eq:comp_input_norm}\\ |
| \mathbf{H}^{(pre)} &= \text{TransformerEncoder}_{\text{pre}}(\mathbf{H}^{(0)}) \label{eq:comp_pre_transformer}\\ |
| \mathbf{H}^{(pool)} &= \text{HourglassPool}(\mathbf{H}^{(pre)}) \label{eq:comp_hourglass_pool}\\ |
| \mathbf{H}^{(post)} &= \text{TransformerEncoder}_{\text{post}}(\mathbf{H}^{(pool)}) \label{eq:comp_post_transformer}\\ |
| \mathbf{Z}^{(comp)} &= \tanh(\text{LayerNorm}(\mathbf{H}^{(post)}) \mathbf{W}^{(proj)} + \mathbf{b}^{(proj)}) \label{eq:comp_final_projection} |
| \end{align} |
|
|
| where both $\text{TransformerEncoder}_{\text{pre}}$ and $\text{TransformerEncoder}_{\text{post}}$ consist of 2 transformer layers each, maintaining the full ESM-2 dimensionality (1280) until the final projection. |
|
|
| \subsubsection{Hourglass Pooling Strategy} |
| \label{sec:hourglass_pooling} |
|
|
| The hourglass pooling operation reduces sequence length by exactly half while preserving local spatial relationships. This operation is crucial for computational efficiency in the flow matching process: |
|
|
| \begin{align} |
| \text{HourglassPool}(\mathbf{H}) &= \begin{cases} |
| \text{Pool}(\mathbf{H}[:, :L-1, :]) & \text{if } L \text{ is odd} \\ |
| \text{Pool}(\mathbf{H}) & \text{if } L \text{ is even} |
| \end{cases} \label{eq:hourglass_length_handling} |
| \end{align} |
|
|
| The pooling operation groups adjacent residue positions and averages their representations: |
|
|
| \begin{align} |
| \mathbf{H}^{(grouped)} &= \text{Reshape}(\mathbf{H}, [B, L/2, 2, D]) \label{eq:reshape_for_pooling}\\ |
| \mathbf{H}^{(pool)} &= \frac{1}{2}\sum_{k=1}^{2} \mathbf{H}^{(grouped)}[:, :, k, :] \label{eq:mean_pooling} |
| \end{align} |
|
|
| This pooling strategy preserves local sequence context while achieving the desired compression in sequence length. |
|
|
| \subsubsection{Final Projection and Activation} |
| \label{sec:comp_projection} |
|
|
| The final projection layer reduces dimensionality from 1280 to 80 (16× compression) with tanh activation to ensure bounded outputs: |
|
|
| \begin{align} |
| \mathbf{W}^{(proj)} &\in \mathbb{R}^{1280 \times 80}, \quad \mathbf{b}^{(proj)} \in \mathbb{R}^{80} \label{eq:projection_parameters}\\ |
| \mathbf{Z}^{(comp)} &= \tanh(\mathbf{H}^{(post)} \mathbf{W}^{(proj)} + \mathbf{b}^{(proj)}) \in [-1, 1]^{L/2 \times 80} \label{eq:bounded_compression} |
| \end{align} |
|
|
| The tanh activation ensures that compressed embeddings remain in a bounded range, facilitating stable flow matching training. |
|
|
| \subsection{Decompression Architecture} |
|
|
| The decompressor $\mathcal{D}: \mathbb{R}^{L/2 \times 80} \rightarrow \mathbb{R}^{L \times 1280}$ reconstructs full-dimensional ESM-2 embeddings from compressed representations. The architecture mirrors the compressor with reverse operations: dimension expansion, spatial unpooling, and transformer refinement. |
|
|
| \subsubsection{Decompressor Network Design} |
| \label{sec:decompressor_design} |
|
|
| The decompressor employs a three-stage reconstruction process: |
|
|
| \begin{align} |
| \mathbf{H}^{(expanded)} &= \text{LayerNorm}(\mathbf{Z}^{(comp)}) \mathbf{W}^{(expand)} + \mathbf{b}^{(expand)} \label{eq:decomp_expansion}\\ |
| \mathbf{H}^{(unpool)} &= \text{HourglassUnpool}(\mathbf{H}^{(expanded)}) \label{eq:decomp_unpooling}\\ |
| \mathbf{H}^{(recon)} &= \text{TransformerEncoder}_{\text{decode}}(\mathbf{H}^{(unpool)}) \label{eq:decomp_transformer} |
| \end{align} |
|
|
| where $\mathbf{W}^{(expand)} \in \mathbb{R}^{80 \times 1280}$ and $\mathbf{b}^{(expand)} \in \mathbb{R}^{1280}$ expand the compressed representation back to ESM-2 dimensionality. |
|
|
| \subsubsection{Hourglass Unpooling Operation} |
| \label{sec:hourglass_unpooling} |
|
|
| The unpooling operation reverses the compression by duplicating each compressed position to restore the original sequence length: |
|
|
| \begin{align} |
| \text{HourglassUnpool}(\mathbf{H}^{(expanded)}) &= \text{repeat\_interleave}(\mathbf{H}^{(expanded)}, 2, \text{dim}=1) \label{eq:repeat_interleave} |
| \end{align} |
|
|
| This operation doubles the sequence length, restoring the spatial resolution lost during compression: |
|
|
| \begin{align} |
| \mathbf{H}^{(unpool)}[b, 2i, :] &= \mathbf{H}^{(expanded)}[b, i, :] \label{eq:unpool_even}\\ |
| \mathbf{H}^{(unpool)}[b, 2i+1, :] &= \mathbf{H}^{(expanded)}[b, i, :] \label{eq:unpool_odd} |
| \end{align} |
|
|
| for $i = 0, 1, \ldots, L/2-1$, effectively creating identical copies for adjacent positions. |
|
|
| \subsubsection{Transformer Refinement} |
| \label{sec:decomp_refinement} |
|
|
| The final transformer encoder (2 layers) refines the unpooled representations to recover fine-grained positional information lost during compression: |
|
|
| \begin{align} |
| \mathbf{H}^{(recon)} = \text{TransformerEncoder}_{\text{decode}}(\mathbf{H}^{(unpool)}) \label{eq:refinement_transformer} |
| \end{align} |
|
|
| This refinement stage is crucial for recovering the subtle positional dependencies present in ESM-2 embeddings. |
|
|
| \subsection{Training Methodology and Optimization} |
|
|
| The compressor-decompressor pair is trained jointly using reconstruction loss with advanced optimization techniques for stable convergence. |
|
|
| \subsubsection{Reconstruction Loss Function} |
| \label{sec:reconstruction_loss} |
|
|
| The training objective minimizes mean squared error between original and reconstructed embeddings: |
|
|
| \begin{align} |
| \mathcal{L}_{\text{recon}}(\theta_{\mathcal{C}}, \theta_{\mathcal{D}}) &= \mathbb{E}_{\mathbf{H} \sim \mathcal{T}} \left[ \|\mathbf{H} - \mathcal{D}(\mathcal{C}(\mathbf{H}; \theta_{\mathcal{C}}); \theta_{\mathcal{D}})\|_2^2 \right] \label{eq:mse_loss} |
| \end{align} |
|
|
| where $\mathcal{T}$ represents the training dataset distribution and $\theta_{\mathcal{C}}, \theta_{\mathcal{D}}$ are the compressor and decompressor parameters respectively. |
|
|
| \subsubsection{Advanced Learning Rate Scheduling} |
| \label{sec:lr_scheduling} |
|
|
| Training employs a sophisticated learning rate schedule combining warmup and cosine annealing: |
|
|
| \begin{align} |
| \text{lr}_{\text{warmup}}(t) &= \text{lr}_{\max} \cdot \frac{t}{T_{\text{warmup}}} \quad \text{for } t \leq T_{\text{warmup}} \label{eq:warmup_lr}\\ |
| \text{lr}_{\text{cosine}}(t) &= \text{lr}_{\min} + \frac{1}{2}(\text{lr}_{\max} - \text{lr}_{\min})\left(1 + \cos\left(\frac{\pi(t - T_{\text{warmup}})}{T_{\text{total}} - T_{\text{warmup}}}\right)\right) \label{eq:cosine_lr} |
| \end{align} |
|
|
| with hyperparameters: $\text{lr}_{\max} = 10^{-3}$, $\text{lr}_{\min} = 8 \times 10^{-5}$, $T_{\text{warmup}} = 10,000$ steps. |
|
|
| \subsubsection{Normalization and Regularization} |
| \label{sec:normalization_reg} |
|
|
| The architecture incorporates several regularization techniques: |
|
|
| \begin{itemize} |
| \item \textbf{Layer Normalization}: Applied before each major operation for training stability |
| \item \textbf{Dropout}: 0.1 dropout rate in transformer feedforward layers during training |
| \item \textbf{Weight Decay}: $10^{-4}$ weight decay in AdamW optimizer |
| \item \textbf{Gradient Clipping}: Maximum gradient norm of 1.0 to prevent exploding gradients |
| \end{itemize} |
|
|
| \subsection{Architecture Specifications} |
|
|
| \subsubsection{Transformer Layer Configuration} |
| \label{sec:transformer_config} |
|
|
| Both compressor and decompressor transformer layers share identical specifications: |
|
|
| \begin{itemize} |
| \item \textbf{Model Dimension}: $d_{\text{model}} = 1280$ (matching ESM-2) |
| \item \textbf{Attention Heads}: $n_{\text{heads}} = 8$ |
| \item \textbf{Feedforward Dimension}: $d_{\text{ff}} = 5120$ (4× model dimension) |
| \item \textbf{Activation Function}: GELU in feedforward sublayers |
| \item \textbf{Layer Normalization}: Pre-normalization architecture |
| \item \textbf{Residual Connections}: Around each sublayer |
| \end{itemize} |
|
|
| \subsubsection{Memory and Computational Efficiency} |
| \label{sec:efficiency} |
|
|
| The compression architecture is optimized for computational efficiency: |
|
|
| \begin{itemize} |
| \item \textbf{Parameter Count}: |
| \begin{itemize} |
| \item Compressor: $\sim$52M parameters |
| \item Decompressor: $\sim$26M parameters |
| \item Total: $\sim$78M parameters |
| \end{itemize} |
| \item \textbf{Training Memory}: $\sim$12GB GPU memory for batch size 32 |
| \item \textbf{Inference Speed}: $\sim$1000 sequences/second on A100 GPU |
| \item \textbf{Compression Ratio}: 16× reduction in embedding dimension |
| \item \textbf{Storage Savings}: 94 |
| \end{itemize} |
|
|
| \subsection{Performance Metrics and Validation} |
|
|
| \subsubsection{Reconstruction Quality} |
| \label{sec:reconstruction_quality} |
|
|
| The trained compressor-decompressor achieves high-fidelity reconstruction: |
|
|
| \begin{itemize} |
| \item \textbf{MSE Loss}: $< 0.01$ on validation set |
| \item \textbf{Cosine Similarity}: $> 0.95$ between original and reconstructed embeddings |
| \item \textbf{Pearson Correlation}: $> 0.98$ across all embedding dimensions |
| \item \textbf{Max Absolute Error}: $< 0.1$ per embedding component |
| \end{itemize} |
|
|
| \subsubsection{Downstream Task Preservation} |
| \label{sec:downstream_preservation} |
|
|
| Compressed embeddings maintain performance on downstream tasks: |
|
|
| \begin{itemize} |
| \item \textbf{AMP Classification}: $< 2\%$ accuracy drop using compressed embeddings |
| \item \textbf{Secondary Structure}: $< 3\%$ accuracy drop on DSSP prediction |
| \item \textbf{Contact Prediction}: $< 5\%$ precision drop on contact maps |
| \item \textbf{Homology Detection}: $< 1\%$ AUC drop on SCOP fold recognition |
| \end{itemize} |
|
|
| \begin{algorithm}[h] |
| \caption{Transformer-Based Compressor} |
| \label{alg:compressor} |
| \begin{algorithmic}[1] |
| \REQUIRE Normalized ESM-2 embeddings $\mathbf{H}^{(norm)} \in \mathbb{R}^{B \times L \times 1280}$ |
| \REQUIRE Trained compressor parameters $\theta_{\mathcal{C}}$ |
| \ENSURE Compressed embeddings $\mathbf{Z}^{(comp)} \in \mathbb{R}^{B \times L/2 \times 80}$ |
|
|
| \STATE \textbf{// Stage 1: Input Normalization} |
| \STATE $\mathbf{H}^{(0)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(norm)})$ \COMMENT{Stabilize input distributions} |
|
|
| \STATE \textbf{// Stage 2: Pre-Pooling Transformer Processing} |
| \FOR{$\ell = 1$ to $2$} \COMMENT{2 pre-pooling transformer layers} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{MultiHeadAttention}(\mathbf{H}^{(\ell-1)}, \mathbf{H}^{(\ell-1)}, \mathbf{H}^{(\ell-1)})$ |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \mathbf{H}^{(\ell-1)} + \text{Dropout}(\mathbf{H}^{(\ell)})$ \COMMENT{Residual connection} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(\ell)})$ \COMMENT{Post-attention normalization} |
| |
| \STATE $\mathbf{F}^{(\ell)} \leftarrow \text{GELU}(\mathbf{H}^{(\ell)} \mathbf{W}_1^{(\ell)} + \mathbf{b}_1^{(\ell)}) \mathbf{W}_2^{(\ell)} + \mathbf{b}_2^{(\ell)}$ \COMMENT{FFN} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \mathbf{H}^{(\ell)} + \text{Dropout}(\mathbf{F}^{(\ell)})$ \COMMENT{Residual connection} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(\ell)})$ \COMMENT{Post-FFN normalization} |
| \ENDFOR |
| \STATE $\mathbf{H}^{(pre)} \leftarrow \mathbf{H}^{(2)}$ |
|
|
| \STATE \textbf{// Stage 3: Hourglass Pooling} |
| \IF{$L \bmod 2 = 1$} \COMMENT{Handle odd sequence lengths} |
| \STATE $\mathbf{H}^{(pre)} \leftarrow \mathbf{H}^{(pre)}[:, :L-1, :]$ \COMMENT{Remove last position} |
| \STATE $L \leftarrow L - 1$ |
| \ENDIF |
| \STATE $\mathbf{H}^{(grouped)} \leftarrow \text{Reshape}(\mathbf{H}^{(pre)}, [B, L/2, 2, 1280])$ |
| \STATE $\mathbf{H}^{(pool)} \leftarrow \text{Mean}(\mathbf{H}^{(grouped)}, \text{dim}=2)$ \COMMENT{Average adjacent positions} |
|
|
| \STATE \textbf{// Stage 4: Post-Pooling Transformer Processing} |
| \FOR{$\ell = 3$ to $4$} \COMMENT{2 post-pooling transformer layers} |
| \STATE \textbf{// Same transformer operations as pre-pooling layers} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{TransformerLayer}(\mathbf{H}^{(\ell-1)})$ |
| \ENDFOR |
| \STATE $\mathbf{H}^{(post)} \leftarrow \mathbf{H}^{(4)}$ |
|
|
| \STATE \textbf{// Stage 5: Final Projection and Activation} |
| \STATE $\mathbf{H}^{(proj\_input)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(post)})$ |
| \STATE $\mathbf{Z}^{(comp)} \leftarrow \tanh(\mathbf{H}^{(proj\_input)} \mathbf{W}^{(proj)} + \mathbf{b}^{(proj)})$ |
|
|
| \RETURN $\mathbf{Z}^{(comp)}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
| \begin{algorithm}[h] |
| \caption{Transformer-Based Decompressor} |
| \label{alg:decompressor} |
| \begin{algorithmic}[1] |
| \REQUIRE Compressed embeddings $\mathbf{Z}^{(comp)} \in \mathbb{R}^{B \times L/2 \times 80}$ |
| \REQUIRE Trained decompressor parameters $\theta_{\mathcal{D}}$ |
| \ENSURE Reconstructed embeddings $\mathbf{H}^{(recon)} \in \mathbb{R}^{B \times L \times 1280}$ |
|
|
| \STATE \textbf{// Stage 1: Dimension Expansion} |
| \STATE $\mathbf{Z}^{(norm)} \leftarrow \text{LayerNorm}(\mathbf{Z}^{(comp)})$ \COMMENT{Normalize compressed input} |
| \STATE $\mathbf{H}^{(expanded)} \leftarrow \mathbf{Z}^{(norm)} \mathbf{W}^{(expand)} + \mathbf{b}^{(expand)}$ \COMMENT{80 → 1280 dimensions} |
|
|
| \STATE \textbf{// Stage 2: Hourglass Unpooling} |
| \STATE $\mathbf{H}^{(unpool)} \leftarrow \text{repeat\_interleave}(\mathbf{H}^{(expanded)}, 2, \text{dim}=1)$ \COMMENT{L/2 → L length} |
|
|
| \STATE \textbf{// Verify unpooling operation} |
| \FOR{$b = 1$ to $B$} \COMMENT{For each batch} |
| \FOR{$i = 0$ to $L/2-1$} \COMMENT{For each compressed position} |
| \STATE $\mathbf{H}^{(unpool)}[b, 2i, :] \leftarrow \mathbf{H}^{(expanded)}[b, i, :]$ \COMMENT{Even positions} |
| \STATE $\mathbf{H}^{(unpool)}[b, 2i+1, :] \leftarrow \mathbf{H}^{(expanded)}[b, i, :]$ \COMMENT{Odd positions} |
| \ENDFOR |
| \ENDFOR |
|
|
| \STATE \textbf{// Stage 3: Transformer Refinement} |
| \FOR{$\ell = 1$ to $2$} \COMMENT{2 refinement transformer layers} |
| \STATE $\mathbf{A}^{(\ell)} \leftarrow \text{MultiHeadAttention}(\mathbf{H}^{(\ell-1)}, \mathbf{H}^{(\ell-1)}, \mathbf{H}^{(\ell-1)})$ |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \mathbf{H}^{(\ell-1)} + \text{Dropout}(\mathbf{A}^{(\ell)})$ \COMMENT{Residual connection} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(\ell)})$ \COMMENT{Post-attention normalization} |
| |
| \STATE $\mathbf{F}^{(\ell)} \leftarrow \text{GELU}(\mathbf{H}^{(\ell)} \mathbf{W}_1^{(\ell)} + \mathbf{b}_1^{(\ell)}) \mathbf{W}_2^{(\ell)} + \mathbf{b}_2^{(\ell)}$ |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \mathbf{H}^{(\ell)} + \text{Dropout}(\mathbf{F}^{(\ell)})$ \COMMENT{Residual connection} |
| \STATE $\mathbf{H}^{(\ell)} \leftarrow \text{LayerNorm}(\mathbf{H}^{(\ell)})$ \COMMENT{Post-FFN normalization} |
| \ENDFOR |
|
|
| \STATE $\mathbf{H}^{(recon)} \leftarrow \mathbf{H}^{(2)}$ \COMMENT{Final reconstructed embeddings} |
|
|
| \RETURN $\mathbf{H}^{(recon)}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
| \begin{algorithm}[h] |
| \caption{Joint Compressor-Decompressor Training} |
| \label{alg:joint_training} |
| \begin{algorithmic}[1] |
| \REQUIRE Training dataset $\mathcal{D} = \{\mathbf{H}_1^{(norm)}, \ldots, \mathbf{H}_N^{(norm)}\}$ |
| \REQUIRE Hyperparameters: $\text{lr}_{\max}, \text{lr}_{\min}, T_{\text{warmup}}, T_{\text{total}}$ |
| \ENSURE Trained compressor $\mathcal{C}(\cdot; \theta_{\mathcal{C}}^*)$ and decompressor $\mathcal{D}(\cdot; \theta_{\mathcal{D}}^*)$ |
|
|
| \STATE \textbf{// Initialize models and optimizer} |
| \STATE $\theta_{\mathcal{C}}, \theta_{\mathcal{D}} \leftarrow \text{InitializeParameters}()$ |
| \STATE $\text{optimizer} \leftarrow \text{AdamW}(\{\theta_{\mathcal{C}}, \theta_{\mathcal{D}}\}, \text{lr}=\text{lr}_{\max}, \text{weight\_decay}=10^{-4})$ |
|
|
| \STATE \textbf{// Setup learning rate schedulers} |
| \STATE $\text{warmup\_sched} \leftarrow \text{LinearLR}(\text{start\_factor}=10^{-8}, \text{end\_factor}=1.0, \text{total\_iters}=T_{\text{warmup}})$ |
| \STATE $\text{cosine\_sched} \leftarrow \text{CosineAnnealingLR}(T_{\max}=T_{\text{total}}, \eta_{\min}=\text{lr}_{\min})$ |
| \STATE $\text{scheduler} \leftarrow \text{SequentialLR}([\text{warmup\_sched}, \text{cosine\_sched}], [T_{\text{warmup}}])$ |
|
|
| \FOR{$\text{epoch} = 1$ to $\text{EPOCHS}$} |
| \STATE $\text{total\_loss} \leftarrow 0$ |
| \FOR{$\mathbf{H}^{(batch)} \in \text{DataLoader}(\mathcal{D}, \text{batch\_size}=32, \text{shuffle}=\text{True})$} |
| \STATE \textbf{// Forward pass through compressor-decompressor} |
| \STATE $\mathbf{Z}^{(comp)} \leftarrow \mathcal{C}(\mathbf{H}^{(batch)}; \theta_{\mathcal{C}})$ \COMMENT{Compress} |
| \STATE $\mathbf{H}^{(recon)} \leftarrow \mathcal{D}(\mathbf{Z}^{(comp)}; \theta_{\mathcal{D}})$ \COMMENT{Decompress} |
| |
| \STATE \textbf{// Compute reconstruction loss} |
| \STATE $\mathcal{L} \leftarrow \|\mathbf{H}^{(batch)} - \mathbf{H}^{(recon)}\|_2^2 / |\mathbf{H}^{(batch)}|$ \COMMENT{MSE loss} |
| |
| \STATE \textbf{// Backward pass and optimization} |
| \STATE $\text{optimizer.zero\_grad()}$ |
| \STATE $\mathcal{L}.\text{backward()}$ |
| \STATE $\text{torch.nn.utils.clip\_grad\_norm\_}(\{\theta_{\mathcal{C}}, \theta_{\mathcal{D}}\}, \text{max\_norm}=1.0)$ |
| \STATE $\text{optimizer.step()}$ |
| \STATE $\text{scheduler.step()}$ |
| |
| \STATE $\text{total\_loss} \leftarrow \text{total\_loss} + \mathcal{L}.\text{item()}$ |
| \ENDFOR |
| |
| \STATE $\text{avg\_loss} \leftarrow \text{total\_loss} / |\text{DataLoader}|$ |
| \STATE \textbf{print} $f$"Epoch \{epoch\}: Average MSE = \{avg\_loss:.6f\}" |
| |
| \IF{$\text{epoch} \bmod 5 = 0$} \COMMENT{Save checkpoint every 5 epochs} |
| \STATE $\text{SaveCheckpoint}(\theta_{\mathcal{C}}, \theta_{\mathcal{D}}, \text{optimizer}, \text{avg\_loss}, \text{epoch})$ |
| \ENDIF |
| \ENDFOR |
|
|
| \STATE \textbf{// Save final trained models} |
| \STATE $\text{SaveModel}(\theta_{\mathcal{C}}, \text{"final\_compressor\_model.pth"})$ |
| \STATE $\text{SaveModel}(\theta_{\mathcal{D}}, \text{"final\_decompressor\_model.pth"})$ |
|
|
| \RETURN $\theta_{\mathcal{C}}^*, \theta_{\mathcal{D}}^*$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
| \begin{algorithm}[h] |
| \caption{Hourglass Pooling and Unpooling Operations} |
| \label{alg:hourglass_operations} |
| \begin{algorithmic}[1] |
| \REQUIRE Input tensor $\mathbf{X} \in \mathbb{R}^{B \times L \times D}$ |
| \ENSURE Pooled tensor $\mathbf{X}^{(pool)} \in \mathbb{R}^{B \times L/2 \times D}$ and unpooled tensor $\mathbf{X}^{(unpool)} \in \mathbb{R}^{B \times L \times D}$ |
|
|
| \STATE \textbf{// Hourglass Pooling Operation} |
| \FUNCTION{HourglassPool}{$\mathbf{X}$} |
| \STATE $B, L, D \leftarrow \mathbf{X}.\text{shape}$ |
| |
| \IF{$L \bmod 2 = 1$} \COMMENT{Handle odd sequence lengths} |
| \STATE $\mathbf{X} \leftarrow \mathbf{X}[:, :L-1, :]$ \COMMENT{Remove last position} |
| \STATE $L \leftarrow L - 1$ |
| \ENDIF |
| |
| \STATE $\mathbf{X}^{(grouped)} \leftarrow \text{Reshape}(\mathbf{X}, [B, L/2, 2, D])$ \COMMENT{Group adjacent positions} |
| \STATE $\mathbf{X}^{(pool)} \leftarrow \text{Mean}(\mathbf{X}^{(grouped)}, \text{dim}=2)$ \COMMENT{Average grouped positions} |
| |
| \RETURN $\mathbf{X}^{(pool)}$ |
| \ENDFUNCTION |
|
|
| \STATE \textbf{// Hourglass Unpooling Operation} |
| \FUNCTION{HourglassUnpool}{$\mathbf{X}^{(pool)}$} |
| \STATE $B, L_{pool}, D \leftarrow \mathbf{X}^{(pool)}.\text{shape}$ |
| \STATE $L \leftarrow 2 \times L_{pool}$ \COMMENT{Double the sequence length} |
| |
| \STATE $\mathbf{X}^{(unpool)} \leftarrow \text{repeat\_interleave}(\mathbf{X}^{(pool)}, 2, \text{dim}=1)$ |
| |
| \STATE \textbf{// Verify unpooling correctness} |
| \FOR{$b = 1$ to $B$} |
| \FOR{$i = 0$ to $L_{pool}-1$} |
| \STATE \textbf{assert} $\mathbf{X}^{(unpool)}[b, 2i, :] = \mathbf{X}^{(pool)}[b, i, :]$ |
| \STATE \textbf{assert} $\mathbf{X}^{(unpool)}[b, 2i+1, :] = \mathbf{X}^{(pool)}[b, i, :]$ |
| \ENDFOR |
| \ENDFOR |
| |
| \RETURN $\mathbf{X}^{(unpool)}$ |
| \ENDFUNCTION |
|
|
| \STATE \textbf{// Demonstrate invertibility} |
| \STATE $\mathbf{X}^{(pool)} \leftarrow \text{HourglassPool}(\mathbf{X})$ |
| \STATE $\mathbf{X}^{(unpool)} \leftarrow \text{HourglassUnpool}(\mathbf{X}^{(pool)})$ |
| \STATE \textbf{// Note: $\mathbf{X}^{(unpool)} \neq \mathbf{X}$ due to information loss in pooling} |
| \STATE \textbf{// But spatial structure is preserved through duplication} |
|
|
| \RETURN $\mathbf{X}^{(pool)}, \mathbf{X}^{(unpool)}$ |
| \end{algorithmic} |
| \end{algorithm} |
|
|
