Hybrid-SSM: Hydra-T5 Encoder–Decoder for Abstractive Summarization

Reproduction of the model in "Lost in Compression: The Failure of SSMs to Propagate Factual Context in Transformer Reconstruction" (Khandoker & Al Zubayer). A Hydra bidirectional State Space Model encoder compresses the input article, a 3-layer MLP adapter aligns its latent states, and a T5-base Transformer decoder generates the summary via cross-attention. The model is trained with a 3-stage progressive-unfreezing regime and evaluated on XSum and CNN/DailyMail.

article ──(BERT tokenizer)──► Hydra SSM encoder ──► H (B,N,768)
                                     │
                             LayerNorm (interface)
                                     │
                          3-layer MLP adapter ──► H̃ (B,N,768)
                                     │
summary ──(T5 tokenizer)──► T5-base decoder (cross-attn over H̃) ──► tokens

The encoder and decoder are both initialized from models pretrained on C4 (Hydra: MLM BERT-style; T5: span-corruption), matching the paper.

Architecture notes

• Encoder — Hydra (bidirectional Mamba-2 via quasiseparable matrix mixers), MosaicBERT-style, C4-pretrained, bert-base-uncased tokenizer, d_model=768.
• Adapter — interface LayerNorm + 3 Linear layers with GELU (LayerNorm → Linear → GELU → Linear → GELU → Linear). Depth/width configurable.
• Decoder — T5-base decoder + lm_head, d_model=768, SentencePiece tokenizer.
• Implemented as HydraT5(T5ForConditionalGeneration) with self.encoder swapped for a Hydra+adapter module, so HuggingFace generate() / beam search works unchanged. Encoder and decoder use different tokenizers; the data collator tokenizes input and target separately.

Repository layout

configs/            model / data / train-stage YAMLs
src/hybrid_ssm/
  encoders/         base + mock (CPU) + hydra (GPU) encoders
  adapters/         MLP adapter (+ interface LayerNorm)
  model.py          HydraT5 + HydraAdapterEncoder
  data/             dataset loading + dual-tokenizer collator
  training/         freezing logic + 3-stage Trainer
  eval/             generation + ROUGE/BERTScore
scripts/            train.py, evaluate.py, download_data.py, run_all_stages.sh
notebooks/          colab_train.ipynb (single-GPU end-to-end)
tests/              offline CPU tests (mock encoder, tiny random T5)

Datasets

Train on each target dataset's own train split and evaluate on its test split (standard summarization setup). Both load from the Hugging Face Hub:

Setup

CPU (development, tests, wiring checks)

pip install -e .
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 pytest -q     # 20 offline tests

The tests use a mock encoder + a tiny random T5, so they need no GPU and no downloads. They verify the adapter, the 3-stage freezing, model forward/backward, generate(), the collator, and an end-to-end train→generate→ROUGE pass.

GPU (real training/eval)

The real Hydra encoder needs CUDA kernels. On a GPU machine:

pip install -e ".[gpu]"                 # mamba-ssm, causal-conv1d, triton, einops
git clone https://github.com/goombalab/hydra.git
pip install -e ./hydra                  # the Hydra SSM package
export HYDRA_BERT_SRC=$PWD/hydra/hydra/bert/src   # Hydra BERT source
python scripts/download_data.py --data-config configs/data/cnn_dailymail.yaml \
    --model-config configs/model/hydra_t5_base.yaml

See notebooks/colab_train.ipynb for a single-GPU walkthrough.

Install order matters: install causal-conv1d before mamba-ssm, and match them to your torch/CUDA build. src/hybrid_ssm/encoders/hydra_encoder.py is the single integration point with the upstream Hydra checkout — adjust its loader there if your checkout exposes a different builder.

Training (3-stage regime)

# Full run + test evaluation on CNN/DailyMail
bash scripts/run_all_stages.sh configs/data/cnn_dailymail.yaml outputs/hydra_t5_cnndm
# XSum
bash scripts/run_all_stages.sh configs/data/xsum.yaml outputs/hydra_t5_xsum

--stage all keeps one model in memory and runs 1→2→3 (each stage resumes the previous stage's weights). Use --stage N --resume <ckpt> to run a single stage.

T5-base baseline (single-variable comparison)

python scripts/train.py --model-config configs/model/hydra_t5_base.yaml \
    --data-config configs/data/cnn_dailymail.yaml --baseline t5 --stage 3 \
    --output-dir outputs/t5_cnndm

Evaluation

python scripts/evaluate.py --model-config configs/model/hydra_t5_base.yaml \
    --data-config configs/data/cnn_dailymail.yaml \
    --checkpoint outputs/hydra_t5_cnndm/stage3/best --split test --bertscore

Reports ROUGE-1/2/L and (optional) BERTScore, and writes a JSON with sample generations. Factuality metrics (entity recall, hallucination via NLI) are the planned Phase-2 addition — see src/hybrid_ssm/eval/ for where they slot in.

What runs where

The CPU path guarantees the pipeline is wired correctly; paper-matching scores require the GPU path with the real Hydra encoder and pretrained checkpoints.