DEV Community: Mushahid Intesum

Can you run a GNN without anyone seeing the data? Journal of our experiments on privacy-preserving GNN inference on microcontrollers

Mushahid Intesum — Mon, 10 Aug 2026 02:56:24 +0000

I have been keenly interested in how to make Graph Neural Networks (GNN) run securely on resource-constrained devices. The question driving this work is deceptively simple: can three microcontrollers jointly compute GNN inference on traffic data such that no single device ever sees the plaintext, not the input features, not the model weights, not the intermediate activations? This post is about the findings I have found so far in this independent research endeavor.

Background on Secure GNN Inference

A standard GCN layer [1] computes: H' = σ(A · H · W + b), where A is the normalized adjacency matrix (graph topology), H is the node feature matrix, W is the learned weight matrix, and σ is an activation function. In a traffic signal coordination setting, A encodes which intersections connect to each other, H holds the per-intersection sensor readings (queue lengths, phase states, flow rates), and the output dictates signal timing decisions.

The privacy concern is real: traffic flow data reveals movement patterns, congestion bottlenecks, and potentially individual vehicle trajectories. If multiple jurisdictions or agencies need to collaboratively optimize signals at their shared borders, no party may want to expose their raw data to the others. Secure Multi-Party Computation (MPC) lets them jointly compute the model output without any single party seeing the full picture.

We use Replicated Secret Sharing (RSS) [2] in a 3-party semi-honest setting. Each value x is split into three shares (x₁, x₂, x₃) such that x₁ + x₂ + x₃ = x mod 2⁶⁴. Party i holds two of the three shares (xᵢ, xᵢ₊₁), so any single party sees only random-looking numbers. Reconstruction requires combining shares from at least two parties.

┌─────────────────────────────────────────────────────┐
│             REPLICATED SECRET SHARING               │
│                                                     │
│  Secret x split into (x₁, x₂, x₃)                 │
│  where x₁ + x₂ + x₃ = x  (mod 2⁶⁴)               │
│                                                     │
│  Party 0 holds: (x₁, x₂)                           │
│  Party 1 holds: (x₂, x₃)                           │
│  Party 2 holds: (x₃, x₁)                           │
│                                                     │
│  → No single party can recover x                    │
│  → Secure addition is FREE (no communication)       │
│  → Secure multiplication costs 1 round of comm      │
└─────────────────────────────────────────────────────┘

The arithmetic properties are what make this practical. Secure addition of two shared values requires zero communication: each party just locally adds their shares. Secure multiplication is more expensive. It requires one round of communication between parties using a PRF-based resharing protocol to maintain the share structure, but the cost is bounded and predictable.

The Key Insight: Public Adjacency, Private Features

Here's the critical observation that makes GNN inference on microcontrollers feasible. In a traffic network, the graph topology is public knowledge. Everyone knows which intersections connect to which. The adjacency matrix A is not secret. Only the node features H (sensor readings) and the model weights W (the learned policy) need protection.

This means the message-passing step A · H is a public-matrix × secret-matrix multiply, which requires zero communication. Each party can locally compute it from their shares alone. Only the feature transformation step H · W (secret × secret) requires the full RSS multiplication protocol.

Compared to CryptGNN [3] which treats everything as secret, this roughly halves the communication cost. In a 2-layer GCN, we have:

Layer 1: A·X (free) → (A·X)·W₁ (1 round) → activation (1 round)
Layer 2: A·H (free) → (A·H)·W₂ (1 round)
Total:   3 communication rounds (vs 5+ if A were secret)

Test Bed

The target hardware is ESP32-S3 microcontrollers, 240MHz dual-core Xtensa processors with WiFi and ~320KB usable SRAM. Three boards form the 3-party RSS computation, communicating over TCP/WiFi. For the model, I use a 2-layer GCN trained on a 20-node synthetic traffic grid (4×5 intersection layout):

4×5 Grid Topology:        Model Architecture:

○─○─○─○─○                 GCN(16 → 16 → 4)
│ │ │ │ │                  • 20 nodes, 16 input features
○─○─○─○─○                 • 16 hidden features, 4 output classes
│ │ │ │ │                  • Polynomial activation: 0.1x² + 0.5x + 0.1
○─○─○─○─○                 • 340 total parameters
│ │ │ │ │                  • Q20 fixed-point arithmetic
○─○─○─○─○

The input features per node are: one-hot phase encoding (4 dims), queue lengths per direction (4 dims), flow rates (4 dims), occupancy, and time-of-day (sin/cos encoding), 16 features total. The model predicts which direction has the highest traffic demand.

An important design choice: we use a polynomial activation 0.1x² + 0.5x + 0.1 instead of ReLU. ReLU requires a comparison (x > 0), which is expensive in MPC because comparisons need bit-decomposition. A degree-2 polynomial can be evaluated using only one secure multiplication (for the x² term), keeping the activation cost to a single communication round. Training with the polynomial activation achieves 52.1% accuracy vs 54.0% with ReLU, a modest 1.9 percentage point drop that we consider acceptable for the massive reduction in protocol complexity.

The Pipeline

The full system goes from training through to secure inference verification:

Train GCN → Export Q20 Fixed-Point → Generate RSS Shares → Secure Inference → Verify Against Plaintext

The Python side handles training (train_model.py), fixed-point export with RSS share generation (export_model.py), and offline analysis (analyze.py + rss_emulator.py). The C side implements the actual protocol: RSS primitives, secure matrix operations, and the 2-layer GCN inference pipeline. Both sides are designed to produce bit-identical results for a given set of shares.

Finding 1: Fixed-point quantization preserves inference fidelity

Before worrying about secret sharing, the first question is whether the float-to-fixed-point conversion introduces unacceptable error. We use Q16 fixed-point (16 fractional bits) for the offline analysis with an 8-node ring graph as a controlled test bed.

Node	Float Output	Q16 Output	Absolute Error
0	0.293579	0.293488	0.000091
1	0.273324	0.273224	0.000100
2	0.276208	0.276123	0.000085
3	0.302083	0.302002	0.000081
4	0.328208	0.328125	0.000083
5	0.354583	0.354492	0.000091
6	0.357505	0.357422	0.000083
7	0.337120	0.337036	0.000084

The mean quantization error is 0.0000875 and the max is 0.0001002. Both output features show identical values per node (the test model uses uniform weights that produce symmetric outputs), confirming the fixed-point pipeline is numerically faithful. All errors are sub-0.01% relative to the output magnitudes, well within acceptable bounds for traffic signal decisions.

Figure 1: Output comparison across nodes showing near-identical float and Q16 fixed-point values. The bars are visually indistinguishable, confirming minimal quantization degradation.

The error is also spatially uniform across nodes. This matters: if quantization error concentrated at specific graph positions, it could systematically bias decisions at those intersections.

Figure 2: Per-node quantization error heatmap. Error values are uniformly distributed across all nodes and output features, staying within the 0.00008 to 0.0001 range.

Finding 2: RSS primitives are bit-exact

The RSS protocol correctness was verified through four escalating tests in the Python emulator:

Test	Description	Result
Share/Reconstruct	Split 42, reconstruct from shares	✅ PASS
Secure Addition	100 + 200 = 300, zero communication	✅ PASS
Secure Multiplication	7 × 6 = 42, with PRF resharing	✅ PASS
Random Multiplication	100 random pairs, full ring arithmetic	100/100 ✅

The last test is the important one. We generate 100 random integer pairs in [0, 10000), secret-share both operands, perform secure multiplication with the full PRF-based resharing protocol, reconstruct, and verify against the plaintext product modulo 2⁶⁴. All 100 pass, confirming that the resharing logic correctly maintains the share invariant x₁ + x₂ + x₃ = x mod 2⁶⁴ through multiplication.

Finding 3: All three parties agree on the output

The C desktop emulator spawns three processes (one per party) communicating over localhost TCP sockets. Each party holds its own shares of the features and weights, runs the full GCN inference protocol (5 steps: public matmul, secure matmul, polynomial activation, public matmul, secure matmul), and then opens the output shares to reconstruct the final result.

All 3 parties produce identical reconstructed outputs, confirming protocol correctness across the inter-process communication boundary. The desktop emulation completes in approximately 0.39 ms average across the three parties (Party 0: 0.43 ms, Party 1: 0.39 ms, Party 2: 0.35 ms).

Figure 3: Inference time per party in the desktop emulator. The slight variation is expected from OS scheduling; the important thing is that all three parties complete within a tight band.

The per-party timing spread (~0.08 ms) is purely from OS scheduling jitter on the desktop. On dedicated ESP32-S3 hardware without a preemptive OS, we expect tighter synchronization.

Finding 4: The secure pipeline mirrors the C implementation

The Python RSS emulator (rss_emulator.py) replicates the exact same protocol as the C implementation:

Step 1: A * X      ← public × shared, FREE
Step 2: (A*X) * W₁ ← shared × shared, 1 communication round
Step 3: Activation  ← polynomial, 1 communication round
Step 4: A * H₁     ← public × shared, FREE
Step 5: (A*H₁) * W₂ ← shared × shared, 1 communication round

Both implementations use identical truncation semantics (arithmetic right shift on signed 64-bit values), the same PRF-based resharing (splitmix64 hash for deterministic randomness), and the same polynomial activation coefficients (Q-format 0.1, 0.5, 0.1). This dual-stack approach, Python for rapid prototyping and verification, C for deployment, lets us catch protocol bugs early. If the Python emulator and the C binary produce different outputs for the same shares, something is wrong in the C implementation.

The Firmware Side

The ESP32-S3 firmware is written in C with ESP-IDF. Each board connects to the others over WiFi/TCP, loads its pre-generated share files (model_shares_p{0,1,2}.bin), and runs the inference protocol. The PRF uses mbedtls hardware-accelerated AES on the ESP32-S3 (replacing the OpenSSL backend used on desktop). The transport layer abstracts over POSIX sockets (desktop) vs lwIP TCP (ESP32), so the core protocol code is identical on both platforms.

This is fully built and compiles, but the actual on-device deployment and benchmarking is the next step. This will be the next part.

Future Directions

This is an ongoing experiment. This blog serves as a journal of my progress so far and will continually be updated once new findings are found. A few concrete next steps:

Hardware deployment and benchmarking: flash the firmware to 3 ESP32-S3 boards, run the full protocol over WiFi, and measure end-to-end latency, per-step communication overhead, and memory usage. The key question: can the protocol complete within a traffic signal cycle (~1-5 seconds) to be practically useful? This will be the next part.
Scale to the full 20-node model: the current desktop verification uses an 8-node ring graph with small weight matrices. The trained traffic model has 20 nodes with 16→16→4 dimensions, significantly larger matrices that will stress both computation and communication budgets on the ESP32.
Real traffic data integration: replace the synthetic training data with real-world traffic signal datasets to evaluate whether the fixed-point quantization error affects decision quality on non-synthetic distributions.
Comparison with CryptGNN: quantitatively benchmark our "public adjacency" optimization against the fully-secret approach to measure the actual communication savings in practice.

Hopefully we'll be able to get a paper through this independent endeavor.

References

[1] T. N. Kipf, M. Welling, "Semi-Supervised Classification with Graph Convolutional Networks," ICLR 2017.
[2] T. Araki, J. Furukawa, Y. Lindell, A. Nof, K. Ohara, "High-Throughput Semi-Honest Secure Three-Party Computation with an Honest Majority," CCS 2016.
[3] R. Ran, W. Wang, Q. Gang, J. Rao, "CryptGNN: Fast Privacy-Preserving Graph Neural Network Inference," 2022.
[4] ESP-IDF Programming Guide, Espressif Systems, https://docs.espressif.com/projects/esp-idf/.

Can you steal a robot's next move by watching its clock? Journal of our experiments on timing side channels in multi agent RL

Mushahid Intesum — Mon, 10 Aug 2026 02:46:47 +0000

I have been keenly interested in the intersection of multi-agent reinforcement learning (MARL) and hardware security. When you deploy a trained RL policy onto a microcontroller, the model runs inference to decide what action to take. But here's the thing: different actions can take different amounts of time to compute. If an adversary can measure that timing, can they figure out what the agent is about to do without ever seeing the input? This post is about the findings I have found so far in this independent research endeavor.

Background on the Threat Model

The core question is straightforward. A MARL policy, say two cooperative agents navigating a grid, runs on an ESP32-S3 microcontroller. The attacker sits on the outside. They can measure:

Total inference duration (and per-operator breakdowns)
Number of inferences per timestep
Network packet timing if WiFi is involved

They cannot see the raw observations fed to the policy, the internal activations, or the weights. Pure black-box timing. The goal: predict the agent's action from timing alone, without ever seeing the observation input.

┌─────────────────────────────────────────────────┐
│                   ATTACKER                      │
│  Can observe:                                   │
│    • Inference duration (total + per-layer)      │
│    • Number of inferences per timestep           │
│    • Network packet timing (if WiFi used)        │
│  Cannot observe:                                │
│    • Raw observations fed to the policy          │
│    • Internal activations or weights             │
│    • Source code (black box timing only)           │
│                                                 │
│  Goal: Predict the agent's ACTION from timing   │
│        without seeing the observation input      │
└─────────────────────────────────────────────────┘

This threat model matters because MARL is increasingly being deployed on edge devices like autonomous drones, warehouse robots, and cooperative IoT networks where an adversary with physical proximity could realistically tap into timing signals. If the leakage is real, it's a serious vulnerability: you can predict what an agent will do before it does it.

Test Bed

For the environments, I set up three test beds with increasing complexity:

1. Cooperative Grid Navigation (Custom, Primary): Two agents on a 5×5 grid with swapped goals. They communicate their position and intended direction. There's a wall with a single gap in the middle, so they need to coordinate to avoid collision. The policy is a small MLP: MLP(9 → 32 → 32 → 5) with 5 discrete actions (stay, up, down, left, right).

A . . . .     2 agents, 5×5 grid, swapped goals
. . . . .     Comm: position + intended direction
# # . # #     Shared reward, collision penalty
. . . . .     
. . . . B     Policy: MLP(9 → 32 → 32 → 5)

2. CartPole (Gymnasium Baseline): The classic single-agent balance task. This serves as a sanity check. It establishes whether timing leakage is a general phenomenon of NN inference on MCUs or specific to multi-agent setups. Policy: MLP(4 → 32 → 32 → 2).

3. MPE Simple Spread (PettingZoo): 3 agents cooperatively covering 3 landmarks. A standard MARL benchmark with a larger policy: MLP(18 → 64 → 64 → 5). This one hasn't been fully analyzed yet.

For both grid_nav and CartPole, the policies are trained using PPO (and IPPO for multi-agent), then exported to TFLite (both FP32 and INT8 quantized variants) for deployment. The export pipeline goes PyTorch → Keras → TFLite, and also generates a .h C header for embedding directly into the ESP32 firmware.

The Pipeline

The full analysis pipeline I built goes from training all the way through to leakage quantification:

Train Policy → Export TFLite → Flash ESP32-S3 → Collect Timing → Analyze Leakage → Generate Report

The analysis module is the core of the project. For each model configuration, it computes:

Mutual Information MI(timing; action): how many bits of action information leak through timing, with bootstrap confidence intervals
Statistical Tests: Kruskal-Wallis H test and one way ANOVA to determine if timing distributions across actions are statistically different
Effect Sizes: Cohen's d between all action pairs to quantify how different the distributions are
Classifier Accuracy: a Random Forest and an MLP trained to predict the action from timing features alone (total cycles + per op breakdown)

If the classifiers beat random baseline by a significant margin, the system is operationally vulnerable: an attacker can build the same model.

Finding 1: Simulated timing shows the pipeline detects leakage when present

Before touching real hardware, I validated the pipeline with simulated timing. The SimulatedCollector runs real TFLite inference on the laptop but generates synthetic cycle counts that are intentionally action dependent: base_cycles + action × bias + noise. This confirms that when there is leakage baked in, the analysis correctly detects it.

The results across 2000 simulated traces per configuration:

Config	MI (bits)	RF Accuracy	KW p-value	Max Cohen's d
cartpole_fp32	0.011	99.5%	1.98e-9	0.274
cartpole_int8	0.022	99.8%	1.71e-9	0.276
grid_nav_fp32	0.976	99.8%	1.46e-281	8.879
grid_nav_int8	1.005	100.0%	1.15e-286	8.838

A few observations jump out immediately:

The grid_nav environment leaks significantly more than CartPole. Grid navigation has 5 actions with MI near 1 bit (out of a max of ~2.32 bits for 5 actions), while CartPole's MI hovers near 0. This makes sense: the simulated timing injects action × 200 cycles of bias, and with 5 spread out actions, there's more room for distinct timing signatures than CartPole's 2 actions.

Classifiers hit near perfect accuracy. Both RF and MLP classifiers achieve 99.5%+ accuracy on the simulated data, far exceeding the 20% random baseline for 5 actions and 50% for 2 actions. This demonstrates that even simple models can operationally exploit timing leakage.

The Kruskal-Wallis p values are absurdly small. Values like 1.46e-281 are effectively zero. The null hypothesis (timing distributions are the same across actions) is annihilated. Combined with Cohen's d values above 8 for grid_nav (anything above 0.8 is conventionally "large"), this signals massive, unmistakable effect sizes.

INT8 quantization doesn't suppress leakage. In fact, the INT8 variants show slightly higher MI and comparable classifier accuracy. This is notable because quantization changes the computational profile of each operator, and we had hypothesized it might reduce timing variance by simplifying the computation paths. At least in simulation, it doesn't.

Finding 2: Per operator timing is where the signal lives

Beyond total inference time, I also measured per operator cycle counts (5 simulated operators: two FullyConnected, two activation, one Softmax). The per operator MI analysis revealed that the leakage isn't uniform. Specific operators (the FullyConnected layers and the final Softmax) carry most of the timing signal, while activation functions contribute very little.

This has practical implications: if you wanted to defend against timing side channels, you'd focus your constant time countermeasures on the matmul and softmax operators specifically, rather than trying to make the entire inference pipeline constant time (which is much harder).

Finding 3: Observation timing correlation exists

Using Spearman correlation between individual observation dimensions and total timing, I found statistically significant correlations in multiple dimensions. This means timing doesn't only leak the output action, it also leaks information about the input observations. In a MARL context, this is arguably worse: an adversary could potentially reconstruct features of the agent's perceived state, not just its decision.

The Firmware Side

The ESP32-S3 firmware is written in C/C++ with ESP-IDF and TFLite Micro. The key features:

Cycle accurate timing via esp_cpu_get_cycle_count() at 240MHz (4.17ns resolution)
Per operator profiling using TFLite's MicroProfiler to attribute timing to individual operators (FullyConnected, ReLU, Softmax)
Binary UART protocol at 921600 baud for high throughput data collection: the laptop sends observation vectors and the ESP32 sends back (action, total_cycles, per_op_cycles[])

This is fully built and compiles, but the actual on device data collection is the next step. This will be covered in the next part.

Future Directions

This is an ongoing experiment. This blog serves as a journal of my progress so far and will continually be updated once new findings are found. A few concrete next steps:

Real hardware timing collection: flash the firmware to the ESP32-S3, collect 10K+ traces per environment, and compare against the simulated results. The key question: does actual hardware exhibit exploitable timing variance, or does the ESP32-S3's SIMD/vector acceleration (ESP-NN) make inference sufficiently constant time? This will be the next part.
Constant time countermeasures: if leakage is confirmed on hardware, experiment with software mitigations: padding inference to a fixed duration, adding random delays, or using constant time matmul implementations.
Cross environment comparison: run the full analysis on Simple Spread (3 agents, larger policy) to see if leakage scales with model size or action space complexity.
Adversarial observation sampling: instead of uniform random observations, use adversarially chosen inputs that maximize timing variance to establish worst case leakage bounds.

Hopefully we'll be able to get a paper through this independent endeavor.

Experiments on Token Transformation of Vision Mamba

Mushahid Intesum — Wed, 06 Aug 2025 17:55:14 +0000

State Space Models such as Mamba have become a promising alternative to attention mechanism, which is difficult to compress due to the quadratic complexity of the attention mechanism. For handling image data, models such as ViM and VMamba have been proposed. Though Mamba has been shown to perform well for long contexts, there is still a lot of understanding to be had on the mechanisms of this model. In this study, I attempt to run some experiments on vision token transformation for the ViM model to see how it performs and try find any possible improvements. This post will serve as a journal of the progress of this independent study.

Background

Attention Mechanism of Transformer

Transformers took the world by storm for its unmatched ability in handling long contexts and scalability. Modern Large Language Models such as ChatGPT, Claude, Gemini, DeepSeek have been built using transformers. These models have billions of parameters, and this can be done because transformers can be very easily parallelized as text is not parsed sequentially, unlike previous seq2seq models. At its core is the self-attention mechanism presented in the following equation.

Here, the input is split into Q, K, V matrices and passed through the above formula to get the attention score for the input. Multiple such attention scores are calculated for each input sequence in a process called multi-headed attention. This provides robust outputs and responses.

Transformers have been adapted to vision tasks by dividing the images into a set number of tokens or patches, inserting positional information and then passing them to the transformer layers. A CLS token is added for identifying which class the image belongs too when needed.

For all its ability of scalability, transformers are notoriously difficult to compress because of the same self-attention mechanism. The quadratic complexity makes it hard for parsing long input sequences under resource constraints. Many optimizations have been proposed to address this.

State Space Models & Mamba

State Space Models (SSM) are a type of mathematical formulation which describes the state of an object in a continuous space with the least number of parameters. It is described using the following equation:

The first equation describes how the system's internal state evolves over time, while the second equation relates the system's internal state to its observable outputs. Neural networks using SSM make matrices A, B and C shared across the network. Matrix D can be thought of as a skip connection and is usually ignored. To make this equation usable for a neural network, the matrices are discretized by introducing another variable delta and transforming the matrices in the following way:

These shared weights allows us to pre-calculate a portion of the equation and then reuse over and over again. Similar to convolutions, the pre-calculated value can be used to calculate proceding representations efficiently. This allows us to save computation costs. This can be visualized using the following figure:

Neural networks using SSMs have been proposed and were shown to be more efficient than transformers, but their performance was not good enough to be considered as a replacement for it. This is because of the static nature of the A, B and C matrices. The model cannot focus on relevant parts of the input like transformers and thus struggles when sequences are longer. The introduction of Mamba changed that notion by addressing the performance issue by making SSMs selective. This notion of selectivity is introduced by making matrices A and C input dependent. That is, each token will have its own A and C matrices instead of having a shared one. This would allow the model to focus on specific input tokens, just like transformers. Though this would not allow us to use the kernel trick, this was addressed by designing a hardware aware algorithm. Mathematical calculations are less time-consuming than memory movement. That is, it takes more time to more data from high-bandwidth SRAM to low bandwidth DRAM and vise versa. In order to make

Mamba for Vision

Different models have been proposed for handling image data using Mamba such as VMamba and ViM. For my experiments, I will be utilizing VMamba. Similar to vision transformers, VMamba divides the image into patches and then the patches are fed through a linear block to get intermediate representations. To properly get positional information, 4-directional SSM is done so that model can get positional information of all the patches properly.

Methodology

This experiment is inspired by Famba-V, which merges tokens after passing through the forward and backward SSM blocks, as can be seen in the following image.

After the tokens pass through the SSM blocks, similar tokens are merged using cosine similarity and a new representation matrix is formed. This allowed for lower training time and memory usage.

Instead of merging, we aim to transform the tokens. As proposed in this paper, just simply pruning or merging vision tokens can lead to relevant information loss. Instead of doing either one, the vision tokens would be transformed using a transformation matrix.

The construction of this transformation matrix first involves calculating attention scores using the attention equivalence scores for mamba models as formulated in this paper, which is:

Afterwards, the transformation matrix is created in the following order:

The attention score of each token with respect to others is summed to create a vector of aggregated attention scores.
Another matrix of size (patch length * patch length) is created by taking the cosine similarity of each token
This new matrix is then softmaxed.

After this, the outputs from the selective scan operation is multiplied with this matrix to get the new latent representations.

This token transformation can be visualized from the following image:

The inner workings of the token transformation can be visualized in the following way:

Potential Pitfalls:

Calculation of this attention matrix involves a quadratic calculation. This might become a bottleneck. In order to alleviate this, we could utilize the linear attention equivalence proposed in EfficientViT.

Conclusion:

This is an ongoing blog and, as I mentioned earlier, is going to act as a journal of whatever updates I have on this project.

Journal of our experiments on VLM token pruning

Mushahid Intesum — Sat, 02 Aug 2025 04:14:39 +0000

I and @oldpilluwu have been keenly interested in how to make Large Vision Models (VLM) work and wanted to know if these really large models can be made more efficient. Because, let's face it, the current AI landscape is dominated by these massive models which can be properly utilized through API calls and using these API can get really expensive really quickly. For that, we started poking into Llava-7B [1] to observe it's characteristics and try to find any possible avenues for efficiency and redundancy removal. This post is about the findings we have found so far in this independent research endeavor.

Background on Llava Model:

Llava[1] is a family of vision language models introduced in 2023 with many size variants. For our experiments, we use the Llava-7B variant. Llava has 2 parts: a vision tower and an LLM. The vision tower processes the image tokens. It does so by breaking the input image into multiple tokens (577 including CLS token), encoding positional information and then passing them through the attention heads in the vision tower. CLIP[3] is used as the vision tower for Llava which has 24 attention heads.

The text is tokenized separately and the both the image and text tokens are passed to the LLM for cross-modal fusion which produces the final output. Llama-7B is used as the LLM for Llava-7B.

Test Bed

For experimentation, we utilized the Llava-7B implementation from transformers[2]. Since LLava 7B requires atleast 24GB VRAM to run, I utilized the GPU resources available at Kaggle. This was a bottleneck because of limited availability. We experimented on different images with various prompts to see the characteristics. We selected 4 random images from the Coco validation and test sets and tried 4 different prompts for our experiments. These are the images being used for testing:

Figure 1: Images used for testing. Image ordering goes from left to right and top to bottom

Finding 1: Significant drop-off in attention entropy in Vision Tower

First, we analyzed the per‐head entropy within the vision tower to quantify each head’s dispersion of focus across image tokens. Entropy, defined as:

represents the normalized attention weight to token (i), serves as a proxy for the breadth of a head’s receptive field: high entropy implies a diffuse attention spread over many tokens, whereas low entropy indicates concentrated focus on a few salient regions. Across a diverse set of natural and synthetic images, we observed a pronounced and reproducible drop in entropy values between attention heads 10 and 11, and head 12, as illustrated in Figure 1. This inflection point suggests that heads 1–11 progressively refine their contextual aggregation, but head 12—and those beyond—begin to collapse onto a narrower subset of tokens. Notably, this pattern held regardless of image content or complexity, implying an architectural or learned bias in the middle‐to‐late vision layers. We hypothesize that low‐entropy heads may prematurely discard secondary but potentially relevant features, thereby propagating suboptimal token representations into deeper layers. In the next section, we systematically probe individual token attention scores to determine whether these low‐entropy dynamics indeed lead to persistent information loss during forward propagation.

Figure 1: Entropy graph of vision tower attention heads. The trend shows a drop in attention scores from the halfway mark.

This was reflected in the CLS-token attention map of each layer. The first half of the layers show a wide distribution of attention whereas the latter ones have sparse distribution.

Figure 2: CLS-token attention maps of vision tower layer. Attention distribution clearly shows that sparse distribution of tokens in latter half of layers across different images

Finding 2: Llava can tolerate visual token pruning pre cross-modal fusion, especially in the latter attention heads of vision tower

Based on the previous findings, we conducted experiments to see how many tokens are below a certain threshold in the vision tower attention heads and if pruning image tokens based on those low scores affect inference performance of the model.

Let n be the number of attention heads we wish to scan and t be the maximum attention score threshold. We want to find what percentage of attention scores in the last n heads are below threshold t.

We tested the percentage of attention scores below certain thresholds and present the results in the table below:

Percentage of Attention Scores Below Various Thresholds

Layer	< 0.00005	< 0.0003	< 0.0004	< 0.0005	< 0.005	< 0.05
1	0.82%	13.96%	19.76%	25.38%	99.65%	100.00%
2	8.08%	30.77%	35.95%	40.95%	97.32%	100.00%
3	6.46%	32.52%	40.81%	48.30%	98.82%	100.00%
4	5.84%	28.39%	37.76%	47.07%	99.31%	100.00%
5	4.76%	30.78%	38.34%	45.55%	97.53%	100.00%
6	2.20%	27.84%	36.09%	43.13%	97.56%	100.00%
7	0.00%	9.57%	17.20%	25.62%	98.51%	100.00%
8	0.03%	6.16%	10.36%	15.82%	98.62%	100.00%
9	0.04%	7.73%	12.86%	18.47%	97.81%	100.00%
10	2.08%	15.49%	21.95%	28.01%	95.00%	100.00%
11	0.04%	10.74%	16.95%	23.33%	95.74%	99.98%
12	0.17%	14.88%	22.18%	28.58%	94.89%	99.98%
13	0.66%	21.53%	31.55%	40.64%	96.86%	99.72%
14	5.89%	47.40%	57.06%	63.80%	97.89%	99.36%
15	3.75%	40.18%	51.04%	60.34%	98.30%	99.35%
16	11.05%	44.39%	53.24%	61.02%	98.13%	99.35%
17	9.72%	46.95%	56.74%	64.59%	97.84%	99.31%
18	15.25%	40.60%	48.07%	54.99%	97.60%	99.36%
19	20.20%	46.34%	53.83%	59.77%	97.49%	99.38%
20	18.83%	43.90%	51.63%	58.17%	97.02%	99.41%
21	23.20%	47.12%	54.46%	60.34%	97.01%	99.45%
22	4.73%	45.62%	52.39%	57.87%	96.78%	99.50%
23	24.97%	51.91%	57.64%	62.88%	95.82%	99.67%
24	19.57%	52.44%	58.51%	63.48%	95.58%	99.53%

Table 1: Attention score percentage under different threshold for all layers in vision tower. All the values are in the range of (0.00005, 0.005)

As the preceding table shows, most of the attention scores are between 0.00005 and 0.005, excluding. Based on this finding, we experiment using 0.0003, 0.0004 and 0.0005 as pruning thresholds.

Patch-wise Pruning:

In this technique, we compute each patch’s importance by averaging its attention to the CLS token over all layers and a selected subset of the final N heads. From the full attention tensor of shape [L, B, H, T, T], we extract CLS-to-patch scores (yielding [L, B, H, P]), retain only the last N heads ([L, B, H′, P]), and collapse layers and heads to produce a [B, P] matrix of mean attention values. Patches whose scores fall below a fixed threshold t are deemed uninformative; their corresponding feature vectors in the [B, P, hidden_dim] image tensor are zeroed out, while all others remain intact.

The results of pruning out tokens using the previously mentioned thresholds are given in the following table:

Image	Prompt	`t` = 0003	`t` = 0004	`t` = 0005
Image 1	Are there any sort of cutlery in this image?	No, there are no cutlery items like forks or knives in the image. The woman is holding two donuts on a paper plate	No, there are no cutlery items like forks or knives in the image. The woman is holding two donuts on a paper plate	Yes, there is a knife in the image.
Image 1	Describe this image	The image features a woman wearing a hat and glasses, standing in front of a doughnut shop. She is holding two plates,	The image features a woman standing in front of a doughnut shop, holding two plates with doughnuts on them. She is sm	The image features a woman wearing a hat and glasses, standing in front of a doughnut shop. She is holding two plates,
Image 1	Is it night or day in this image?	It is night in the image.	It is daytime in the image.	It is daytime in the image.
Image 1	List all the objects here	1. Donuts 2. Ice cream cones 3. A woman 4. A doughnut 5. A	The objects in the image include a woman holding two donuts, a doughnut on a plate, a doughnut on a napkin,	1. Donuts 2. Ice cream 3. Woman 4. Donuts 5. Ice cream 6. Woman
Image 2	Are there any sort of cutlery in this image?	Yes, there is a pair of scissors in the image.	No, there are no knives or any other sort of cutlery visible in the image. The focus is on the woman holding a kite	Yes, there is a woman in the image who is holding a pair of scissors.
Image 2	Describe this image	The image features a woman standing in a grassy field, holding a kite in her hand. She is wearing a red shirt and blue	The image features a woman standing in a grassy field, holding a kite in her hand. She is wearing a red shirt and appears	The image features a woman standing in a grassy field, holding a kite in her hand. She is wearing a red shirt and blue
Image 2	Is it night or day in this image?	Day	It is daytime in the image.	It is daytime in the image.
Image 2	List all the objects here	[Empty response with newlines]	In the image, there are several objects, including a woman holding a kite, a car, a truck, a bus, a boat,	[Empty response with newlines]
Image 3	Are there any sort of cutlery in this image?	Yes, there is a knife in the image.	No, there are no cutlery items in this image. The image features a red and white ambulance and a red bus parked next	Yes, there is a cutlery in the image.
Image 3	Describe this image	The image depicts a city street at night, with a red and white ambulance parked next to a red and black bus. The	The image depicts a city street at night, with a red and white bus driving down the road. A white ambulance is parked	The image shows a city street at night, with a red and white ambulance parked next to a red and black bus. The bus is
Image 3	Is it night or day in this image?	It is night in this image.	It is night in this image.	It is night in this image.
Image 3	List all the objects here	1. Ambulance 2. Bus 3. Tree 4. Firemen 5. Fire truck 6. Bus	In the image, there are two buses, a truck, and a fire hydrant. The buses are parked next to each other	1. Ambulance 2. Bus 3. tree 4. firemen 5. fire truck 6.
Image 4	Are there any sort of cutlery in this image?	Yes, there are two cats in the image.	Yes, there is a knife in the image.	Yes, there are two cats in the image.
Image 4	Describe this image	The image features two cats lying on a pink couch. One cat is located on the left side of the couch, while the other	The image features two cats lying on a pink couch. One cat is positioned on the left side of the couch, while the	The image features two cats lying on a pink couch. One cat is on the left side of the couch, while the other cat
Image 4	Is it night or day in this image?	It is daytime in the image.	It is night in this image, as the two cats are sleeping on the couch.	It is daytime in the image.
Image 4	List all the objects here	1. Cat 2. Remote 3. Couch	The objects in the image include a cat, a remote control, and a couch.	The objects in the image include a cat, a remote control, a couch, and a blanket.

Table 2: Patch-wise pruning on different thresholds

Experimentally, we found that pruning the last 8 heads yielded the best results. Pruning more than 8 caused significant performance degradation and pruning any less did not give any performance enhancements. Experiments showed that different pruning thresholds gave good or bad results based on prompts. For example, for the prompt "Are there any cutlery in the image?" for Image 2, only pruning threshold 0.0004 was able to answer properly, whereas the others failed. For the same prompt, all pruning thresholds caused the model to hallucinate in some cases. There were also some instances where only threshold of 0.0004 was able to generate a response and the others were not able to.

Layer-wise Pruning:

In this layer‐wise pruning scheme, for each layer l we extract the CLS‐to‐patch attention scores from the full tensor of shape [L, B, H, T, T], yielding [L, B, H, P], and retain only the final N heads to form [L, B, H′, P]. We then average over those N heads within each layer to produce a [B, P] matrix of mean attention values per layer, apply a fixed threshold to identify uninformative patch–head pairs, and expand the resulting binary mask into a full [B, P, hidden_dim] feature mask by assigning each head to its contiguous hidden‐dimension segment. Finally, this mask is applied element‐wise to the image features [B, P, hidden_dim], zeroing out only those head‐specific features that fall below the threshold while preserving all others.

The table below shows the generation results for different threshold values:

Table 3: Layer-wise pruning on different thresholds

For layer-wise pruning, we found that pruning the last 9 attention heads yielded the best results. Unlike patch-wise pruning all thresholds were able to produce responses, whether it be correct or incorrect. Responses were more or less consistent in all cases. However, just like patch-wise pruning, the model hallucinated responses in all cases.

In the following table, we provide generation time across unmodified model and our two pruning strategies:

Image 1

Prompt	No Change (s)	Layer-wise Pruning	Patch-wise Pruning
Are there any sort of cutlery in this image?	19.50	17.20	7.54
Describe this image	18.95	18.19	18.62
Is it night or day in this image?	5.70	5.44	5.48
List all the objects here	19.64	18.30	18.58

Image 2

Prompt	No Change (s)	Layer-wise Pruning	Patch-wise Pruning
Are there any sort of cutlery in this image?	18.45	9.12	12.53
Describe this image	19.21	18.28	18.62
Is it night or day in this image?	5.64	5.46	5.56
List all the objects here	19.13	18.44	18.56

Image 3

Prompt	No Change (s)	Layer-wise Pruning	Patch-wise Pruning
Are there any sort of cutlery in this image?	18.65	13.31	7.83
Describe this image	19.51	18.38	19.11
Is it night or day in this image?	5.44	4.92	4.92
List all the objects here	20.37	18.62	18.76

Image 4

Prompt	No Change (s)	Layer-wise Pruning	Patch-wise Pruning
Are there any sort of cutlery in this image?	32.89	32.13	13.68
Describe this image	17.90	17.73	26.79
Is it night or day in this image?	12.05	5.28	5.31
List all the objects here	11.63	18.07	13.70

Table 4: Comparison of generation time across all models for different prompt per image

From the table, we can see that patch-wise pruning achieved 17.7% faster generation times compared to the baseline, with some individual queries running up to 4x faster. Even the more conservative layer-wise pruning approach showed 4.4% improvements over baseline processing.

Not all tasks benefit equally from masking:

-Object Detection Queries: Patch-wise pruning dominated, reducing "cutlery detection" times from 22.37s (baseline) to just 10.40s on average
-Simple Binary Questions: "Is it night or day?" queries were consistently fast across all approaches (5-6 seconds)
-Complex Description Tasks: Showed more mixed results, with layer-wise pruning often performing best
-Object Listing: Layer-wise pruning provided the most consistent performance

While patch-wise pruning won on raw speed, layer-wise pruning proved most reliable, winning 62.5% of individual comparisons. This approach rarely produced the slowest times and maintained steady performance across different image types and question complexity.

Future Directions

This is an ongoing experiment. This blog serves as a journal of our progress so far and will continually be updated once new findings are found. A lot more things need to be added; a more comprehensive reference section and more thorough experimentation are the two things that come to mind. Hopefully we'll be able to get a paper through this independent endeavor of ours.

References

[1] H. Liu, C. Li, Q. Wu, Y. J. Lee, “Visual Instruction Tuning,” presented at the A. Oh, T. Naumann, A. Globerson, K. Saenko, M.Hardt, S. Levine (Eds.), Advances in Neural Information Processing Systems, vol. 36, pp. 34892–34916 (2023).

[2] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, A. M. Rush, “Transformers: State-of-the-Art Natural Language Processing,” presented at the Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp.38–45 (2020 Oct.).

[3] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, I. Sutskever, “Learning Transferable Visual Models From Natural Language Supervision,” (2021), URL https://arxiv.org/abs/2103.00020.

Image	Prompt	`t` = 0003	`t` = 0004	`t` = 0005
Image 1	Are there any sort of cutlery in this image?	No, there are no cutlery items in the image. The focus is on the two donuts being held by the person.	No, there are no cutlery items in the image. The focus is on the two donuts being held by the person.	No, there are no cutlery items in the image. The focus is on the two donuts being held by the person.
Image 1	Describe this image	The image features a woman standing in front of a doughnut stand, holding up two doughnuts for display. One doughnut is	The image features a woman standing in front of a doughnut stand, holding up two doughnuts for display. One doughnut is	The image features a woman standing in front of a doughnut stand, holding up two doughnuts for display. One doughnut is
Image 1	Is it night or day in this image?	It is daytime in the image.	It is daytime in the image.	It is daytime in the image.
Image 1	List all the objects here	1. Donuts 2. Doughnuts 3. Doughnuts 4. Doughnuts 5.	1. Donuts 2. Doughnuts 3. Doughnuts 4. Doughnuts 5.	1. Donuts 2. Doughnuts 3. Doughnuts 4. Doughnuts 5.
Image 2	Are there any sort of cutlery in this image?	Yes, there is a knife in the image.	Yes, there is a knife in the image.	Yes, there is a pair of scissors in the image.
Image 2	Describe this image	The image features a woman standing in a grassy field, holding a kite that is flying high in the sky. She is wearing a red	The image features a woman standing in a grassy field, holding a kite that is flying high in the sky. She is wearing a red	The image features a woman wearing a red and blue dress, standing on a grassy field. She is holding a kite, which is flying
Image 2	Is it night or day in this image?	It is daytime in the image.	It is daytime in the image.	It is daytime in the image.
Image 2	List all the objects here	1. Kite 2. Woman 3. Pants 4. Shirt 5. Sunglasses 6.	1. Kite 2. Woman 3. Pants 4. Shirt 5. Sunglasses 6.	The objects present in the image include a woman wearing a red shirt and holding a kite, a kite, a field, a par
Image 3	Are there any sort of cutlery in this image?	Yes, there is a knife in the image, which is located near the bus.	Yes, there is a knife in the image, which is located near the bus.	Yes, there is a knife in the image, which is located on the side of the road.
Image 3	Describe this image	The image depicts a city street at night, with a large red bus parked on the side of the road. The bus is positioned	The image depicts a city street at night, with a large red bus parked on the side of the road. The bus is positioned	The image depicts a city street at night, with a large red bus parked on the side of the road. The bus is positioned
Image 3	Is it night or day in this image?	It is night in this image.	It is night in this image.	It is night in this image.
Image 3	List all the objects here	1. Bus 2. Tree 3. Fire hydrant 4. Person 5. Bus stop 6. Fire truck	1. Bus 2. Tree 3. Fire hydrant 4. Person 5. Bus stop 6. Fire truck	1. Bus 2. Tree 3. Fire hydrant 4. Bus stop 5. Fire truck 6. Bus
Image 4	Are there any sort of cutlery in this image?	Yes, there is a knife in the image, which is placed on the couch along with the two cats.	Yes, there is a knife in the image, which is placed on the couch along with the two cats.	No, there are no cutlery items in the image. The image features a cat and a remote control, with the cat lying on a p
Image 4	Describe this image	The image features a couch with two cats lying on it. One cat is positioned towards the left side of the couch, while the	The image features a couch with two cats lying on it. One cat is positioned towards the left side of the couch, while the	The image features a couch with two cats lying on it. One cat is positioned towards the left side of the couch, while the
Image 4	Is it night or day in this image?	It is daytime in the image.	It is daytime in the image.	It is daytime in the image.
Image 4	List all the objects here	1. Cat 2. Cat 3. Remote control 4. Remote control 5. Cat 6. Cat	1. Cat 2. Cat 3. Remote control 4. Remote control 5. Cat 6. Cat	1. Cat 2. Cat 3. Remote control 4. Remote control 5. Cat 6. Cat