Authors: Dhruv Gupta (guptdhru) & Lucas Stewart (ludastew)
This project builds and evaluates 9 deep learning models on the CIFAR-10 dataset across three architectural families — fully connected (FC), convolutional (CNN), and transformer (Attention) — and across three task variants: plain classification, robustness to 16×16 patch shuffling (D-shuffletruffle), and robustness to 8×8 patch shuffling (N-shuffletruffle). The core objective for us is to demonstrate the limits of each architecture type and the critical importance of selecting the right architectural inductive bias for a given task.
python main.py --epochs 200 --model_class 'Plain-Old-CIFAR10-FC' --batch_size 1024 --learning_rate 3e-4 --l2_regularization 0.01--epochs(int): Number of training epochs (default: 100)--model_class(str): Which model to train (default:Plain-Old-CIFAR10-FC). We have provided the following options:- FC:
Plain-Old-CIFAR10-FC,D-shuffletruffle-FC,N-shuffletruffle-FC - CNN:
Plain-Old-CIFAR10-CNN,D-shuffletruffle-CNN,N-shuffletruffle-CNN - Attention:
Plain-Old-CIFAR10-Attention,D-shuffletruffle-Attention,N-shuffletruffle-Attention
- FC:
--batch_size(int): Batch Sizefor training (default: 128)--learning_rate(float): Learning rate for the optimizer (default: 0.01)--l2_regularization(float): Weight decay for AdamW (default: 0.0)
The CIFAR-10 training set contains 50,000 images. We split this into a 40,000-image training subset and a 10,000-image validation subset, then compute normalisation statistics exclusively from the training subset to prevent data leakage into validation or test evaluation.
A seeded random permutation of all 50,000 indices is generated using torch.Generator().manual_seed(0). The first 40,000 indices become the training set; the remaining 10,000 become the validation set. The mean and standard deviation are computed per channel over the 40,000 training images only, via loadOrComputeNormStats() from compute_norm_stats.py. These values are saved to norm_stats.npz and reloaded on subsequent runs to ensure reproducibility across machines.
Computed Statistics (from training subset only):
| Channel | Mean | Std |
|---|---|---|
| Red | 0.49135 | 0.24711 |
| Green | 0.48225 | 0.24364 |
| Blue | 0.44674 | 0.26164 |
We are using torch.randperm with a fixed seed to ensure the same split is reproduced on every machine and across every run, which is important as we are training and comparing multiple models. We are also computing statistics only from the training data as using the full 50,000 images would allow the model to indirectly observe validation distribution during normalisation.
The goal here is to maximise classification accuracy on the standard CIFAR-10 test set using three architecturally distinct models, each with at least three layers. No constraints are placed on the test sets used here beyond the clean CIFAR-10 test set — the models are free to exploit all spatial structure in the image.
Each model is trained by running main.py with the appropriate --model_class flag. The training loop evaluates the validation set every epoch. The best checkpoint is selected by maximising validation accuracy and saved as {model_class}.pth. After training, the best checkpoint is reloaded and evaluated on the clean CIFAR-10 test set as well as both patch-shuffled test sets.
Training augmentation for all Task 1 models: RandomCrop(32, padding=4) + RandomHorizontalFlip() + RandAugment(num_ops=2, magnitude=9).
Optimiser: AdamW. Scheduler: linear warmup for the first epochs // 10 epochs (start factor 0.1), then cosine annealing to eta_min=1e-6.
The image is flattened from 3×32×32 = 3072 dimensions. Four fully connected layers progressively compress the representation, each followed by BatchNorm1d, GELU activation, and Dropout:
Flatten → Linear(3072, 2048) → BatchNorm1d → GELU → Dropout(0.2)
→ Linear(2048, 1024) → BatchNorm1d → GELU → Dropout(0.2)
→ Linear(1024, 512) → BatchNorm1d → GELU → Dropout(0.1)
→ Linear( 512, 256) → BatchNorm1d → GELU → Dropout(0.1)
→ Linear( 256, 10)
extractEmbedding(x) returns the 256-dimensional output of featureExtractor (pre-classifier).
| Parameter | Value |
|---|---|
| Epochs | 200 |
| Batch Size | 1024 |
| Learning Rate | 3e-4 |
| L2 Regularization | 0.01 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 1.129 | 65.62% |
| Patch-16 | 2.139 | 29.69% |
| Patch-8 | 2.377 | 21.81% |
Training Loss
Accuracy
The FC network achieves 65.62% clean accuracy. Its performance collapses completely on both shuffled test sets because every weight in the network is tied to a specific pixel position. Shuffling patches relocates pixels to positions the network has never associated them with; the learned spatial correlations are entirely destroyed. This collapse is the expected and intended behaviour for a plain FC network, demonstrating that positional specificity is a fundamental limitation of this architecture family.
We chose GELU over ReLU because it provides smoother gradients, which helps with the deeper network. BatchNorm1d before each activation stabilises training and acts as an additional regulariser. Dropout rates of 0.2 in the wider early layers and 0.1 in the narrower later layers reflect the higher redundancy in the wider layers.
The architecture is a standard 3+1 CNN architecture consisting of 3x3 kernels with appropriate stride and padding. Each convolution is followed by a batch normalization, ReLU, and then max pooled with a 2x2 kernel. The number of channels is doubled each time, from 32 to 64 to 128.
The 2D network is then flattened for the fully-connected portion which consists of one final hidden layer (again doubling the effective channels from 128 to 256) before the final output layer. A relatively aggressive dropout of 0.5 was used, in accordance with my experience with standard practice in NLP.
For testing, the default hyperparameters were used.
# CNN Portion ---]
Conv2d( 3, 32, kernel_size=3, stride=1, padding=1)
-> BatchNorm2d -> ReLU -> MaxPool2d(kernel_size=2, stride=2) ->
Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
-> BatchNorm2d -> ReLU -> MaxPool2d(kernel_size=2, stride=2) ->
Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
-> BatchNorm2d -> ReLU
---|
# FC Portion ---]
Flatten -> Linear(8192, 256) -> BatchNorm1d -> ReLU -> Dropout(0.5) ->
-> Linear( 256, 10)
---|
| Parameter | Value |
|---|---|
| Epochs | 100 |
| Batch Size | 128 |
| Learning Rate | 1e-2 |
| L2 Regularization | 0.00 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.631 | 88.02% |
| Patch-16 | 1.430 | 55.84% |
| Patch-8 | 2.192 | 34.58% |
Training Loss
Accuracy
The CNN achieved much greater accuracy than the FC network throughout despite a simpler structure and activation function. Clean accuracy was 71.59% vs FC's 65.62%, while Patch-16 was 40.50% vs FC's 29.69%. This reflects strongly on the CNN's inherent advantage in image processing due to the localization of its kernels, passing the structure of the image down through the layers instead of each pixel being treated in isolation.
We divided the image into non-overlapping 4×4 patches, producing 64 tokens. Each token is a flattened patch of dimension 3×4×4 = 48. A learnable CLS token and learnable positional embeddings are added before the transformer blocks:
Image → 64 patches (4×4×3 = 48 dims each)
→ Linear(48, 512) [patch embedding]
→ prepend CLS token + positional embedding
→ Dropout(0.1)
→ 8 × TransformerEncoderBlock(512, 8 heads, mlpDim=1024)
→ LayerNorm(512)
→ CLS token → Linear(512, 10)
Each TransformerEncoderBlock uses pre-norm design:
x → LayerNorm → MultiheadAttention → + residual
→ LayerNorm → FFN(512→1024→512) → + residual
extractEmbedding(x) returns the 512-dimensional LayerNorm-processed CLS token.
| Parameter | Value |
|---|---|
| Epochs | 300 |
| Batch Size | 1024 |
| Learning Rate | 6e-4 |
| L2 Regularization | 0.05 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.677 | 88.24% |
| Patch-16 | 1.894 | 55.16% |
| Patch-8 | 2.059 | 50.92% |
Training Loss
Accuracy
Our transformer achieves 88.24% clean accuracy — significantly above the FC network and well above our own preliminary run (Run 1, same architecture without CNN patch encoder or warmup scheduler) which achieved 75.94% [1]. This demonstrates the substantial impact of the architectural and training improvements made between runs. It collapses on both shuffled test sets because the learnable positional embedding encodes a token's location as part of its identity. After shuffling, patch 5 (for example) carries different visual content but still receives the "position 5" signal — the model receives irreconcilably contradictory information.
We chose small 4×4 patches over larger alternatives to give the model 64 tokens for richer spatial coverage than 8×8 patches (16 tokens) or 16×16 patches (4 tokens). The CLS token aggregates a global representation across all 64 tokens. Pre-norm (LayerNorm before attention) is more stable than post-norm when training from scratch. The learning rate of 6e-4 follows the linear scaling rule from the batch size of 1024 (base 3e-4 at batch 256). Warmup is essential for transformers as random attention weights at initialisation produce destructively large gradients without it.
| Model | Clean Accuracy | Clean Loss |
|---|---|---|
Net_FC |
65.62% | 1.129 |
Net_CNN |
71.59% | 1.173 |
Net_Attention |
88.24% | 0.677 |
The D-shuffletruffle has a large receptive field — each lens covers a coarse region. We model this as a model that operates on 16×16 patches (4 patches total from a 32×32 image). The requirement is that accuracy on the test_patch_16.npz test set differs from clean test accuracy by less than 1%.
We discovered that robustness to patch shuffling cannot be learned from unshuffled data, it must be guaranteed architecturally. A model that has never seen shuffled data during training can only be invariant if its architecture makes the output independent of token order by construction.
This model uses the Deep Sets framework: each patch is encoded independently by a shared MLP, then the set of patch embeddings is aggregated using symmetric statistics before classification. Since symmetric functions produce the same output regardless of input order, permutation invariance is mathematically guaranteed.
Image → 4 patches (16×16×3 = 768 dims each)
→ PatchMLPEncoder (shared weights, applied to each patch independently):
Linear(768, 1024) → BatchNorm1d → GELU → Dropout(0.1)
Linear(1024,1024) → BatchNorm1d → GELU → Dropout(0.1)
Linear(1024, 512) → BatchNorm1d → GELU → Dropout(0.1)
→ PatchSetAggregator:
mean(tokens, dim=1) and std(tokens, dim=1) → concat → dim 1024
Linear(1024, 512) → BatchNorm1d → GELU → Dropout(0.1)
Linear( 512, 256) → BatchNorm1d → GELU → Dropout(0.1)
→ Linear(256, 10)
extractEmbedding(x) returns the 256-dimensional output of setAggregator (pre-classifier).
| Parameter | Value |
|---|---|
| Epochs | 200 |
| Batch Size | 1024 |
| Learning Rate | 3e-4 |
| L2 Regularization | 0.01 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.975 | 67.88% |
| Patch-16 | 0.975 | 67.88% |
| Patch-8 | 2.239 | 33.69% |
Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
The 0.00% gap between clean and patch-16 accuracy is exact and confirmed from epoch 1 throughout all 200 training epochs — it is architectural, not learned. The mean and standard deviation of a set of vectors are both symmetric functions: mean(permutation(T)) = mean(T) and std(permutation(T)) = std(T) for any permutation. Since the D-model uses 16×16 patches and the test file shuffles at 16×16 granularity, shuffling only rearranges the token order — the set itself is unchanged, so the model's output is identical.
The model is not invariant to 8×8 shuffling because a 16×16 patch has its internal content modified when 8×8 sub-regions within it are rearranged. This changes the input to the patch encoder, not just the token order — the set is genuinely different.
A notable result: this model achieves 67.88% clean accuracy, exceeding Net_FC (65.62%) despite using far fewer and coarser features (4 patches of 16×16 vs the full 3072-dimensional vector). The per-patch shared MLP forces the model to learn generalisable local representations before aggregating, which acts as an implicit structural regulariser that benefits clean accuracy.
RandAugment was removed from the training transform for this model (and kept for Net_FC). RandAugment applies photometric operations — solarize, posterize, equalize — that corrupt pixel values. For a 16×16 patch with 768 raw values, this corruption is disproportionately harmful compared to a full 32×32 image. RandomCrop and HorizontalFlip were retained as they do not corrupt patch content. Standard cross-entropy (no label smoothing) was used because the simpler architecture is not at risk of overconfident logits.
The architecture of the CNN is identical to the original Net_CNN with two caveats: Flatten is replaced by AdaptiveAvgPool2d((1,1)) to account for the batches, and the forward() method divides the training image into quadrants which are used as elements of a batch for training, as detailed in the Discussion section below. This architecture was refactored into a Shuffletruffle_CNN superclass for reuse with Net_N_shuffletruffle_CNN as well.
| Parameter | Value |
|---|---|
| Epochs | 100 |
| Batch Size | 128 |
| Learning Rate | 1e-2 |
| L2 Regularization | 0.00 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.985 | 80.63% |
| Patch-16 | 0.985 | 80.63% |
| Patch-8 | 1.428 | 59.63% |
Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
Various architectures were tried before settling on the one above. Initial attempts tried drastically reducing the feature space (both in the final or initial layers of convolution), increasing the feature space through reflection padding, and a custom averaging kernel that only considered its four corners. The best performing of these was the reflection padding, with Clean 69.84% accuracy vs 41.24% for Patch-16.
Since averaging the quadrants would remove any relative locality information, this method was further pursued. After consulting Gemini for ideas, it suggested treating the quadrants as distinct images within a batch instead of as parts of the image. The batch itself could then be averaged for the final output. This is the approach which I went with and which was refactored to allow for patches of various sizes.
That this relied on averaging the four quadrants against one another explains the impressively low accuracy gap. The average of a non-shuffled and shuffled image should be the same. However, the strong performance in the Patch-8 case as well is surprising since information is still scattered across its (relative) subpatches.
This model performs better than either the FC or ViT architectures, despite ViT performing better in Task 1. Moreover, it exceeds its own performance from Task 1 at 80.63% vs 71.59%. Given that locality information is lost in the averaging, perhaps subjects existed in different corners of the image, making learning generalized patterns difficult. I.e. off-center subjects became centered.
The positional embedding present in Net_Attention is removed entirely — if the model knows where a token came from, it cannot be invariant to token reordering. A CNN-based patch encoder (PatchCNNEncoder) replaces the linear projection to extract richer per-patch features. Mean pooling over all tokens replaces the CLS token for aggregation:
Image → 4 patches (16×16, kept as spatial tensors)
→ PatchCNNEncoder (shared, applied to each patch independently):
Conv2d(3, 384, 3×3, padding=1) → GELU
Conv2d(384, 384, 3×3, padding=1) → GELU
AdaptiveAvgPool2d(1) → squeeze → Linear(384, 768) → Dropout(0.1)
→ 10 × TransformerEncoderBlock(768, 8 heads, mlpDim=2048)
→ LayerNorm(768)
→ mean over 4 tokens → Linear(768, 10)
extractEmbedding(x) returns the 768-dimensional mean-pooled LayerNorm output.
| Parameter | Value |
|---|---|
| Epochs | 300 |
| Batch Size | 1024 |
| Learning Rate | 6e-4 |
| L2 Regularization | 0.02 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.980 | 77.66% |
| Patch-16 | 0.980 | 77.66% |
| Patch-8 | 1.669 | 56.72% |
Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
Removing positional embeddings and using mean pooling guarantees permutation invariance through the same mathematical argument: mean(permutation(tokens)) = mean(tokens). The gap holds to within 0.01% throughout training, converging to 0.00% as the model stabilises — confirmed across all 300 epochs in the training log.
With only 4 tokens, there is an information bottleneck — each token must carry a large amount of semantic content. This motivates the wider embedding dimension (768 vs 512 for the plain model), deeper transformer stack (10 layers vs 8), and wider MLP dimension (2048 vs 1024). The CNN patch encoder (two Conv3×3 layers + AdaptiveAvgPool) extracts local edge and texture features from each patch before the transformer processes them, providing richer token representations than a simple linear projection. AdaptiveAvgPool2d(1) collapses spatial dimensions to a single vector per patch without a large convolution. A naive alternative — Conv2d(384, 768, kernel_size=16, stride=16) — would have 384 × 768 × 16 × 16 + 768 = 75,498,240 parameters in that layer alone. The actual design uses AdaptiveAvgPool2d(1) followed by Linear(384, 768) = 295,680 parameters: a reduction of approximately 255× computed directly from the architecture parameters in main.py [2].
A lower weight decay (0.02 vs 0.05 for the plain model) was used because the information bottleneck at 4 tokens already limits overfitting — heavier regularisation would harm fitting capacity that is already constrained by design.
| Model | Clean Accuracy | Clean Loss | Patch-16 Accuracy | Patch-16 Loss | Gap (clean vs patch-16) |
|---|---|---|---|---|---|
D-shuffletruffle-FC |
67.88% | 0.975 | 67.88% | 0.975 | 0.00% |
D-shuffletruffle-CNN |
80.63% | 0.985 | 80.63% | 0.985 | 0.00% |
D-shuffletruffle-Attention |
77.66% | 0.980 | 77.66% | 0.980 | 0.00% |
The N-shuffletruffle has a smaller receptive field — each lens covers a finer region. We model this as a model that operates on 8×8 patches (16 patches from a 32×32 image). The requirement is that accuracy differs by less than 1% from clean on both the test_patch_8.npz test set. Crucially, the N-model must also be invariant to 16×16 shuffling, because a 16×16 shuffle is simply a permutation of groups of four adjacent 8×8 tokens — the set of 8×8 tokens is unchanged. The D-model must NOT share this same profile, which is guaranteed by using different patch sizes.
Identical structure to Net_D_shuffletruffle_FC but with 8×8 patches, producing 16 tokens instead of 4:
Image → 16 patches (8×8×3 = 192 dims each)
→ PatchMLPEncoder (shared):
Linear(192, 1024) → BatchNorm1d → GELU → Dropout(0.1)
Linear(1024,1024) → BatchNorm1d → GELU → Dropout(0.1)
Linear(1024, 512) → BatchNorm1d → GELU → Dropout(0.1)
→ PatchSetAggregator:
mean + std → concat (dim 1024)
Linear(1024, 512) → BatchNorm1d → GELU → Dropout(0.1)
Linear( 512, 256) → BatchNorm1d → GELU → Dropout(0.1)
→ Linear(256, 10)
extractEmbedding(x) returns the 256-dimensional output of setAggregator (pre-classifier).
| Parameter | Value |
|---|---|
| Epochs | 200 |
| Batch Size | 1024 |
| Learning Rate | 3e-4 |
| L2 Regularization | 0.01 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 0.878 | 70.62% |
| Patch-16 | 0.878 | 70.62% |
| Patch-8 | 0.878 | 70.62% |
Accuracy gap (clean vs patch-8): 0.00% Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
The N-model achieves identical loss (0.878) and accuracy (70.62%) on all three test sets — the same checkpoint, identical predictions regardless of whether patches are unshuffled, shuffled at 8×8, or shuffled at 16×16 granularity. This perfect tri-invariance arises because: (1) the 8×8 patch size means 8×8 shuffling only permutes tokens, and (2) a 16×16 shuffle simply permutes groups of four 8×8 tokens, which is still a permutation of the full 16-token set. Both are handled by the same symmetric aggregation.
The N-model (70.62%) outperforms the D-model (67.88%) on clean accuracy because 16 tokens carry substantially more spatial information than 4 tokens. The mean and standard deviation statistics computed across 16 samples are also more stable estimators than those computed from 4 samples, giving the aggregation step richer material to work with.
The architecture is identical to Net_D_shuffletruffle_CNN as it uses the same superclass, ShuffleTruffle_CNN, with patch_size=8.
| Parameter | Value |
|---|---|
| Epochs | 100 |
| Batch Size | 128 |
| Learning Rate | 1e-2 |
| L2 Regularization | 0.00 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 1.097 | 75.39% |
| Patch-16 | 1.097 | 75.39% |
| Patch-8 | 1.097 | 75.39% |
Accuracy gap (clean vs patch-8): 0.00% Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
As with the D_shuffletruffle_CNN, there is no longer an accuracy discrepancy between the shuffled and un-shuffled images. Notably, the performance is nearly on par with Net_N_shuffletruffle_Attention, <1.50%.
Same design philosophy as Net_D_shuffletruffle_Attention but with 8×8 patches, yielding 16 tokens:
Image → 16 patches (8×8, kept as spatial tensors)
→ PatchCNNEncoder (shared):
Conv2d(3, 256, 3×3, padding=1) → GELU
Conv2d(256, 256, 3×3, padding=1) → GELU
AdaptiveAvgPool2d(1) → squeeze → Linear(256, 512) → Dropout(0.1)
→ 8 × TransformerEncoderBlock(512, 8 heads, mlpDim=1024)
→ LayerNorm(512)
→ mean over 16 tokens → Linear(512, 10)
extractEmbedding(x) returns the 512-dimensional mean-pooled LayerNorm output.
| Parameter | Value |
|---|---|
| Epochs | 300 |
| Batch Size | 1024 |
| Learning Rate | 6e-4 |
| L2 Regularization | 0.02 |
| Test Set | Loss | Accuracy |
|---|---|---|
| Clean | 1.036 | 76.80% |
| Patch-16 | 1.036 | 76.80% |
| Patch-8 | 1.036 | 76.80% |
Accuracy gap (clean vs patch-8): 0.00%
Accuracy gap (clean vs patch-16): 0.00%
Training Loss
Accuracy
The N-Attention model achieves identical results across all three test sets, confirming complete permutation invariance through the same reasoning as the N-FC model. It uses a narrower encoder (innerDim 256 vs 384 for D-model) because 8×8 patches are spatially smaller — there are fewer local features to extract from a 64-pixel region — and the 16-token information budget makes up for any per-token compression.
The D-model's accuracy on patch-16 (77.66%) and D-model's accuracy on patch-8 (56.72%) differ, while the N-model achieves the same accuracy on all three sets (76.80%). This exactly satisfies the assignment's requirement that D and N models have different accuracy profiles across the two shuffled test sets. The difference arises from the fundamental architecture choice: D uses 16×16 patches and is structurally invariant only to 16×16 shuffling, while N uses 8×8 patches and is structurally invariant to both.
| Model | Clean Accuracy | Clean Loss | Patch-8 Accuracy | Patch-8 Loss | Gap (clean vs patch-8) |
|---|---|---|---|---|---|
N-shuffletruffle-FC |
70.62% | 0.878 | 70.62% | 0.878 | 0.00% |
N-shuffletruffle-CNN |
75.39% | 1.097 | 75.39% | 1.097 | 0.00% |
N-shuffletruffle-Attention |
76.80% | 1.036 | 76.80% | 1.036 | 0.00% |
All 9 models evaluated on all 3 test sets (test accuracy % / test loss):
| Model | Clean Accuracy | Clean Loss | Patch-16 Accuracy | Patch-16 Loss | Patch-8 Accuracy | Patch-8 Loss |
|---|---|---|---|---|---|---|
Net_FC |
65.62% | 1.129 | 29.69% | 2.139 | 21.81% | 2.377 |
D-shuffletruffle-FC |
67.88% | 0.975 | 67.88% | 0.975 | 33.69% | 2.239 |
N-shuffletruffle-FC |
70.62% | 0.878 | 70.62% | 0.878 | 70.62% | 0.878 |
Net_CNN |
88.02% | 0.631 | 55.84% | 1.430 | 34.58% | 2.192 |
D-shuffletruffle-CNN |
80.63% | 0.985 | 80.63% | 0.985 | 59.63% | 1.428 |
N-shuffletruffle-CNN |
75.39% | 1.097 | 75.39% | 1.097 | 75.39% | 1.097 |
Net_Attention |
88.24% | 0.677 | 55.16% | 1.894 | 50.92% | 2.059 |
D-shuffletruffle-Attention |
77.66% | 0.980 | 77.66% | 0.980 | 56.72% | 1.669 |
N-shuffletruffle-Attention |
76.80% | 1.036 | 76.80% | 1.036 | 76.80% | 1.036 |
25 images were selected from the original CIFAR-10 test set (indices 0–24, labeled 1–25 in the plots). For each image, a 16×16-shuffled version and an 8×8-shuffled version were created using the same shufflePatches() function used in training, yielding a 75-example dataset. All three versions of each image retain the original class label.
For each available model, extractEmbedding() was called on all 75 examples to produce a matrix of shape 75×embeddingDim. The three CNN models were skipped as their checkpoints are not yet available. PCA with n_components=2 was applied to project each matrix into 2D. The 75 points are plotted as a scatter plot coloured by class label (10 CIFAR-10 classes).
The embedding dimensions per model are:
| Model | Embedding Dim |
|---|---|
Net_FC |
256 |
D-shuffletruffle-FC |
256 |
N-shuffletruffle-FC |
256 |
Net_CNN |
256 |
D-shuffletruffle-CNN |
256 |
N-shuffletruffle-CNN |
256 |
Net_Attention |
512 |
D-shuffletruffle-Attention |
768 |
N-shuffletruffle-Attention |
512 |
The PCA explained variance numbers (PC1 + PC2 combined) from our run are:
| Model | PC1 | PC2 | Total |
|---|---|---|---|
Net_FC |
0.131 | 0.124 | 25.5% |
D-shuffletruffle-FC |
0.185 | 0.119 | 30.4% |
N-shuffletruffle-FC |
0.176 | 0.151 | 32.7% |
Net_CNN |
0.175 | 0.110 | 28.5% |
D-shuffletruffle-CNN |
0.293 | 0.195 | 48.8% |
N-shuffletruffle-CNN |
0.328 | 0.187 | 51.5% |
Net_Attention |
0.203 | 0.143 | 34.6% |
D-shuffletruffle-Attention |
0.160 | 0.118 | 27.8% |
N-shuffletruffle-Attention |
0.213 | 0.147 | 36.0% |
Plain models (Net_FC, Net_Attention, Net_CNN):
Some clustering can be observed within each model depending on image class. This is most readily apparent in the Net_Attention results as three distinct regions are formed. The image variants, however, are generally not clustered, with a few light exceptions such as the Original images for the classes airplane and ship under Net_Attention. This, alongside the total percentages, implies that the models are generally learning rich feature embeddings as reduction to two dimensions does not adequately distinguish the classes.
D-shuffletruffle models:
In both the FC and Attention variants, the circle (original) and triangle (patch-16) markers from the same image coincide at the same point — they are not just close, they are identical to floating point precision. This directly visualises the architectural guarantee: the model computes the same embedding regardless of 16×16 token order, so the two scatter onto exactly the same 2D coordinate. The square (patch-8) markers scatter to different positions because 8×8 shuffling modifies the content inside each 16×16 patch, not just the token order. D-shuffletruffle-Attention has lower explained variance (27.8%) than Net_Attention (34.6%), reflecting the information bottleneck of operating on only 4 tokens.
With the CNN model, however, this overlap does not exist. This indicates that the model, as in the various experiments, is still learning to distinguish the shuffled and un-shuffled image variants. That this has not affected the overall accuracy at all, alongside the high variance explained by the two principle components, indicates that there may be an error in the embedding extraction for the CNNs. The average of the patches should have removed any locality information that would have allowed the two variants to be distinguished, but that is plainly not evident in the generated graphs.
N-shuffletruffle models: All three markers from the same image — circle, triangle, and square — coincide at a single point. This is the strongest form of invariance: the model produces identical embeddings whether it sees the original image, a 16×16-shuffled version, or an 8×8-shuffled version. N-shuffletruffle-Attention achieves the highest explained variance of all six models (36.0%), suggesting that its 16-token mean-pooled representation produces the most structured and class-separable embedding space. The class clusters in the N-Attention PCA plot are visibly tighter and better-separated than in any other model, while each image still appears as a single point regardless of which shuffle variant was used.
As with the D-shuffletruffle-CNN, there is generally clear class clustering but not a collapsing of the variants into a single component. The explained variance, meanwhile, has jumped to 51.5%, the highest of any of the models tested. As before, due to the averaging of the patches against each other, the shuffled and un-shuffled images should be interpreted identically. Thus, there is likely an issue with the final embedding extraction. Alternatively, since the patches are interpeted as independent images within a batch, it is possible that during training the order of the patches may somehow be learned by the model. I do not see how it could be architecturally, but if so that would explain the discrepancy since shuffled and un-shuffled images would have patches in differnet orders within a batch.
The PCA plots confirm that permutation invariance is not merely a test-set accuracy statistic — it is a geometric property of the embedding space. The invariant models literally collapse three input variants onto one output point, while the non-invariant models scatter them across the space.
The fundamental challenge is achieving high clean accuracy while maintaining permutation invariance. Several architectural directions could push beyond our current results:
Bag-of-words CNN with global average pooling: A CNN that uses small convolutional filters operating within patch boundaries, followed by global average pooling rather than spatial max pooling or fully connected layers. Global average pooling computes the spatial mean of each feature map, making it insensitive to the positions of the detected features. The convolutions would still respect patch boundaries, maintaining invariance while leveraging the feature extraction power of convolutional operations.
Capsule Networks with dynamic routing: Capsule networks group neurons into capsules that represent entities and their properties rather than scalar feature activations. If capsules are defined at the patch level with routing that is order-agnostic, the network could in principle learn richer patch representations than a flat MLP while remaining invariant to patch permutations.
Deeper Deep Sets with cross-patch interaction via attention but without positional encoding: Our current Deep Sets models process patches independently before aggregation. Adding a step of cross-patch attention (without positional encodings) before mean pooling would allow patches to condition their representations on what other patches look like — providing the interaction benefits of attention while preserving invariance through the final orderless pooling.
Equivariant neural networks: Networks designed using the theory of group-equivariant convolutions can be constructed to be exactly invariant to specific symmetry groups, including permutation groups. This is a more principled approach to invariance than architectural tricks and could achieve both higher accuracy and provable invariance simultaneously.
The bulk of the code was done by Dhruv. He created the FC and Attention models for all three tasks. The automated training and PCA code were written by him too.
The CNN portions were done by Lucas. He drafted the core architecture and adapted it for the N_ and D_shuffletruffle variants.
Each author wrote the portion of this README corresponding to their respective models. Comparative analysis between models was done by each as appropriate. Additional writing, including references, was written by Dhruv.
Each attests the accuracy of this statement of contributions below.
[x] Dhruv Gupta
[x] Lucas Stewart
[1] Preliminary run (Job 6639579, Big Red 200, Sun Mar 15 2026) — Net_Attention with linear patch embedding and ReduceLROnPlateau scheduler, no warmup. Logged val accuracy 75.94% at epoch 101 (early stopped). Training log available in logs/Vit_1.out.
[2] Parameter count derived from main.py, class PatchCNNEncoder: innerDim = embeddingDim // 2 = 384, patchSize = 16. Naive alternative Conv2d(384, 768, 16, 16): 384 × 768 × 16 × 16 + 768 = 75,498,240 params. Actual design Linear(384, 768): 384 × 768 + 768 = 295,680 params. Ratio: 75,498,240 ÷ 295,680 ≈ 255.3×.
[3] Zaheer, M., Kottur, S., Ravanbhakhsh, S., Póczos, B., Salakhutdinov, R., & Smola, A. J. (2017). Deep Sets. Advances in Neural Information Processing Systems, 30. — Theoretical foundation for the permutation-invariant aggregation used in Net_D_shuffletruffle_FC and Net_N_shuffletruffle_FC.
[4] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Zhai, X., Unterthiner, T., Dehghani, M., ... & Houlsby, N. (2020). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021. — Foundation for Net_Attention, Net_D_shuffletruffle_Attention, and Net_N_shuffletruffle_Attention.
[5] Cubuk, E. D., Zoph, B., Shlens, J., & Le, Q. V. (2020). RandAugment: Practical Automated Data Augmentation with a Reduced Search Space. CVPR 2020 Workshops. — Data augmentation strategy used for FC and Attention plain models.
[6] Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019. — AdamW optimiser used across all models.