Cosine-Gated LSTM for Sequential Learning

1. Overview

Cosine similarity meets recurrent gating

CGLSTM adds a cosine-controlled gate to the vanilla LSTM so the network can emphasise inputs that align with the current trajectory and down-weight irrelevant signals. It replaces the traditional sigmoid-only gating with a similarity-aware mechanism, delivering consistent gains across synthetic sequence tasks, vision benchmarks, NLP datasets, and the SocNavGym prediction challenge.

Chart

99.07% Accuracy

Row-wise MNIST (↑0.65 pts vs LSTM/GRU)

Math

0.0225 MAE

Adding problem (T=1000) lowest error

Robot

35% ↓ MAE @ k=10

SocNavGym vs Transformer baseline

Architecture overview from my viva slides: vanilla LSTM (left), cosine-gated module (centre), and integrated output path (right).

2. Research Motivation

• Problem observed in SocNavGym world models: vanilla LSTM-based predictors drifted as the horizon increased, producing inaccurate latent rollouts for Predictive World Models research.
• Hypothesis: a gate that measures alignment between the input trajectory and the current hidden state can decide whether to retain, blend, or suppress information.
• Goal: improve long-range dependency retention without the parameter overhead of Transformers, while keeping inference latency acceptable for real-time robotics.

3. Architecture

Cosine-gated layer

• Computes cosine similarity between the transformed input and both the previous and current hidden states (gate_ic, gate_co).
• Uses those similarities as scaling coefficients before fusing the input projection with the candidate hidden update.
• Concatenates the similarity-aware vector with the vanilla LSTM output and applies a linear projection, giving the RL policy or downstream head a richer latent signal.

def cosine_gate(x_proj, h_prev, h_curr, eps=1e-8):
    def cos(a, b):
        return torch.sum(a * b, dim=-1, keepdim=True) / (
            torch.norm(a, dim=-1, keepdim=True) * torch.norm(b, dim=-1, keepdim=True) + eps)

    gate_ic = cos(x_proj, h_prev)
    gate_co = cos(x_proj, h_curr)

    gated = gate_ic * x_proj + h_curr
    gated = gate_co * gated
    return torch.cat([gated, h_curr], dim=-1)

Model footprint

• 118,794 trainable parameters (≈44% fewer than the Transformer baseline evaluated in the paper).
• Drop-in replacement for LSTM blocks in seq2seq, world-model, or classification stacks.

4. Benchmark Evaluation

The IEEE PRML 2024 paper reports results across synthetic sequence modelling, vision, and NLP tasks (mean across three seeds). The table summarises the headline metrics and deltas against the strongest recurrent baseline per task.

* CGLSTM matches or exceeds recurrent baselines but Transformer remains ahead on IMDB and PTB in the published results. This is noted as a limitation below.

Benchmark comparison extracted from my viva deck: accuracy and MAE highlights across MNIST, Fashion-MNIST, IMDB, Penn Treebank, and the adding problem.

5. SocNavGym Performance

Cosine-gating was originally motivated by social navigation prediction. The IEEE PRML paper reports the following results when predicting latent trajectories in the SocNavGym environment configured with four moving humans, a table, and a flowerpot.

Key takeaways:

• 35% reduction in MAE at k=10 relative to the Transformer baseline, while keeping inference under 8 ms.
• Visual inspection (Fig. 3 in the paper) shows blue dominance for CGLSTM, indicating tighter alignment with ground-truth trajectories.
• The model was subsequently integrated into the Predictive World Models pipeline to replace the world-model backbone.

SocNavGym qualitative comparison from my viva: CGLSTM predictions (right) retain blue overlays indicating accurate future states, whereas vanilla LSTM and GRU (left) accumulate red error regions.

6. Implementation & Usage

Drop-in usage for PyTorch recurrent stacks

from cglstm import CosineGatedLSTMCell

class CosineGatedLSTM(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.cell = CosineGatedLSTMCell(input_size, hidden_size)

    def forward(self, inputs, hx=None):
        outputs = []
        h, c = hx
        for x in inputs:  # sequence-first tensor
            h, c = self.cell(x, (h, c))
            outputs.append(h)
        return torch.stack(outputs), (h, c)

# Swap into a world model or classifier
vae = VAE(latent_dim=64)
cglstm = CosineGatedLSTM(latent_dim=64, hidden_size=128)

Practical guidance

• Use cosine gating to improve long-horizon predictions before world-model rollouts (as applied in FallingBallEnv and SocNavGym).
• Pair with balanced training schedules; the paper highlights gains when bounce events are oversampled in FallingBallEnv.
• Monitor compute budget: this cell is ~2× heavier than GRU but far cheaper than the evaluated Transformer (0.007 s vs 0.028 s per prediction in SocNavGym).

7. Research Impact

• Published in IEEE PRML 2024, with proceedings pages 8–15.
• Serves as the latent model in the Predictive World Models and Adaptive Prediction Horizons contributions.
• Demonstrated improvements across vision and NLP benchmarks: +0.44 pts on Fashion-MNIST, +0.37 pts on MNIST, +9.5% lower MAE on the adding problem.
• Provided publicly accessible implementation (GitHub) supporting replication and extension.

8. Limitations & Next Steps

• Transformers remain stronger on language modelling: the Encoder Transformer attained lower perplexity on Penn Treebank (103.38 vs 104.43).
• Training cost: CGLSTM takes ~40% longer to train on Fashion-MNIST compared with GRU, so further optimisation or GPU-efficient kernels would help.
• IMDB accuracy plateau: GRU slightly outperformed CGLSTM (86.48% vs 86.30%); integrating cosine gates into GRU variants is proposed as future work in the paper.

Planned improvements include cosine-gated GRUs, multi-task transfer experiments, and deeper analysis of horizon-adaptive strategies for social navigation.

Slide excerpt summarising trade-offs: Transformer still leads on PTB perplexity, and training cost increases relative to GRU.

9. Resources & Citation

Technologies: Python • PyTorch • TensorFlow • Sequential Modelling • Cosine Similarity

Key Links:

• IEEE PRML 2024 Paper (PDF)
• CosineGatedLSTM GitHub
• FallingBallEnv Environment
• Publication Details

Citation (BibTeX)

@inproceedings{oguzie2024cglstm,
  title={Cosine-Gated Long Short-Term Memory for Sequential Learning},
  author={Oguzie, Goodluck},
  booktitle={IEEE 5th International Conference on Pattern Recognition and Machine Learning (PRML 2024)},
  pages={8--15},
  year={2024},
  organization={IEEE}
}