StoryReasoning

A small CNN/LSTM/GRU architecture that approximates selected ideas from the StoryReasoning framework — multimodal representation learning, temporal reasoning, grounding, and cross-frame entity consistency — sized to train on a free Google Colab GPU.

The problem#

Story understanding needs more than treating frames independently. A system has to keep track of visual entities, their textual descriptions, the temporal relationships between frames, and — the hard part — which character or object is which across the whole sequence.

The StoryReasoning paper tackles this with a large system: cross-frame object re-identification, face recognition, object and landmark detection, chain-of-thought reasoning, and grounding of textual references to visual entities, all on top of Qwen2.5-VL 7B. The paper reports that its fine-tuned Qwen Storyteller reduced hallucinations from 4.06 to 3.56 (about −12.3%) compared with the non-fine-tuned model. That result belongs to the paper's system; I use it only as context for what the full approach achieves.

The assessment asked for something different: show an understanding of the paper and its dataset, then design a reasonable smaller architecture inspired by it.

The dataset#

The StoryReasoning dataset gives each story several signals at once:

• images — a sequence of story frames
• story — frame-level descriptions in GDI-style structured tags
• chain_of_thought — Markdown reasoning with per-frame tables of characters, objects, and settings, each with an entity ID (char1, obj5, bg26, …) and a bounding box

That last field is what makes grounding possible at toy scale. When char1 appears in frames 1, 2, and 3, the dataset is telling you it's the same person, and where they are in each frame.

Architecture#

The task: given frames 1–4 and their descriptions, predict frame 5 and its description.

Training is staged, so each piece can be debugged and checkpointed on its own:

1. Text autoencoder. Text is tokenised with the bert-base-uncased tokenizer, but the encoder and decoder are small LSTMs (128-dimensional embedding, 128-dimensional latent, one layer), trained with teacher forcing and cross-entropy.
2. Visual autoencoder. Frames are resized to 60 × 125 and passed through a three-layer CNN (32 → 64 → 128 channels) into a 128-dimensional latent, with a transpose-convolution decoder trained on L1 reconstruction loss.
3. Sequence predictor. The pretrained encoders turn each context frame and its description into a 128-d visual latent and a 128-d text latent, which are concatenated into a 256-d fused vector per step and run through a GRU. Decoders then predict the next frame and the next description.
4. Grounding signals. Two extra losses add the paper's ideas in miniature (below).

The modification: CoT grounding and toy re-identification#

The baseline architecture uses only global frame and text embeddings, which can lose track of which entity is which. I added two grounding signals drawn from the dataset's chain-of-thought tables:

• A CoT grounding parser extracts entity IDs, descriptions, categories, and bounding boxes for every frame from the Markdown tables.
• A cross-frame re-identification loss. When an entity ID appears in two frames, the model crops both regions using the bounding boxes, encodes each crop with the visual encoder, and penalises the distance (MSE) between the two latents. It's a lightweight proxy for the paper's re-identification, not a tracking system.

A third, simpler loss aligns the mean visual latent of a sequence with the mean text latent of its descriptions. Having the same 128-d size doesn't put the two modalities in a shared semantic space on its own; this loss is what adds a weak shared-space constraint.

Results#

Both autoencoders learned steadily over five epochs:

The sequence predictor (164,992 trainable parameters) is the more honest story:

Total loss barely moves, so it isn't a useful headline on its own. Breaking it down: image loss stays small, text loss falls substantially in some epochs, and the alignment and re-identification losses fluctuate. Five epochs of training is unstable, and the components demonstrate the intended architecture more than they demonstrate performance — which is why qualitative inspection matters here.

The predicted frame is blurry and greyish, and the generated text follows the dataset's style without preserving entity identity. Both are expected. The visual decoder is a tiny CNN reconstructing 60 × 125 images through a 128-d bottleneck, so it loses high-frequency detail; its job is to give the sequence model a compact representation, not to compete with diffusion models or ViTs. The LSTM decoder is similarly small and loses long-range narrative structure.

Limitations#

What this implementation does not do, compared with the paper:

These gaps are deliberate, for feasibility on free Colab and for interpretability. Building the grounded pieces at small scale also shows why large pretrained vision-language models are necessary for high-quality grounded storytelling.

Code#

The repository is structured so the notebook is a thin, reproducible entry point: model definitions live in src/model.py, parsing and checkpoint utilities in src/utils.py, the staged training pipeline in src/train.py, and every hyperparameter in config.yaml. The notebook mounts Drive, installs dependencies, and calls run("config.yaml"); phases with an existing checkpoint are skipped.