Add comprehensive documentation: compressor_decompressor_latex.tex

documentation/compressor_decompressor.tex

@@ -0,0 +1,390 @@
+\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% reduction in embedding storage requirements
+\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}