Jigsaw Puzzle Reconstruction on STL-10

Problem

Exam project for the Deep Learning course of the MSc in Artificial Intelligence at the University of Bologna (2026). Each STL-10 image is cut into a 3×3 grid, every cell is centre-cropped from 32×32 to 28×28 so patches never share an edge, and the 9 patches are shuffled. The model must output the original 96×96 image, scored by mean absolute error (MAE) on 10,000 held-out images.

The rules: neural networks only, no pretrained weights, fewer than 6 M trainable parameters, and it has to run on Colab.

Method

• OrderNet — a small shared CNN embeds each patch into a 128-d token; three self-attention blocks (no positional encoding, so input order carries no information) let the patches compare with each other, and a dense head scores every (patch, cell) pair into a 9×9 matrix.
• Sinkhorn layer — 20 alternating row/column log-softmax normalisations turn the scores into a doubly-stochastic matrix. It is a parameter-free layer inside the network, not a discrete solver, so gradients flow through it.
• Differentiable assembly — each canvas cell is a convex combination of the patches weighted by that matrix. Where the network is confident it is a permutation for all practical purposes; where two patches are equally plausible the cell becomes a blend.
• RefineNet — a 3-level U-Net receives the canvas plus a mask of known pixels and inpaints only the 2-pixel seams (23% of the pixels). Known pixels are copied straight through.
• Label-free loss — placement MAE on the known pixels plus reconstruction MAE on the final image. Because each patch is an exact crop of the target, the placement term is zero only when every patch sits in its own cell, so routing is learned from the images alone.
• Temperature annealing — the Sinkhorn temperature drops linearly from 1.0 to 0.2 over 60 epochs so the assignment starts soft and sharpens gradually.

Results

Metric	Mean-patch baseline	Ours
Test MAE	0.1824	0.0428
Routing accuracy (patch in its exact cell)	–	78.6%
Puzzles fully solved	–	45.9%
Trainable parameters	0	3.26 M

A correctly solved puzzle leaves only a faint grid of error along the inpainted seams. Most remaining routing errors come from near-uniform patches (sky, walls, letterbox bars) that are almost interchangeable and cost almost no MAE when swapped.

Evaluation is exactly reproducible from the released weights; training takes about two hours on a Colab T4.