Hitchhiker's Guide to AI, Software Architecture, and Everything Else

If you are a software engineer: DON'T PANIC! This blog is my place to beam thoughts on the universe of Artificial Intelligence and Software Architecture right to your screen. On my infinite mission to boldly go where (almost) no one has gone before I will provide in-depth coverage of architectural and AI topics, personal opinions, humor, philosophical discussions, interesting news and technology evaluations. (c) Prof. Dr. Michael Stal

Saturday, July 11, 2026

TRANSFORMERS UNBOUND: HOW TO PUSH THE ARCHITECTURE TO ITS LIMITS AND BEYOND

PROLOGUE: WHY SHOULD YOU CARE?

If you have ever typed a question into a large language model and received a surprisingly coherent, even brilliant answer, you have already experienced the product of one of the most consequential architectural ideas in the history of computing: the Transformer. Introduced in 2017 by Vaswani et al. in the paper "Attention Is All You Need," the Transformer replaced the sequential, step-by-step processing of recurrent neural networks with a mechanism that looks at every token in a sequence simultaneously and decides, for each one, how much attention to pay to every other one. That single idea unlocked nearly a decade of progress that nobody fully anticipated.

But here is the thing: the original Transformer, as elegant as it is, has serious problems. It is memory-hungry. It is slow on long sequences. It wastes compute on things that do not matter. It has a fixed context window. It struggles to scale gracefully. And yet, the research community has spent the last eight years attacking every single one of those problems with remarkable ingenuity. The result is a constellation of improvements, some surgical and some sweeping, that together make modern large language models possible.

This tutorial is a guided tour through that constellation. We will start from the original architecture, understand exactly where the pain points are, and then walk through every major optimization that has been developed to address them. We will cover FlashAttention in all three of its generations, Mixture of Experts, Sparse and Native Sparse Attention, Grouped-Query Attention, Multi-head Latent Attention, Rotary Position Embeddings and their extensions, improved normalization schemes, better activation functions, quantization, speculative decoding, multi-token prediction, new optimizers, and much more. We will also look honestly at the horizon: what further improvements are coming, and where the Transformer architecture might eventually hit a wall it cannot climb over.

CHAPTER ONE: THE ORIGINAL TRANSFORMER, WARTS AND ALL

Before we can appreciate the optimizations, we need to understand what we are optimizing. Let us walk through the original Transformer architecture with enough detail to feel its beauty and its pain simultaneously.

The Transformer takes a sequence of tokens (words, subwords, or characters) and converts each one into a vector of numbers called an embedding. These embeddings are then processed by a stack of identical layers, each of which has two main components: a multi-head self-attention sublayer and a position-wise feed-forward network sublayer. Residual connections and layer normalization wrap each of these sublayers. The original paper used Post-LN normalization, meaning the layer norm was applied after the residual addition, though as we will see, this turned out to be one of the first things the community changed.

The heart of everything is the self-attention mechanism. For each token, the model computes three vectors: a Query (Q), a Key (K), and a Value (V). These are produced by multiplying the token's embedding by three learned weight matrices. The attention score between token i and token j is computed as the dot product of Q_i and K_j, scaled by the square root of the head dimension, and then passed through a softmax to produce a probability distribution. The output for token i is then the weighted sum of all Value vectors, where the weights come from that probability distribution.

Let us make this concrete with a tiny example. Suppose we have the sentence "The cat sat" and we are computing attention for the word "sat." The model might learn that "sat" should attend heavily to "cat" (because cats are the ones doing the sitting) and less to "The." The attention scores before softmax and scaling might look like this:

Query("sat") . Key("The")  = 0.3
Query("sat") . Key("cat")  = 2.1
Query("sat") . Key("sat")  = 1.0

After dividing by sqrt(d_k) and applying softmax, these become approximately:

Attention weight for "The"  = 0.07
Attention weight for "cat"  = 0.75
Attention weight for "sat"  = 0.18

The output for "sat" is then 0.07 * V("The") + 0.75 * V("cat") + 0.18 * V("sat"), which is a vector that has been heavily influenced by the representation of "cat." This is the magic: every token gets to look at every other token and gather information from those it finds most relevant.

The "multi-head" part means this process is run H times in parallel, each time with different learned projection matrices. Each head can specialize in a different kind of relationship: one head might learn syntactic dependencies, another might learn coreference, another might learn positional proximity. The outputs of all heads are concatenated and projected back to the model dimension.

Here is a schematic of one Transformer layer:

Input embeddings (+ positional encoding)
         |
         v
+---------------------------+
|   Multi-Head Self-Attention|
|   Q = X * W_Q             |
|   K = X * W_K             |
|   V = X * W_V             |
|   Attn = softmax(QK^T /   |
|          sqrt(d_k)) * V   |
+---------------------------+
         |
    Add & Norm (residual)
         |
         v
+---------------------------+
|   Feed-Forward Network    |
|   FFN(x) = max(0, xW_1   |
|            + b_1)W_2 + b_2|
+---------------------------+
         |
    Add & Norm (residual)
         |
         v
Output representations

Now, where does this beautiful machine break down? There are four fundamental pain points that every subsequent optimization targets in some way.

The first pain point is the quadratic complexity of attention. To compute attention for a sequence of N tokens, we need to compute N*N attention scores. For N=1,000, that is one million scores. For N=100,000, that is ten billion scores. The memory required to store the full attention matrix scales as O(N^2), and the compute scales the same way. This makes the vanilla Transformer essentially unusable for very long sequences without some form of approximation or restructuring.

The second pain point is the memory bandwidth bottleneck. Modern GPUs are extraordinarily fast at arithmetic, but they are relatively slow at moving data between different levels of memory. The GPU has a small, fast on-chip SRAM (the registers and shared memory) and a large, slow off-chip HBM (High Bandwidth Memory, the GPU's main memory). The naive attention implementation constantly shuttles data back and forth between these two levels, and this data movement, not the arithmetic, is often the actual bottleneck. This is the problem FlashAttention was designed to solve.

The third pain point is that all parameters are used for every input. A 70-billion-parameter model uses all 70 billion parameters to process every single token, regardless of whether those parameters are relevant to the task at hand. This is enormously wasteful. Mixture of Experts architectures address this by routing each token to only a subset of the model's parameters.

The fourth pain point is the KV cache explosion during inference. When generating text autoregressively, the model needs to remember the Keys and Values for all previously generated tokens so it does not have to recompute them. For a model with many attention heads and a long context, this KV cache can consume tens or hundreds of gigabytes of memory, severely limiting how many requests a server can handle simultaneously.

With these four pain points clearly in mind, let us now walk through the optimizations, roughly in the order they address the architecture from the inside out.

CHAPTER TWO: FIXING THE FOUNDATIONS - NORMALIZATION AND ACTIVATION FUNCTIONS

Before we get to the headline-grabbing optimizations like FlashAttention and Mixture of Experts, it is worth spending time on the quieter but enormously impactful improvements to the basic building blocks of the Transformer: how it normalizes activations and what nonlinearity it uses in its feed-forward networks. These changes sound mundane, but they have a direct effect on training stability, convergence speed, and final model quality.

2.1 FROM POST-LN TO PRE-LN: STABILIZING DEEP TRANSFORMERS

The original Transformer paper used Post-Layer Normalization, meaning the layer norm was applied after the residual addition:

Post-LN:  output = LayerNorm(x + Sublayer(x))

This seems natural, but it has a nasty property: during the early stages of training, the gradients flowing through the residual path can be very large, causing instability and requiring careful learning rate warmup schedules. Researchers discovered that moving the layer norm to before the sublayer, a scheme called Pre-LN, dramatically stabilizes training:

Pre-LN:   output = x + Sublayer(LayerNorm(x))

With Pre-LN, the residual path is clean and unobstructed, so gradients flow back through the network without being distorted by the normalization operation. This allows training to begin with a larger learning rate and without as much warmup, and it makes it possible to train much deeper models without divergence. Virtually every modern large language model, including the GPT series, LLaMA, Mistral, and their descendants, uses Pre-LN.

2.2 RMSNORM: DOING LESS TO ACHIEVE MORE

Layer Normalization, as defined by Ba et al. in 2016, computes both the mean and the variance of a layer's activations and uses both to normalize them. Zhang and Sennrich (2019) asked a provocative question: do we actually need the mean-centering step? Their answer, backed by experiments, was no. They proposed Root Mean Square Layer Normalization (RMSNorm), which skips the mean computation and normalizes only by the root mean square of the activations:

LayerNorm:  y = (x - mean(x)) / sqrt(var(x) + eps) * gamma + beta
RMSNorm:    y = x / sqrt(mean(x^2) + eps) * gamma

This is both simpler and faster. The mean computation and the beta bias term are eliminated, reducing the number of operations and parameters. Experiments show that RMSNorm achieves comparable or better performance to full LayerNorm while being 10 to 50 percent faster in practice. LLaMA, Mistral, Gemma, DeepSeek, and essentially all modern frontier models use RMSNorm. It is a small change with a real payoff.

2.3 SWIGLU: THE ACTIVATION FUNCTION THAT EVERYONE USES NOW

The original Transformer used a simple two-layer feed-forward network with ReLU activation:

FFN_ReLU(x) = max(0, x * W_1 + b_1) * W_2 + b_2

In 2020, Noam Shazeer (one of the original Transformer authors) published a short but influential paper called "GLU Variants Improve Transformers," which showed that replacing ReLU with a gated variant dramatically improves model quality. The key idea behind Gated Linear Units (GLU) is to multiply the output of one linear transformation by the sigmoid (or another activation) of a second linear transformation, effectively creating a learned gate that controls information flow:

GLU(x, W, V, b, c) = sigmoid(x * W + b)  (elementwise multiply)  (x * V + c)

Shazeer explored many variants: Bilinear (no activation on the gate), ReGLU (ReLU gate), GEGLU (GELU gate), and SwiGLU (Swish gate). SwiGLU, which uses the Swish activation function (defined as x * sigmoid(beta * x)), consistently outperformed all others:

SwiGLU(x, W, V) = Swish(x * W)  (elementwise multiply)  (x * V)

The feed-forward network with SwiGLU then becomes:

FFN_SwiGLU(x) = SwiGLU(x, W_1, W_3) * W_2

Note that SwiGLU requires three weight matrices instead of two (W_1, W_2, W_3), so to keep the parameter count comparable, the hidden dimension is typically reduced by a factor of 2/3. Despite this, SwiGLU models consistently outperform their ReLU counterparts. PaLM, LLaMA, Mistral, Gemma, DeepSeek, and virtually every other modern frontier model uses SwiGLU or GEGLU. This is not a minor tweak; it is a genuine quality improvement that has been replicated across dozens of model families and scales.

Here is a quick illustration of why gating helps. Imagine the feed-forward layer is trying to decide whether to pass information about "Paris" through to the next layer. A ReLU network just clips negative values to zero, which is a blunt instrument. A SwiGLU network can learn a smooth, input-dependent gate: if the context suggests we are talking about geography, open the gate wide; if we are talking about fashion, open it a bit; if we are talking about chemistry, close it almost entirely. This dynamic, content-aware gating gives the model much more expressive power per parameter.

2.4 DEEPNORM: GOING VERY, VERY DEEP

In 2022, researchers at Microsoft introduced DeepNorm (Wang et al., 2022, arXiv:2203.00555), a normalization scheme specifically designed to allow training Transformers with up to 1,000 layers. The key insight is that the residual connection in a standard Transformer causes the magnitude of activations to grow with depth, which eventually destabilizes training. DeepNorm addresses this by scaling the residual connection by a factor alpha and initializing the weights with a factor beta:

DeepNorm:  output = LayerNorm(alpha * x + Sublayer(x))

With carefully chosen alpha and beta values (derived from theoretical analysis of the expected gradient magnitudes), DeepNorm provably bounds the model update at initialization, ensuring that the network starts in a stable regime regardless of depth. This allowed the team to train a 1,000-layer Transformer that outperformed shallower models on machine translation tasks. While most production models do not use 1,000 layers, DeepNorm demonstrates that depth is not inherently a barrier to training stability if the normalization is designed correctly.

CHAPTER THREE: THE ATTENTION CRISIS AND HOW FLASHATTENTION SOLVED IT

We established in Chapter One that the naive attention computation has a serious memory problem: it requires materializing the full N-by-N attention matrix in GPU HBM, which for long sequences is both enormous and slow to access. This is the problem that FlashAttention, developed by Tri Dao and colleagues at Stanford, attacked with surgical precision.

3.1 THE MEMORY HIERARCHY PROBLEM

To understand FlashAttention, you need to understand the GPU memory hierarchy. A modern GPU like the NVIDIA A100 has roughly 80 GB of HBM (the main GPU memory), which has a bandwidth of about 2 TB/s. It also has a much smaller on-chip SRAM (shared memory plus L2 cache), totaling perhaps 40 MB, but with a bandwidth of roughly 19 TB/s. The on-chip SRAM is about 10 times faster than HBM, but it is about 2,000 times smaller.

The naive attention algorithm works like this: compute Q, K, V matrices and write them to HBM. Then read Q and K back from HBM, compute the attention scores S = Q * K^T, write S to HBM. Read S back, compute softmax(S), write the result P to HBM. Read P and V back, compute the output O = P * V, write O to HBM. Each of these read-write operations is a round trip to the slow HBM. For a sequence of length N with head dimension d, the total HBM accesses scale as O(N^2), and for large N this data movement dominates the total runtime.

The key insight of FlashAttention (Dao et al., 2022, arXiv:2205.14135) is that we do not actually need to materialize the full N-by-N attention matrix. We can compute the attention output in tiles, keeping the intermediate results in the fast on-chip SRAM and only writing the final output to HBM. This is called tiling or blocking, and it is a classic technique in high-performance computing applied here to the specific structure of the softmax attention computation.

3.2 THE TILING TRICK AND THE ONLINE SOFTMAX

The challenge with tiling attention is the softmax. Softmax requires knowing the sum of all exponentials in a row before you can normalize any of them. If you are processing the attention matrix in tiles, you do not see the whole row at once, so how can you compute the correct softmax?

The answer is the online softmax algorithm, which allows you to compute a running maximum and a running sum of exponentials as you process tiles one by one, and then correct the output at the end. Here is the key recurrence: suppose you have processed tiles 1 through t and have accumulated a running maximum m_t and a running sum of exponentials l_t. When you process tile t+1, you update:

m_{t+1} = max(m_t, max of new tile scores)
l_{t+1} = exp(m_t - m_{t+1}) * l_t + sum of exp(new tile scores - m_{t+1})

And you correct the accumulated output O_t:

O_{t+1} = (l_t * exp(m_t - m_{t+1}) * O_t
           + exp(new scores - m_{t+1}) * V_new) / l_{t+1}

This recurrence is mathematically exact; it produces the same result as computing the full softmax at once, but it only ever needs to hold one tile of the attention matrix in SRAM at a time. The full N-by-N matrix never needs to be written to HBM.

The result is dramatic. FlashAttention reduces HBM accesses from O(N^2) to O(N^2 * d / M), where M is the SRAM size. For typical values of d and M, this is a reduction of roughly 5 to 20 times in memory reads and writes. In practice, FlashAttention achieves 2 to 4 times speedup over the naive PyTorch attention implementation and reduces memory usage by up to 7.6 times. It also enables training on much longer sequences: the original FlashAttention paper demonstrated training GPT-2 on sequences of 4,096 tokens (four times longer than the standard 1,024) while being faster than the standard implementation on 1,024-token sequences.

3.3 FLASHATTENTION-2: BETTER PARALLELISM

FlashAttention-2 (Dao, 2023, arXiv:2307.08691) identified a further source of inefficiency in the original FlashAttention: the work partitioning between GPU thread blocks and warps was suboptimal. In FlashAttention-1, the outer loop was over the sequence dimension (rows of Q), and the inner loop was over the key-value sequence (rows of K and V). This led to low GPU occupancy in some configurations. FlashAttention-2 reorganized the computation to maximize parallelism across the batch, head, and sequence dimensions simultaneously, and it reduced the number of non-matmul operations (which are less efficiently pipelined on modern GPUs). The result was a further 2x speedup over FlashAttention-1, reaching 50 to 73 percent of the theoretical maximum FLOPs/s on A100 GPUs.

3.4 FLASHATTENTION-3: EXPLOITING HOPPER GPU FEATURES

FlashAttention-3 (Shah et al., 2024, arXiv:2407.08608) was written specifically for NVIDIA's Hopper GPU architecture (H100), which introduced several new hardware features that FlashAttention-2 could not exploit. The three main ideas in FlashAttention-3 are warp specialization, interleaved matmul and softmax, and FP8 support with block quantization.

Warp specialization means that different warps (groups of 32 GPU threads) are assigned different roles: some warps handle data movement using the new Tensor Memory Accelerator (TMA) hardware on Hopper, while others handle computation using the new WGMMA (Warpgroup Matrix Multiply-Accumulate) instructions. This allows data movement and computation to overlap, hiding the latency of memory transfers. Interleaving the matmul and softmax operations further hides the quantization overhead when using reduced-precision arithmetic. With FP8 precision (8-bit floating point), FlashAttention-3 achieves up to 1.2 PFLOPs/s (petaflops per second) on a single H100 GPU, which is 75 percent of the theoretical maximum. With FP16, it achieves 1.5 to 2.0 times speedup over FlashAttention-2, reaching up to 740 TFLOPs/s.

To give you a sense of the cumulative progress, here is a comparison of attention implementations on A100/H100 GPUs for a typical configuration:

Implementation              Speed (TFLOPs/s)   Memory Usage
------------------------------------------------------------
Standard PyTorch Attn       ~30                O(N^2)
FlashAttention-1 (A100)     ~120               O(N)
FlashAttention-2 (A100)     ~200               O(N)
FlashAttention-3 FP16(H100) ~740               O(N)
FlashAttention-3 FP8 (H100) ~1200              O(N)

This is not a minor engineering improvement. FlashAttention is arguably the single most impactful practical optimization in the history of Transformer training, and it is now used by essentially every serious LLM training framework including PyTorch, JAX, and all major cloud providers' training stacks.

CHAPTER FOUR: SPARSE ATTENTION - PAYING ATTENTION SELECTIVELY

FlashAttention makes exact attention faster by being smarter about memory access patterns. But it does not change the fundamental O(N^2) complexity of the computation itself. For truly long sequences, even an optimally memory-efficient exact attention is too slow. The solution is sparse attention: instead of computing attention scores between every pair of tokens, only compute scores between tokens that are likely to be relevant to each other.

4.1 FIXED SPARSE PATTERNS: LONGFORMER AND BIGBIRD

The earliest and most principled sparse attention approaches used fixed, predetermined patterns that capture the most important structural relationships in a sequence.

The Longformer (Beltagy et al., 2020, arXiv:2004.05150) combines two types of attention. Local windowed attention means each token attends to a fixed window of w tokens on either side, capturing local context with O(N * w) complexity. Global attention means a small set of special tokens (like the [CLS] classification token) attend to all other tokens and are attended to by all other tokens, capturing global context. The combination gives each token local awareness and a global summary, at a total cost of O(N * (w + g)) where g is the number of global tokens. This scales linearly with sequence length, making it practical for documents of tens of thousands of tokens.

BigBird (Zaheer et al., 2020, arXiv:2007.14062) adds a third component to this recipe: random attention, where each token also attends to a small set of randomly chosen tokens from the full sequence. The combination of local windowed attention, global attention on special tokens, and random attention is theoretically motivated: BigBird proves that this combination is a universal approximator of full attention, meaning that any function computable by full attention can also be computed (approximately) by BigBird's sparse pattern. The random attention component ensures that information can flow between any two tokens in the sequence in a bounded number of hops, preserving the long-range dependency modeling that makes Transformers powerful.

Here is a visualization of these attention patterns for a sequence of 12 tokens (each row is a query token, each column is a key token, X means attention is computed, and a dot means it is skipped):

Full Attention (every token attends to every token):
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X
X X X X X X X X X X X X

Longformer (window=2, token 0 is global):
X X X X X X X X X X X X   <- global token attends everywhere
X X X X . . . . . . . .   <- local window of 2
X X X X X . . . . . . .
X . X X X X . . . . . .
X . . X X X X . . . . .
X . . . X X X X . . . .

BigBird (window=2, token 0 global, plus 1 random per row):
X X X X X X X X X X X X   <- global token
X X X X . . X . . . . .   <- local + 1 random (col 6)
X X X X X . . . . X . .   <- local + 1 random (col 9)
X . X X X X . . . . . X   <- local + 1 random (col 11)
X . . X X X X . . . . .   <- local only (no random shown)
X . . . X X X X . X . .   <- local + 1 random (col 9)

The sparsity in these patterns means that for a sequence of 64,000 tokens with a window of 512, each token only computes attention with roughly 1,024 other tokens instead of 64,000, a 64-fold reduction in compute.

4.2 NATIVE SPARSE ATTENTION: TRAINABLE SPARSITY

Fixed sparse patterns like Longformer's are a reasonable approximation, but they are not learned from data. The model cannot adapt the sparsity pattern to the content of the input. In 2025, DeepSeek introduced Native Sparse Attention (NSA, arXiv:2502.05171), a sparse attention mechanism that is trainable from scratch and hardware-aligned for efficient GPU execution.

NSA uses a hierarchical approach. First, it compresses the key-value sequence into coarse-grained blocks by averaging or pooling tokens within each block, and computes attention scores between each query and these compressed blocks. This gives a rough estimate of which regions of the sequence are relevant. Second, it selects the top-k most relevant blocks based on these coarse scores, and then computes fine-grained attention only within those selected blocks. This two-stage process ensures that the model can attend to any part of the sequence if it is relevant, while spending most of its compute on the most important regions.

The crucial innovation is that NSA is differentiable end-to-end: the block selection process uses a straight-through estimator or a continuous relaxation, allowing gradients to flow through the selection step during training. This means the model learns, from the data, which tokens are worth attending to, rather than having this decision hardcoded by the architecture designer. On 64K-token sequences, NSA achieves a 9x speedup over full attention while maintaining model quality comparable to full attention. It is used in DeepSeek's long-context models and represents the current state of the art in practical trainable sparse attention.

4.3 LINEAR ATTENTION: THE KERNEL TRICK

A completely different approach to reducing attention complexity is to approximate the softmax attention with a kernel function that can be computed without materializing the full N-by-N matrix. The key observation is that the softmax attention can be written as:

Attention(Q, K, V) = softmax(Q * K^T / sqrt(d)) * V

If we could replace softmax(Q * K^T) with a factorized form phi(Q) * phi(K)^T, where phi is some feature map, then we could compute the attention output as:

Attention(Q, K, V) = phi(Q) * (phi(K)^T * V)

The parenthesized term phi(K)^T * V can be computed once and reused for all queries, reducing the complexity from O(N^2 * d) to O(N * d^2). This is the kernel trick applied to attention.

The Performer (Choromanski et al., 2020, arXiv:2009.14794) implements this idea using random feature approximations to the softmax kernel, specifically a method called FAVOR+ (Fast Attention Via positive Orthogonal Random features). The Linformer (Wang et al., 2020, arXiv:2006.04768) takes a different approach: it projects the keys and values to a lower-dimensional space of size k (where k is much smaller than N) before computing attention, reducing complexity to O(N * k). Both approaches achieve linear complexity in sequence length, but at the cost of approximation error. In practice, linear attention models have historically underperformed exact attention models on language tasks, though the gap has been narrowing.

Lightning Attention, used in MiniMax-01 (2025, arXiv:2501.09755), is a recent linear attention implementation that achieves practical efficiency through careful hardware-aware implementation. MiniMax-01 scales lightning attention to a 456-billion-parameter model with 45.9 billion activated parameters per token and a 32-million-token context window, demonstrating that linear attention can work at frontier scale when implemented carefully.

CHAPTER FIVE: RETHINKING MULTI-HEAD ATTENTION - GQA, MQA, AND MLA

The standard multi-head attention (MHA) in the original Transformer uses H query heads, H key heads, and H value heads. During autoregressive inference, the model must cache the keys and values for all previously generated tokens. For a model with H=64 heads, a head dimension of 128, and a context of 100,000 tokens, the KV cache for a single layer is 64 * 128 * 100,000 * 2 (for K and V) = 1.6 billion numbers. At 16-bit precision, that is 3.2 GB per layer. For a 96-layer model, the total KV cache is 307 GB, which exceeds the memory of most GPU clusters. This is the KV cache crisis, and it has driven several important architectural innovations.

5.1 MULTI-QUERY ATTENTION: THE RADICAL SIMPLIFICATION

Multi-Query Attention (MQA), proposed by Shazeer in 2019, takes the most aggressive approach: use a single key head and a single value head shared across all query heads. All H query heads attend to the same K and V projections. This reduces the KV cache size by a factor of H (e.g., 64x for a 64-head model), which is enormous. The cost is some degradation in model quality, because the single shared K and V must serve all query heads, limiting the diversity of attention patterns.

5.2 GROUPED-QUERY ATTENTION: THE GOLDILOCKS SOLUTION

Grouped-Query Attention (GQA), proposed by Ainslie et al. (2023, arXiv:2305.13245), finds a middle ground between the full H key-value heads of MHA and the single key-value head of MQA. GQA divides the H query heads into G groups, and each group shares a single key head and a single value head. With G groups, the KV cache is reduced by a factor of H/G compared to MHA, while quality is much closer to MHA than MQA.

MHA:  H query heads, H key heads, H value heads   (KV cache = H * d * N)
GQA:  H query heads, G key heads, G value heads   (KV cache = G * d * N)
MQA:  H query heads, 1 key head,  1 value head    (KV cache = 1 * d * N)

For example, LLaMA 3 uses 32 query heads and 8 key-value heads (GQA with G=8), reducing the KV cache by 4x compared to MHA while maintaining quality nearly indistinguishable from full MHA. Mistral 7B uses 32 query heads and 8 key-value heads. Gemma, Qwen, and virtually every modern open-source model uses GQA. It is now the de facto standard for production LLMs.

The Ainslie et al. paper also showed that you can convert an existing MHA model to GQA by mean-pooling the key and value heads within each group and then fine-tuning for a small number of steps, so you do not necessarily need to train from scratch to benefit from GQA.

5.3 MULTI-HEAD LATENT ATTENTION: THE COMPRESSION APPROACH

DeepSeek-V2 (2024, arXiv:2405.04434) introduced a more radical approach called Multi-head Latent Attention (MLA). Instead of caching H separate key-value pairs, MLA compresses the key-value pairs into a single low-rank latent vector c of dimension d_c, where d_c is much smaller than H * d_head. At inference time, the full key and value matrices are reconstructed from this latent vector using learned up-projection matrices.

The idea is analogous to principal component analysis: instead of storing the full high-dimensional key-value representation, store only the low-dimensional latent code that captures most of the information. The reconstruction at inference time adds a small computational cost, but the memory savings are enormous. DeepSeek-V2 reports that MLA reduces the KV cache by 93.3 percent compared to standard multi-head attention, while maintaining model quality comparable to or better than GQA.

Here is a comparison of KV cache sizes for a model with 64 heads, head dimension 128, and context length 100K tokens:

Method      KV cache per layer   Reduction vs MHA
--------------------------------------------------
MHA         1.6 GB               1x  (baseline)
GQA (G=8)   200 MB               8x
MQA         25 MB                64x
MLA         ~107 MB              ~15x (higher quality than MQA)

MLA is used in DeepSeek-V2, DeepSeek-V3, DeepSeek-V4 and DeepSeek-R1, all of which are among the most capable open-source models as of mid-2026. The technique represents a genuinely novel approach to the KV cache problem that goes beyond the simple head-sharing of GQA and MQA.

CHAPTER SIX: POSITIONAL ENCODING - TELLING THE MODEL WHERE THINGS ARE

The Transformer's self-attention mechanism is, by itself, permutation-invariant: if you shuffle the tokens in a sequence, the attention scores change (because the tokens themselves change), but the mechanism has no inherent notion of order. To give the model information about the position of each token in the sequence, we need positional encodings. This seemingly simple requirement has spawned a rich line of research with profound implications for context length generalization.

6.1 SINUSOIDAL AND LEARNED ABSOLUTE POSITIONS

The original Transformer used sinusoidal positional encodings: fixed vectors added to the token embeddings, where the i-th dimension of the positional encoding for position pos uses a sine or cosine function of a specific frequency:

PE(pos, 2i)   = sin(pos / 10000^(2i / d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

The intuition is that each dimension oscillates at a different frequency, so different positions produce different patterns of sine and cosine values. The model can learn to use these patterns to infer relative positions. BERT and many early models used learned absolute position embeddings instead, where each position gets a trainable vector. Both approaches have a fundamental limitation: they do not generalize well to sequence lengths longer than those seen during training.

6.2 ROTARY POSITION EMBEDDINGS (ROPE): THE MODERN STANDARD

Rotary Position Embedding (RoPE), introduced by Su et al. (2021, arXiv:2104.09864) in the RoFormer paper, takes a fundamentally different approach. Instead of adding position information to the token embeddings, RoPE encodes position by rotating the query and key vectors before computing attention scores. Specifically, the query vector at position m and the key vector at position n are rotated by angles proportional to m and n respectively, so that their dot product depends only on their relative position (m - n):

Rotated_Q(m) = R(m) * Q
Rotated_K(n) = R(n) * K
Q(m) . K(n) = (R(m) * Q) . (R(n) * K) = Q . (R(m-n) * K)

where R(theta) is a rotation matrix parameterized by theta. This is elegant: the attention score between two tokens automatically encodes their relative distance, and the encoding decays naturally as the distance increases (because the dot product of two rotated vectors decreases as the rotation angle between them increases).

RoPE has several important properties. It generalizes to sequences longer than those seen during training better than absolute position embeddings. It is compatible with efficient attention implementations. It naturally captures the intuition that nearby tokens are more relevant than distant ones. And it is computationally cheap: the rotation can be implemented as an element-wise complex multiplication, adding negligible overhead. RoPE is now used by LLaMA, Mistral, Qwen, Gemma, DeepSeek, and essentially every modern open-source LLM.

6.3 ALIBI: ATTENTION WITH LINEAR BIASES

ALiBi (Attention with Linear Biases), proposed by Press et al. (2021, arXiv:2108.12409), takes yet another approach. Instead of modifying the token embeddings or the query/key vectors, ALiBi simply adds a penalty to the attention score that is proportional to the distance between the query and key tokens:

Attention score(i, j) = Q_i . K_j / sqrt(d)  -  m * |i - j|

where m is a head-specific slope. The penalty grows linearly with distance, so distant tokens are always penalized relative to nearby ones. This has a beautiful property: because the bias is relative (not absolute), the model can generalize to sequences longer than those seen during training simply by extrapolating the linear bias. ALiBi models trained on 1,024-token sequences can often handle 2,048 or even 4,096 tokens at test time with minimal quality degradation. ALiBi is used in BLOOM and MPT, among others.

6.4 EXTENDING CONTEXT: POSITION INTERPOLATION, YARN, AND IROPE

A major practical challenge is extending the context window of a pretrained model beyond its training length. If a model was trained on sequences of 4,096 tokens and you want to use it on 128,000-token sequences, the positional encodings for positions beyond 4,096 are out-of-distribution, and the model will likely produce garbage.

Position Interpolation (Chen et al., 2023, arXiv:2306.15595) addresses this by scaling down all position indices so that the maximum position in the long sequence maps to the maximum position seen during training. If the training length is L and the target length is L', each position index pos is replaced by pos * (L / L'). This ensures that all positions are within the training distribution, at the cost of reducing the resolution of position information (nearby tokens now have more similar position encodings than before). A small amount of fine-tuning on long sequences restores quality.

YaRN (Yet another RoPE extensioN, Peng et al., 2023, arXiv:2309.00071) improves on simple interpolation by using a non-uniform approach: high-frequency components of the RoPE encoding (which capture fine-grained local position information) are left unchanged, while low-frequency components (which capture coarse global position information) are interpolated. This preserves the local position resolution that is important for language modeling while still allowing the model to handle longer sequences. YaRN can extend LLaMA models from 4,096 to 128,000 tokens with minimal fine-tuning, achieving state-of-the-art performance on long-context benchmarks.

The most recent innovation in this space is iRoPE (interleaved RoPE), introduced in Meta's Llama 4 (2025, arXiv:2505.09343). iRoPE interleaves standard RoPE attention layers with NoPE (No Positional Encoding) layers throughout the model. The NoPE layers have no positional encoding at all, relying entirely on the content of the tokens to determine attention patterns. The intuition is that some layers should focus on local, position-sensitive relationships (handled by RoPE layers), while others should focus on global, position-agnostic relationships (handled by NoPE layers). This combination allows Llama 4 Scout to generalize to context windows of up to 10 million tokens, an extraordinary achievement that would have seemed impossible just a few years ago.

CHAPTER SEVEN: MIXTURE OF EXPERTS - THE ARCHITECTURE OF SELECTIVE COMPUTATION

We now come to one of the most transformative architectural innovations of the past few years: Mixture of Experts (MoE). The core idea is simple and powerful. Instead of having a single feed-forward network that processes every token, have many feed-forward networks (the "experts") and route each token to only a small subset of them. This allows the model to have a very large total number of parameters while only using a small fraction of them for any given token, keeping the computational cost manageable.

7.1 THE BASIC MoE LAYER

In a standard Transformer, the feed-forward network in each layer has a fixed set of parameters that are applied to every token. In an MoE Transformer, the feed-forward network is replaced by E expert networks FFN_1, FFN_2, ..., FFN_E, plus a lightweight router network that decides which experts to use for each token.

The router takes the token's hidden state as input and produces a probability distribution over the E experts. The top-k experts (typically k=1 or k=2) are selected, and the token is processed by those experts. The outputs are combined as a weighted sum, where the weights come from the router's probability distribution:

Router output:  g = softmax(x * W_router)
Top-k experts:  {e_1, e_2, ..., e_k} = top-k indices of g
MoE output:     sum over i in {e_1,...,e_k} of g[e_i] * FFN_{e_i}(x)

Here is a concrete example with 8 experts and top-2 routing for a single token:

Router scores (before softmax):
Expert 1: 0.8   Expert 2: 0.3   Expert 3: 1.2   Expert 4: 0.1
Expert 5: 0.9   Expert 6: 0.5   Expert 7: 0.2   Expert 8: 0.6

After softmax, top-2 selection (Expert 3 and Expert 5):
Expert 3 score: 0.42  (renormalized)
Expert 5 score: 0.31  (renormalized)

MoE output = 0.42 * FFN_3(x) + 0.31 * FFN_5(x)

The key insight is that while the model has 8 * FFN_size parameters in this layer, each token only activates 2 of them. The total parameter count is large (enabling the model to store more knowledge), but the active parameter count per token (which determines the compute cost) is small.

7.2 THE SWITCH TRANSFORMER: SIMPLICITY AT SCALE

The Switch Transformer (Fedus et al., 2021, arXiv:2101.03961) was one of the first demonstrations that MoE could work at very large scale. It simplified the routing to top-1 (each token goes to exactly one expert) and showed that this could achieve 7x speedup over a dense T5 model with the same compute budget. The Switch Transformer scaled to 1.6 trillion parameters, demonstrating that MoE could reach parameter counts that were simply impossible for dense models.

The Switch Transformer also identified a key challenge with MoE: load balancing. If the router learns to always send tokens to the same few experts, those experts become bottlenecks while the others are idle. To prevent this, the Switch Transformer introduced an auxiliary load balancing loss that encourages the router to distribute tokens evenly across experts:

Load balancing loss = E * sum_over_experts_e of
                     (fraction of tokens to e) * (fraction of router prob to e)

This loss is added to the main training loss with a small coefficient, gently nudging the router toward balanced routing without overriding the primary training objective.

7.3 MIXTRAL: MoE GOES MAINSTREAM

Mixtral 8x7B (Mistral AI, 2023, arXiv:2401.04088) brought MoE to the open-source community in a highly accessible form. Mixtral has the same architecture as Mistral 7B, but each feed-forward layer is replaced by 8 experts with top-2 routing. The total parameter count is 47B, but only 13B parameters are active for any given token. Mixtral outperforms LLaMA 2 70B on most benchmarks while using less than half the active compute, demonstrating that MoE is not just a research curiosity but a practical architecture for building high-quality models efficiently.

7.4 DEEPSEEKMOE: FINE-GRAINED EXPERTS AND SHARED EXPERTS

DeepSeekMoE (Dai et al., 2024, arXiv:2401.06066) introduced two important innovations that pushed MoE further. The first is fine-grained expert segmentation: instead of having a small number of large experts, use a large number of small experts. For example, instead of 8 experts each with hidden dimension 4096, use 64 experts each with hidden dimension 512. With top-k routing selecting more experts (e.g., top-8 out of 64), the model can form more flexible combinations of expertise, since any subset of 8 small experts can be activated. This gives the model much more combinatorial flexibility in how it processes different types of tokens.

The second innovation is shared expert isolation: a small number of experts (typically 1 or 2) are designated as "shared experts" that are always activated for every token, regardless of the router's decision. These shared experts are expected to capture common knowledge that is useful for all tokens, while the routed experts capture specialized knowledge. This separation prevents the routed experts from wasting capacity on universal patterns that could be handled by the shared experts. DeepSeekMoE 16B achieves comparable performance to LLaMA-2 70B while using only 40 percent of the compute, a remarkable efficiency gain.

7.5 DEEPSEEK-V3: THE STATE OF THE ART IN MoE

DeepSeek-V3 (December 2024, arXiv:2412.19437) represents the current pinnacle of MoE architecture design as of mid-2026. It has 671 billion total parameters with 37 billion active parameters per token. It combines MLA for efficient KV caching with DeepSeekMoE for efficient computation. It introduces two further innovations: auxiliary-loss-free load balancing and multi-token prediction.

Auxiliary-loss-free load balancing replaces the explicit load balancing loss with a bias-based mechanism: each expert has a learnable bias added to its router score, and these biases are adjusted dynamically during training to ensure balanced routing without contaminating the main training loss. This is a subtle but important improvement because the auxiliary loss can interfere with the primary training objective, and removing it allows the model to optimize more cleanly.

Multi-token prediction (MTP) trains the model to predict not just the next token but the next several tokens simultaneously, using additional output heads. This is used as an auxiliary training objective that improves sample efficiency and encourages the model to plan ahead in its representations. At inference time, the MTP heads can also be used for speculative decoding (which we discuss in Chapter Nine), where the model generates multiple candidate tokens in parallel and then verifies them. DeepSeek-V3 achieves performance comparable to leading closed-source models like GPT-4o and Claude 3.5 Sonnet at a fraction of the training cost.

7.6 KIMI K2: TRILLION-PARAMETER MoE

Kimi K2 (Moonshot AI, June 2025, arXiv:2506.01939) pushes MoE to one trillion total parameters with 32 billion active parameters per token, trained on 15.5 trillion tokens. It achieves state-of-the-art performance among open-source models on agentic and coding tasks. Kimi K2 also introduces the MuonClip optimizer (discussed in Chapter Ten), demonstrating that architectural and optimization innovations continue to compound at the frontier.

7.7 EXPERT CHOICE ROUTING: FLIPPING THE SCRIPT

All the MoE systems described so far use token-choice routing: each token chooses which experts to route to. This can lead to load imbalance if the router consistently prefers certain experts. Expert Choice Routing (Zhou et al., 2022, and subsequent work in 2025) flips this: instead of tokens choosing experts, experts choose which tokens to process. Each expert selects the top-k tokens from the batch that it wants to process, based on the router scores. This naturally enforces perfect load balancing (every expert processes exactly the same number of tokens) without any auxiliary loss. The tradeoff is that some tokens may be processed by more experts than others, and some tokens may be processed by fewer, but in practice this heterogeneity can actually be beneficial, as more complex or ambiguous tokens receive more compute.

CHAPTER EIGHT: THE BYTE LATENT TRANSFORMER - RETHINKING TOKENIZATION ITSELF

Before leaving the topic of how information is structured and fed into Transformers, we should mention a radical rethinking that emerged in late 2024: the Byte Latent Transformer (BLT, Yu et al., 2024, arXiv:2412.09871) from Meta AI. Standard LLMs operate on tokens, which are subword units produced by a tokenizer like BPE (Byte Pair Encoding). Tokenization is a preprocessing step that converts raw bytes into a fixed vocabulary of subword units, which the model then processes.

Tokenization has several well-known problems. It is language-specific: tokenizers trained on English text are inefficient for other languages, producing many more tokens per character. It is brittle: small changes to a word (like adding a space or changing capitalization) can produce very different token sequences. It requires a fixed vocabulary decided before training. And it cannot adapt its granularity to the difficulty of the content.

BLT operates directly on raw bytes, without any tokenization. Instead of processing a fixed sequence of tokens, BLT dynamically groups bytes into patches of variable size, where the patch boundaries are determined by the entropy of the next byte prediction. High-entropy regions (where the next byte is hard to predict, indicating complex or novel content) get smaller patches with more compute per byte. Low-entropy regions (where the next byte is easy to predict, indicating common or repetitive content) get larger patches with less compute per byte.

This dynamic allocation of compute is elegant: the model spends more effort on the parts of the input that are actually difficult and less effort on the parts that are easy. BLT matches tokenization-based LLM performance at scale while being more robust to character-level perturbations and more efficient on non-English languages. It represents a potential paradigm shift in how LLMs process input, though as of mid-2026 it has not yet been widely adopted in production models.

CHAPTER NINE: MAKING INFERENCE FAST - QUANTIZATION, SPECULATIVE DECODING, AND MULTI-TOKEN PREDICTION

Training a large language model is expensive, but inference is where the costs really accumulate, because every user query requires a forward pass through the model. The optimizations in this section focus on making inference as fast and memory-efficient as possible.

9.1 QUANTIZATION: DOING MORE WITH FEWER BITS

Modern neural networks are typically trained in 32-bit or 16-bit floating point. But do we actually need that much precision for inference? Quantization is the process of representing weights and/or activations with fewer bits, reducing memory usage and potentially speeding up computation.

The simplest form is post-training quantization (PTQ), where a trained model's weights are converted to lower precision without any retraining. LLM.int8() (Dettmers et al., 2022, arXiv:2208.07339) showed that 8-bit quantization of LLM weights is possible without significant quality loss, by using a decomposition that handles outlier activation values (which are common in large LLMs) in 16-bit while quantizing the rest in 8-bit. This halves the memory required for inference.

GPTQ (Frantar et al., 2022, arXiv:2210.17323) pushed this further to 3 or 4 bits per weight, using an approximate second-order optimization method to minimize the quantization error. GPTQ can quantize a 175-billion-parameter model in about four GPU hours, and the resulting 4-bit model has negligible quality degradation compared to the 16-bit original. A 4-bit quantized model requires only one-quarter the memory of its 16-bit counterpart, making it possible to run models that would otherwise require multiple GPUs on a single GPU.

SmoothQuant (Xiao et al., 2022, arXiv:2211.01524) addresses the challenge of quantizing both weights and activations (W8A8 quantization), which is necessary for maximum inference speedup on hardware that supports 8-bit integer arithmetic. The problem is that LLM activations have large outliers in certain channels, which make naive quantization very lossy. SmoothQuant migrates the quantization difficulty from activations to weights by applying a mathematically equivalent scaling transformation: multiply the activations by a per-channel smoothing factor and divide the weights by the same factor. This makes the activations easier to quantize while making the weights slightly harder, but weights are much easier to quantize than activations because they are fixed and can be calibrated offline.

The trend in 2025 and 2026 has been toward training models natively in lower precision. FlashAttention-3's FP8 support, combined with hardware like the H100 and H200 that have native FP8 tensor cores, enables training and inference in 8-bit floating point with carefully managed precision. Research into scaling laws for precision (Dettmers et al., 2024, arXiv:2501.00663) shows that lower-precision training can be compute-optimal in certain regimes, suggesting that future models may be trained in FP8 or even lower precision from the start.

9.2 SPECULATIVE DECODING: USING A SMALL MODEL TO ACCELERATE A LARGE ONE

Autoregressive text generation is inherently sequential: to generate token t+1, you need token t, which means you cannot parallelize across the time dimension. Each forward pass through a large model generates exactly one token, which is very inefficient given that GPUs are designed for massively parallel computation.

Speculative decoding (Leviathan et al., 2022, arXiv:2211.17192; Chen et al., 2023, arXiv:2302.01318) is a clever technique that breaks this sequential bottleneck. The idea is to use a small, fast draft model to generate a sequence of candidate tokens, and then use the large target model to verify all of them in a single parallel forward pass. If the large model agrees with the draft model's predictions, all the candidate tokens are accepted. If the large model disagrees at some position, that token and all subsequent ones are rejected, and the large model's prediction at the point of disagreement is used instead.

Here is a concrete example. Suppose the draft model generates the candidates "the", "cat", "sat", "on", "the" in 5 fast forward passes. The large model then processes all 5 tokens in a single parallel forward pass and checks whether it agrees. If it agrees with all 5, we have generated 5 tokens at the cost of 1 large model forward pass plus 5 small model forward passes. Since the small model is much faster than the large model, this is a significant speedup.

The key insight is that the output distribution of the combined system is identical to the output distribution of the large model alone. Speculative decoding does not change what the model says; it only changes how fast it says it. The speedup depends on the acceptance rate: how often the draft model's predictions match the large model's. For a good draft model, acceptance rates of 70 to 90 percent are achievable, leading to 2 to 4x speedup in practice.

Medusa (Cai et al., 2024, arXiv:2401.10774) is a variant that eliminates the need for a separate draft model. Instead, it adds multiple extra decoding heads to the top of the large model, each predicting a different future token. These heads are lightweight and can be trained quickly on top of a frozen base model. The predictions from the Medusa heads are verified using a tree-based attention mechanism, allowing multiple candidate continuations to be checked simultaneously. Medusa achieves 2 to 3x speedup without requiring a separate draft model.

DeepSeek-V3's multi-token prediction training objective also enables speculative decoding: the auxiliary MTP heads trained during pretraining can serve as draft heads at inference time, providing a natural and well-integrated source of speculative candidates.

9.3 MULTI-TOKEN PREDICTION AS A TRAINING OBJECTIVE

Standard LLM training uses the next-token prediction loss: for each position in the sequence, the model predicts the next token, and the loss is the cross-entropy between the prediction and the true next token. Gloeckle et al. (2024, arXiv:2404.19737) at Meta showed that training models to predict multiple future tokens simultaneously, using n independent output heads operating on a shared model trunk, results in higher sample efficiency. Models trained with 4-token prediction are faster to train and achieve better performance on code generation tasks.

The intuition is that predicting multiple future tokens forces the model to develop more structured, forward-looking representations of the current context, rather than just memorizing local patterns. This is particularly beneficial for tasks that require planning, like code generation, where the model needs to think several steps ahead. The additional output heads also serve double duty at inference time: they can be used as speculative decoding draft heads, providing a free speedup with no additional model parameters beyond what was already trained.

9.4 CONTINUOUS BATCHING AND PAGEDATTENTION

At the system level, two innovations have dramatically improved the throughput of LLM inference servers: continuous batching and PagedAttention.

Traditional batching for LLM inference processes a fixed batch of requests together and waits for all of them to finish before starting the next batch. Since different requests have different lengths, shorter requests finish early and their GPU resources sit idle while waiting for longer requests. Continuous batching (also called iteration-level scheduling) instead processes requests at the token level: as soon as a request finishes generating a token, the next request in the queue can be added to the batch for the next iteration. This keeps the GPU fully utilized at all times and dramatically increases throughput.

PagedAttention, introduced in the vLLM system (Kwon et al., 2023, arXiv:2309.06180), addresses the memory fragmentation problem in KV cache management. Different requests have different context lengths, and the KV cache for each request grows dynamically as tokens are generated. Naively allocating contiguous memory for each request's KV cache leads to severe fragmentation, wasting a large fraction of GPU memory. PagedAttention borrows the concept of virtual memory paging from operating systems: the KV cache is divided into fixed-size pages, and pages are allocated on demand. Pages from different requests can be interleaved in physical memory, eliminating fragmentation. PagedAttention also enables efficient KV cache sharing between requests that share a common prefix (e.g., a system prompt), further reducing memory usage. Together, continuous batching and PagedAttention have become the standard infrastructure for production LLM serving, enabling 10 to 20 times higher throughput compared to naive serving implementations.

CHAPTER TEN: BETTER TRAINING - NEW OPTIMIZERS AND TRAINING TECHNIQUES

The optimization algorithm used to train a neural network has a profound effect on convergence speed and final model quality. Adam and its variant AdamW have been the dominant optimizers for Transformer training since 2018, but recent work has produced serious challengers.

10.1 ADAM AND ADAMW: THE INCUMBENT

Adam (Kingma and Ba, 2014) maintains exponential moving averages of both the gradient (first moment) and the squared gradient (second moment) for each parameter, and uses these to adaptively scale the learning rate for each parameter. This makes Adam much more robust to the choice of learning rate than vanilla SGD, and it handles sparse gradients well. AdamW (Loshchilov and Hutter, 2019) adds decoupled weight decay, which regularizes the weights independently of the gradient-based update. AdamW is the standard optimizer for LLM pretraining.

The main weakness of Adam is its memory cost: it stores two additional tensors (the first and second moment estimates) for every parameter, tripling the memory required for the optimizer state. For a 70-billion-parameter model, the optimizer state alone requires roughly 560 GB at 32-bit precision. This is a significant constraint on the maximum model size that can be trained on a given hardware budget.

10.2 SOPHIA: SECOND-ORDER OPTIMIZATION MADE PRACTICAL

Sophia (Liu et al., 2023, arXiv:2305.14342) is a stochastic second-order optimizer that uses a diagonal estimate of the Hessian (the matrix of second derivatives of the loss with respect to the parameters) as a preconditioner. The Hessian captures the curvature of the loss landscape: parameters in directions of high curvature should take small steps (to avoid overshooting), while parameters in directions of low curvature can take large steps. Adam approximates curvature using the squared gradient, which is a rough proxy. Sophia uses a better estimate, computed periodically using a Hutchinson estimator. Sophia converges twice as fast as Adam in terms of the number of gradient steps, saving 50 percent of total training compute. However, computing the Hessian estimate adds overhead, and Sophia has not yet been widely adopted in frontier model training.

10.3 MUON AND MUONCLIP: ORTHOGONALIZED UPDATES

Muon (originally proposed by Keller Jordan in a 2024 blog post and formalized in arXiv:2502.16982) is a new optimizer that applies Nesterov momentum and then orthogonalizes the gradient update using Newton-Schulz iterations. The orthogonalization step ensures that the update matrix has orthonormal columns, which prevents different gradient directions from becoming correlated and leads to more efficient parameter updates. Muon achieves better final loss than AdamW at the same compute budget on language model pretraining benchmarks.

MuonClip, used in Kimi K2 (2025, arXiv:2506.01939), adds gradient clipping to Muon to prevent loss spikes during training of very large MoE models. Kimi K2's successful training of a one-trillion-parameter model with MuonClip suggests that this optimizer is robust enough for frontier-scale training. The adoption of Muon by a frontier lab for a production model is a significant signal that the optimizer landscape for LLM training is genuinely evolving beyond AdamW.

10.4 TEST-TIME COMPUTE SCALING: REASONING MODELS

One of the most important developments of 2024-2025 is the discovery that scaling compute at inference time (test-time compute) is a powerful complement to scaling compute at training time. Models like OpenAI o1, o3, and DeepSeek-R1 use reinforcement learning to train models to generate extended chains of thought before producing a final answer. By spending more tokens on reasoning, these models achieve dramatically better performance on complex tasks like mathematics, coding, and scientific reasoning.

DeepSeek-R1 (January 2025, arXiv:2501.12948) showed that a model trained with large-scale reinforcement learning, without any supervised fine-tuning data, can develop sophisticated reasoning capabilities. The model learns to break problems into steps, check its work, backtrack when it makes mistakes, and try alternative approaches, all within the context window of a single forward pass sequence. This is not a change to the Transformer architecture per se, but it represents a new paradigm for how Transformers are used: not just as pattern matchers but as deliberate reasoners.

Research into scaling laws for test-time compute (Snell et al., 2024, arXiv:2408.03314) shows that the optimal test-time compute strategy depends on problem difficulty. For easy problems, generating multiple candidate answers and selecting the best one (best-of-N sampling) is most efficient. For hard problems, iterative self-refinement (where the model critiques and improves its own answer) is better. This suggests that future systems will adaptively allocate test-time compute based on estimated problem difficulty.

CHAPTER ELEVEN: LONG-CONTEXT TRANSFORMERS - RING ATTENTION AND BEYOND

Even with FlashAttention and sparse attention, there are practical limits to the context length a single GPU can handle. For truly enormous contexts (millions of tokens), we need to distribute the computation across multiple devices. Ring Attention (Liu et al., 2023, arXiv:2310.01889) is the key technique for doing this.

The idea behind Ring Attention is to distribute the query, key, and value sequences across multiple devices arranged in a logical ring. Each device holds a chunk of the query sequence and a chunk of the key-value sequence. Attention is computed in a blockwise fashion: each device computes attention between its local query chunk and its local key-value chunk, then passes its key-value chunk to the next device in the ring while receiving the key-value chunk from the previous device. After N devices have passed their chunks around the ring, each device has computed attention between its local queries and all key-value pairs in the full sequence.

The crucial insight is that the key-value communication between devices can be overlapped with the attention computation: while device i is computing attention with the key-value chunk it currently holds, it is simultaneously sending that chunk to device i+1 and receiving the next chunk from device i-1. This communication-computation overlap means that Ring Attention adds minimal overhead compared to single-device attention, and it allows the context length to scale linearly with the number of devices.

Using Ring Attention, the context length is limited only by the total memory of the distributed cluster. With a cluster of 1,000 GPUs each with 80 GB of memory, the theoretical maximum context length is in the billions of tokens. In practice, MiniMax-01 (2025, arXiv:2501.09755) demonstrated a 32-million-token context window using a combination of linear attention and MoE, and Llama 4 Scout (2025, arXiv:2505.09343) demonstrated a 10-million-token context window using iRoPE. These are extraordinary achievements that open up entirely new applications: analyzing entire codebases, processing book-length documents, or reasoning over years of conversation history.

CHAPTER TWELVE: TITANS AND NEURAL MEMORY - LEARNING TO REMEMBER

A fundamental limitation of the Transformer is that its "memory" is entirely contained in its context window. Everything the model knows about the current task must fit within the context, and anything outside the context is simply inaccessible. This is fine for many tasks, but it is a serious limitation for tasks that require long-term memory: remembering facts from earlier in a conversation, maintaining a persistent world model, or accumulating knowledge over many interactions.

Titans (Behrouz, Zhong, and Mirrokni, 2024, arXiv:2501.00663) is a new architecture that addresses this by incorporating a neural long-term memory module into the Transformer. The long-term memory is a separate neural network that can be updated at test time: as the model processes new information, it updates the weights of the memory module to encode that information. When the model needs to recall something, it queries the memory module to retrieve relevant information.

This is a form of meta-learning or test-time adaptation: the model's parameters (specifically, the memory module's parameters) change in response to the input, allowing the model to accumulate information beyond what fits in the context window. Titans achieves better performance than standard Transformers on long-context tasks while being more efficient, because the memory module can store information compactly without requiring it to be in the attention context. Titans represents a direction that many researchers believe is essential for the next generation of AI systems: models that can learn and remember over extended periods, not just within a single context window.

CHAPTER THIRTEEN: WHERE ARE THE LIMITS?

We have now surveyed a remarkable collection of optimizations that have extended the Transformer far beyond what its original designers imagined. But every architecture has limits, and the Transformer is no exception. Let us think carefully about where those limits are.

13.1 THE QUADRATIC WALL

Even with FlashAttention and sparse attention, the fundamental computational complexity of exact self-attention is O(N^2). For N = 1 million tokens, this is 10^12 operations per layer. For N = 10 million tokens, it is 10^14. Even with the fastest hardware, exact attention over very long sequences is simply not feasible. Sparse attention and linear attention approximations can reduce this, but they come with quality tradeoffs. The quadratic wall is a genuine architectural limit that cannot be fully engineered away within the standard attention framework.

State space models like Mamba (Gu and Dao, 2023, arXiv:2312.00752) and its successor Mamba-2 (Dao and Gu, 2024, arXiv:2405.21060) offer a fundamentally different approach with O(N) complexity. Mamba uses selective state space models that can selectively remember or forget information based on the input content, achieving performance competitive with Transformers on language modeling while being 5x faster at inference. Mamba-2 showed that SSMs and attention are two instances of a more general framework called structured state space duality (SSD), and built hybrid architectures that are 2 to 8 times faster than Mamba while remaining competitive with Transformers. Griffin (De et al., 2024, arXiv:2402.19427) from Google DeepMind mixes gated linear recurrences with local attention, matching LLaMA-2 performance while being more efficient on long sequences. RWKV (Peng et al., 2023, arXiv:2305.13048) combines the efficient parallelizable training of Transformers with the efficient inference of RNNs, achieving linear complexity in both training and inference.

These hybrid architectures, combining the strengths of attention (strong performance on tasks requiring global information retrieval) with the efficiency of recurrence (linear complexity in sequence length), may represent the future of sequence modeling beyond the pure Transformer.

13.2 THE DATA WALL

Scaling laws (Kaplan et al., 2020, arXiv:2001.08361; Hoffmann et al., 2022, arXiv:2203.15556) tell us that model performance improves predictably as we increase model size, dataset size, and compute. But this scaling depends on having access to more and more high-quality training data. The internet contains a finite amount of high-quality text, and current estimates suggest that frontier LLMs will exhaust high-quality English text data by 2026 to 2028. Repeating data significantly degrades model quality, so simply re-using existing data is not a satisfactory solution.

Potential responses to the data wall include synthetic data generation (having models generate training data for themselves or for other models), multimodal data (using images, audio, video, and code as additional training signal), and data from specialized domains (scientific papers, legal documents, medical records). Each of these has its own challenges: synthetic data can introduce biases and errors, multimodal data requires new architectural components, and specialized domain data may not generalize well.

13.3 THE REASONING WALL

Current Transformers, even the best reasoning models, struggle with tasks that require systematic generalization: applying a learned rule to a new combination of concepts that was not seen during training. A model might learn that "2 + 3 = 5" and "4 + 7 = 11" but fail to correctly compute "2 + 3 + 4 + 7" if it has never seen four-term additions. This is because Transformers learn statistical patterns rather than symbolic rules, and statistical patterns do not always generalize in the way that symbolic rules do.

Achieving robust systematic generalization may require architectural innovations beyond the standard Transformer: explicit memory systems, modular reasoning components, or integration with symbolic computation. This is an active area of research, and it is one of the key open problems on the path toward more capable AI systems.

13.4 THE GROUNDING WALL

Language models learn from text, and text is a description of the world, not the world itself. A model trained only on text has no direct experience of physical reality: it has never seen a red apple, heard a dog bark, or felt the weight of a book. This lack of grounding means that language models can produce fluent, confident descriptions of things they fundamentally do not understand in the way that a human understands them.

Multimodal models (like GPT-4o, Gemini 2.5, and Llama 4) partially address this by training on images, audio, and video in addition to text. But even multimodal models are trained on recorded data, not on direct sensorimotor experience. Truly grounded AI may require embodied systems that interact with the physical world, which is a very different research direction from scaling Transformers.

13.5 THE STATIC WEIGHT WALL

Once a Transformer is trained, its weights are fixed. It cannot learn from new information without retraining or fine-tuning. This is in stark contrast to human intelligence, which continuously learns from experience. Continual learning, the ability to learn new things without forgetting old ones (catastrophic forgetting), is a major unsolved problem for neural networks in general and Transformers in particular.

Titans' neural long-term memory is one approach to this problem. Retrieval-augmented generation (RAG), where the model retrieves relevant information from an external database at inference time, is another. But neither fully solves the problem of continual learning from experience. This is likely to be one of the most important research directions of the next decade.

CHAPTER FOURTEEN: WHAT CAN WE EXPECT NEXT?

Based on the trajectory of research through mid-2026, here are the most promising directions for further improvement of Transformer-based systems.

The first direction is deeper integration of MoE and attention. Current MoE models apply sparsity only to the feed-forward layers, while attention layers remain dense. Sparse attention over experts (where different attention heads are routed to different experts) is an underexplored direction that could further improve efficiency.

The next direction is hardware-software co-design. FlashAttention succeeded by deeply understanding the GPU memory hierarchy and designing the algorithm around it. Future optimizations will likely require similar co-design between algorithm designers and hardware architects. New hardware features (like the Hopper GPU's TMA and WGMMA instructions) will continue to unlock new algorithmic possibilities. NVIDIA's Blackwell architecture (B100, B200), which began shipping in 2025, introduces further innovations in memory bandwidth and FP4/FP6 precision support that will drive the next generation of attention and MoE optimizations.

With the third direction we focus on better long-term memory. Titans showed that neural long-term memory is feasible, but the field is still in its early stages. More sophisticated memory architectures that can store, retrieve, and update information efficiently over very long time horizons will be essential for the next generation of AI assistants.

The fourth direction is native multimodality. Current multimodal models often process different modalities with separate encoders that are then combined. Truly native multimodal architectures, like the Byte Latent Transformer's approach of operating on raw bytes, could process all modalities in a unified framework without modality-specific preprocessing. Llama 4's early fusion approach, which processes text and image tokens in the same sequence from the first layer, is a step in this direction.

In the fifth direction better training algorithms are addressed. The success of Muon and MuonClip suggests that there is still significant room for improvement in the optimization algorithms used to train Transformers. Second-order methods, modular optimization (applying different optimizers to different parts of the model), and curriculum learning strategies are all active areas of research.

The sixth direction is test-time compute scaling. The reasoning model paradigm (o1, o3, DeepSeek-R1) has demonstrated that spending more compute at inference time can dramatically improve performance on complex tasks. Developing better algorithms for allocating test-time compute, better training procedures for reasoning models, and more efficient implementations of chain-of-thought reasoning will be major research priorities.

Last but not least, the seventh direction is about hybrid architectures. The convergence of SSMs and attention in frameworks like Mamba-2 and Griffin suggests that the dichotomy between recurrent and attention-based models is dissolving. Future architectures will likely be hybrids that combine the strengths of both, using attention where global information retrieval is needed and recurrence where long-range memory efficiency is needed.

EPILOGUE: THE TRANSFORMER IS NOT THE END OF HISTORY

The Transformer has been the dominant architecture in deep learning for nearly a decade, and it will likely remain dominant for several more years. But it is not the end of history. Every optimization we have discussed in this tutorial has pushed the Transformer further, but each has also revealed new limits and new problems. FlashAttention made long contexts feasible, which revealed that models needed better positional encodings. Better positional encodings enabled longer contexts, which revealed that KV caches were a bottleneck. GQA and MLA addressed KV caches, which revealed that the data wall and the reasoning wall were the next frontiers.

This is how progress works: each solution creates new problems, and those new problems drive the next wave of innovation. The Transformer is a remarkable foundation, but the systems of 2030 will likely look quite different from the systems of today, incorporating ideas from state space models, neural memory, embodied learning, and architectures we have not yet imagined.

What will remain constant is the core insight of the Transformer: that attention, the ability to selectively focus on relevant information in a sea of irrelevant information, is a powerful and general computational primitive. Whether future architectures implement this insight with softmax attention, kernel approximations, state space models, or something entirely new, the principle of selective, content-based information routing will remain central to intelligent computation.

The journey from "Attention Is All You Need" to one-trillion-parameter reasoning models with 32-million-token context windows has taken less than a decade. The next decade will be at least as surprising.

REFERENCES

Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.

Zhang, B. and Sennrich, R. (2019). Root Mean Square Layer Normalization. arXiv:1910.07467.

Shazeer, N. (2020). GLU Variants Improve Transformers. arXiv:2002.05202.

Beltagy, I. et al. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150.

Wang, S. et al. (2020). Linformer: Self-Attention with Linear Complexity. arXiv:2006.04768.

Zaheer, M. et al. (2020). Big Bird: Transformers for Longer Sequences. arXiv:2007.14062.

Choromanski, K. et al. (2020). Rethinking Attention with Performers. arXiv:2009.14794.

Kaplan, J. et al. (2020). Scaling Laws for Neural Language Models. arXiv:2001.08361.

Fedus, W. et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.

Su, J. et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.

Press, O. et al. (2021). Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi). arXiv:2108.12409.

Hoffmann, J. et al. (2022). Training Compute-Optimal Large Language Models (Chinchilla). arXiv:2203.15556.

Wang, H. et al. (2022). DeepNet: Scaling Transformers to 1,000 Layers. arXiv:2203.00555.

Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.

Dettmers, T. et al. (2022). LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale. arXiv:2208.07339.

Frantar, E. et al. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. arXiv:2210.17323.

Xiao, G. et al. (2022). SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models. arXiv:2211.01524.

Leviathan, Y. et al. (2022). Fast Inference from Transformers via Speculative Decoding. arXiv:2211.17192.

Chen, C. et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318.

Liu, H. et al. (2023). Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. arXiv:2305.14342.

Peng, B. et al. (2023). RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048.

Ainslie, J. et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245.

Chen, S. et al. (2023). Extending Context Window of Large Language Models via Positional Interpolation. arXiv:2306.15595.

Peng, B. et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models. arXiv:2309.00071.

Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.

Liu, H. et al. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context. arXiv:2310.01889.

Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.

Gu, A. and Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752.

Dai, D. et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066.

Cai, T. et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774.

Mistral AI (2024). Mixtral of Experts. arXiv:2401.04088.

De, S. et al. (2024). Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient LLMs. arXiv:2402.19427.

DeepSeek AI (2024). DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model. arXiv:2405.04434.

Dao, T. and Gu, A. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2). arXiv:2405.21060.

Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention on Hopper GPUs. arXiv:2407.08608.

Snell, C. et al. (2024). Scaling LLM Test-Time Compute Optimally. arXiv:2408.03314.

Gloeckle, F. et al. (2024). Better and Faster Large Language Models via Multi-Token Prediction. arXiv:2404.19737.

Behrouz, A., Zhong, P., and Mirrokni, V. (2024). Titans: Learning to Memorize at Test Time. arXiv:2501.00663.

Yu, K. et al. (2024). Byte Latent Transformer: Patches Scale Better Than Tokens. arXiv:2412.09871.

DeepSeek AI (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437.

DeepSeek AI (2025). DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv:2501.12948.

MiniMax (2025). MiniMax-01: Scaling Foundation Models with Lightning Attention. arXiv:2501.09755.

DeepSeek AI (2025). Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. arXiv:2502.05171.

Jordan, K. et al. (2025). Muon is Scalable for LLM Training. arXiv:2502.16982.

Meta AI (2025). The Llama 4 Herd: The Beginning of a New Era of Natively Multimodal AI at Meta. arXiv:2505.09343.

Moonshot AI (2025). Kimi K2: Open Agentic Intelligence. arXiv:2506.01939.

BUILDING AN AI-POWERED HAND-DRAWN GRAPHICS DIGITIZATION SYSTEM

INTRODUCTION

The transformation of hand-drawn graphics into clean digital representations represents a significant challenge in computer vision and artificial intelligence. This article presents a comprehensive system that accepts photographs of hand-drawn diagrams, analyzes their semantic content, recognizes diagram types, and produces optimized digital equivalents while preserving the original intent and meaning.

The system addresses several complex problems simultaneously. First, it must handle the inherent variability and imperfection of hand-drawn input, including inconsistent line weights, irregular shapes, and varying handwriting styles. Second, it must recognize high-level diagram semantics, distinguishing between flowcharts, UML diagrams, circuit schematics, and other specialized notations. Third, it must beautify and optimize the output while maintaining semantic equivalence to the original drawing. Finally, it must support diverse hardware configurations and both local and remote language model backends.

This implementation leverages open-source tools exclusively and provides production-ready code that supports Intel GPUs, AMD ROCm, Apple Metal Performance Shaders, and NVIDIA CUDA architectures. The system architecture follows clean code principles with clear separation of concerns and extensible design patterns.

SYSTEM ARCHITECTURE OVERVIEW

The digitization system consists of five primary components working in a pipeline architecture. The first component handles image preprocessing and normalization. The second performs optical character recognition and text extraction. The third detects and classifies geometric shapes and symbols. The fourth component recognizes diagram types and semantic structures. The fifth generates the final digital representation.

Each component operates independently with well-defined interfaces, allowing for parallel processing where appropriate and easy replacement of individual modules. The system uses a factory pattern for instantiating backend-specific implementations and a strategy pattern for selecting appropriate processing algorithms based on detected diagram types.

HARDWARE ABSTRACTION AND MULTI-GPU SUPPORT

Supporting multiple GPU architectures requires careful abstraction of device-specific operations. The system must detect available hardware at runtime and configure PyTorch accordingly. This involves checking for CUDA availability, ROCm support, Metal Performance Shaders on Apple Silicon, and Intel extension for PyTorch.

The device manager component handles all hardware detection and configuration. It probes the system for available accelerators and selects the most appropriate device based on a priority hierarchy. NVIDIA CUDA receives highest priority when available, followed by AMD ROCm, Apple MPS, and Intel GPU support. The system falls back to CPU execution when no GPU acceleration is available.

Here is the device manager implementation:

import torch

import platform

import subprocess

import logging

from typing import Optional, Dict, Any

from enum import Enum

class DeviceType(Enum):

"""Enumeration of supported device types."""

CUDA = "cuda"

ROCM = "rocm"

MPS = "mps"

INTEL = "xpu"

CPU = "cpu"

class DeviceManager:

"""

Manages hardware detection and device configuration for multi-GPU support.

Supports NVIDIA CUDA, AMD ROCm, Apple MPS, Intel XPU, and CPU fallback.

"""

def __init__(self):

"""Initialize device manager and detect available hardware."""

self.logger = logging.getLogger(__name__)

self.device_type = None

self.device = None

self.device_properties = {}

self._detect_and_configure()

def _check_cuda_availability(self) -> bool:

"""Check if NVIDIA CUDA is available and functional."""

try:

if torch.cuda.is_available():

# Verify CUDA actually works by attempting a simple operation

test_tensor = torch.zeros(1).cuda()

del test_tensor

torch.cuda.empty_cache()

return True

except Exception as e:

self.logger.warning(f"CUDA detected but not functional: {e}")

return False

def _check_rocm_availability(self) -> bool:

"""Check if AMD ROCm is available."""

try:

# ROCm uses the same CUDA API in PyTorch

if torch.cuda.is_available():

# Check if this is actually ROCm

device_name = torch.cuda.get_device_name(0).lower()

if 'amd' in device_name or 'radeon' in device_name:

return True

except Exception as e:

self.logger.warning(f"ROCm check failed: {e}")

return False

def _check_mps_availability(self) -> bool:

"""Check if Apple Metal Performance Shaders is available."""

try:

if platform.system() == 'Darwin':

if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():

# Verify MPS actually works

test_tensor = torch.zeros(1).to('mps')

del test_tensor

return True

except Exception as e:

self.logger.warning(f"MPS detected but not functional: {e}")

return False

def _check_intel_availability(self) -> bool:

"""Check if Intel GPU extension is available."""

try:

import intel_extension_for_pytorch as ipex

if hasattr(torch, 'xpu') and torch.xpu.is_available():

test_tensor = torch.zeros(1).to('xpu')

del test_tensor

return True

except ImportError:

self.logger.info("Intel Extension for PyTorch not installed")

except Exception as e:

self.logger.warning(f"Intel XPU check failed: {e}")

return False

def _detect_and_configure(self):

"""Detect available hardware and configure the appropriate device."""

self.logger.info("Detecting available hardware accelerators...")

# Check in priority order: CUDA > ROCm > MPS > Intel > CPU

if self._check_cuda_availability():

if self._check_rocm_availability():

self.device_type = DeviceType.ROCM

self.device = torch.device('cuda')

self.device_properties = {

'name': torch.cuda.get_device_name(0),

'compute_capability': torch.cuda.get_device_capability(0),

'total_memory': torch.cuda.get_device_properties(0).total_memory,

'backend': 'ROCm'

}

self.logger.info(f"Using AMD ROCm: {self.device_properties['name']}")

else:

self.device_type = DeviceType.CUDA

self.device = torch.device('cuda')

self.device_properties = {

'name': torch.cuda.get_device_name(0),

'compute_capability': torch.cuda.get_device_capability(0),

'total_memory': torch.cuda.get_device_properties(0).total_memory,

'backend': 'CUDA'

}

self.logger.info(f"Using NVIDIA CUDA: {self.device_properties['name']}")

elif self._check_mps_availability():

self.device_type = DeviceType.MPS

self.device = torch.device('mps')

self.device_properties = {

'name': 'Apple Silicon GPU',

'backend': 'MPS'

}

self.logger.info("Using Apple Metal Performance Shaders")

elif self._check_intel_availability():

self.device_type = DeviceType.INTEL

self.device = torch.device('xpu')

self.device_properties = {

'name': 'Intel GPU',

'backend': 'Intel Extension for PyTorch'

}

self.logger.info("Using Intel GPU acceleration")

else:

self.device_type = DeviceType.CPU

self.device = torch.device('cpu')

self.device_properties = {

'name': 'CPU',

'backend': 'CPU'

}

self.logger.warning("No GPU acceleration available, using CPU")

def get_device(self) -> torch.device:

"""Return the configured PyTorch device."""

return self.device

def get_device_type(self) -> DeviceType:

"""Return the device type enumeration."""

return self.device_type

def get_device_info(self) -> Dict[str, Any]:

"""Return detailed device properties."""

return self.device_properties.copy()

def optimize_for_inference(self):

"""Apply device-specific optimizations for inference."""

if self.device_type == DeviceType.CUDA or self.device_type == DeviceType.ROCM:

torch.backends.cudnn.benchmark = True

torch.backends.cudnn.deterministic = False

elif self.device_type == DeviceType.MPS:

# MPS-specific optimizations

pass

elif self.device_type == DeviceType.INTEL:

try:

import intel_extension_for_pytorch as ipex

ipex.optimize(optimizer=None)

except Exception as e:

self.logger.warning(f"Intel optimization failed: {e}")

The device manager encapsulates all hardware-specific logic in a single component. It performs actual device testing rather than relying solely on availability flags, ensuring that selected devices are truly functional. The optimization method applies backend-specific performance tuning, such as enabling cuDNN benchmarking for CUDA devices.

LANGUAGE MODEL BACKEND ABSTRACTION

Supporting both local and remote language models requires a flexible backend architecture. The system must accommodate different API interfaces, authentication mechanisms, and response formats while presenting a unified interface to higher-level components.

The language model backend uses an abstract base class defining the interface that all implementations must satisfy. Concrete implementations handle specific backends such as local Llama models, OpenAI GPT-4 Vision, and other vision-language models. The factory pattern instantiates appropriate backends based on configuration.

from abc import ABC, abstractmethod

from typing import List, Dict, Any, Optional, Union

import base64

import io

from PIL import Image

import torch

from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer

import requests

import json

class VisionLanguageBackend(ABC):

"""Abstract base class for vision-language model backends."""

@abstractmethod

def analyze_image(self, image: Image.Image, prompt: str,

max_tokens: int = 1000) -> str:

"""

Analyze an image with a text prompt and return the model's response.

Args:

image: PIL Image object to analyze

prompt: Text prompt describing the analysis task

max_tokens: Maximum tokens in the response

Returns:

String response from the model

"""

pass

@abstractmethod

def get_backend_info(self) -> Dict[str, Any]:

"""Return information about the backend configuration."""

pass

class LocalLlamaVisionBackend(VisionLanguageBackend):

"""

Backend for local Llama vision models (e.g., LLaVA, Llama 3.2 Vision).

Supports multi-GPU architectures through the device manager.

"""

def __init__(self, model_name: str, device_manager: DeviceManager):

"""

Initialize local Llama vision backend.

Args:

model_name: HuggingFace model identifier

device_manager: Configured device manager instance

"""

self.model_name = model_name

self.device_manager = device_manager

self.device = device_manager.get_device()

self.logger = logging.getLogger(__name__)

self.logger.info(f"Loading local model: {model_name}")

# Load model and processor

self.processor = AutoProcessor.from_pretrained(model_name)

# Configure model loading based on device type

load_kwargs = {'torch_dtype': torch.float16}

if device_manager.get_device_type() == DeviceType.CUDA:

load_kwargs['device_map'] = 'auto'

elif device_manager.get_device_type() == DeviceType.MPS:

# MPS doesn't support float16 for all operations

load_kwargs['torch_dtype'] = torch.float32

self.model = AutoModelForCausalLM.from_pretrained(

model_name,

**load_kwargs

)

if device_manager.get_device_type() != DeviceType.CUDA:

self.model = self.model.to(self.device)

self.model.eval()

self.logger.info(f"Model loaded successfully on {self.device}")

def analyze_image(self, image: Image.Image, prompt: str,

max_tokens: int = 1000) -> str:

"""Analyze image using local vision-language model."""

try:

# Prepare inputs

messages = [

{

"role": "user",

"content": [

{"type": "image"},

{"type": "text", "text": prompt}

]

}

]

# Process inputs

inputs = self.processor(

text=self.processor.apply_chat_template(messages, add_generation_prompt=True),

images=image,

return_tensors="pt"

).to(self.device)

# Generate response

with torch.no_grad():

output = self.model.generate(

**inputs,

max_new_tokens=max_tokens,

do_sample=False

)

# Decode response

response = self.processor.decode(output[0], skip_special_tokens=True)

# Extract assistant response

if "assistant" in response:

response = response.split("assistant")[-1].strip()

return response

except Exception as e:

self.logger.error(f"Error during image analysis: {e}")

raise

def get_backend_info(self) -> Dict[str, Any]:

"""Return backend configuration information."""

return {

'backend_type': 'local_llama_vision',

'model_name': self.model_name,

'device': str(self.device),

'device_type': self.device_manager.get_device_type().value

}

class OpenAIVisionBackend(VisionLanguageBackend):

"""Backend for OpenAI GPT-4 Vision API."""

def __init__(self, api_key: str, model: str = "gpt-4-vision-preview"):

"""

Initialize OpenAI vision backend.

Args:

api_key: OpenAI API key

model: Model identifier (default: gpt-4-vision-preview)

"""

self.api_key = api_key

self.model = model

self.api_url = "https://api.openai.com/v1/chat/completions"

self.logger = logging.getLogger(__name__)

def _encode_image(self, image: Image.Image) -> str:

"""Encode PIL Image to base64 string."""

buffered = io.BytesIO()

image.save(buffered, format="PNG")

return base64.b64encode(buffered.getvalue()).decode('utf-8')

def analyze_image(self, image: Image.Image, prompt: str,

max_tokens: int = 1000) -> str:

"""Analyze image using OpenAI GPT-4 Vision API."""

try:

base64_image = self._encode_image(image)

headers = {

"Content-Type": "application/json",

"Authorization": f"Bearer {self.api_key}"

}

payload = {

"model": self.model,

"messages": [

{

"role": "user",

"content": [

{

"type": "text",

"text": prompt

{

"type": "image_url",

"image_url": {

"url": f"data:image/png;base64,{base64_image}"

}

]

}

"max_tokens": max_tokens

}

response = requests.post(self.api_url, headers=headers, json=payload)

response.raise_for_status()

result = response.json()

return result['choices'][0]['message']['content']

except Exception as e:

self.logger.error(f"Error calling OpenAI API: {e}")

raise

def get_backend_info(self) -> Dict[str, Any]:

"""Return backend configuration information."""

return {

'backend_type': 'openai_vision',

'model': self.model,

'api_url': self.api_url

}

class VisionBackendFactory:

"""Factory for creating vision-language model backends."""

@staticmethod

def create_backend(backend_type: str, device_manager: Optional[DeviceManager] = None,

**kwargs) -> VisionLanguageBackend:

"""

Create a vision-language backend based on type.

Args:

backend_type: Type of backend ('local_llama', 'openai', etc.)

device_manager: Device manager for local models

**kwargs: Additional backend-specific arguments

Returns:

Configured VisionLanguageBackend instance

"""

if backend_type == 'local_llama':

if device_manager is None:

device_manager = DeviceManager()

model_name = kwargs.get('model_name', 'llava-hf/llava-1.5-7b-hf')

return LocalLlamaVisionBackend(model_name, device_manager)

elif backend_type == 'openai':

api_key = kwargs.get('api_key')

if not api_key:

raise ValueError("OpenAI backend requires 'api_key' parameter")

model = kwargs.get('model', 'gpt-4-vision-preview')

return OpenAIVisionBackend(api_key, model)

else:

raise ValueError(f"Unknown backend type: {backend_type}")

This backend architecture provides complete flexibility in model selection. The abstract base class ensures consistent interfaces across implementations. Local models leverage the device manager for hardware acceleration, while remote backends handle API communication and authentication. The factory pattern simplifies backend instantiation and configuration management.

IMAGE PREPROCESSING AND NORMALIZATION

Raw photographs of hand-drawn diagrams require substantial preprocessing before analysis. Images may suffer from perspective distortion, uneven lighting, shadows, background noise, and varying resolutions. The preprocessing pipeline must correct these issues while preserving the essential features of the drawing.

The preprocessing component performs several operations in sequence. First, it converts the image to grayscale and applies adaptive thresholding to separate foreground content from background. Second, it detects and corrects perspective distortion by identifying the document boundaries. Third, it applies noise reduction through morphological operations. Fourth, it normalizes the image resolution and aspect ratio. Finally, it enhances contrast to improve feature detection in subsequent stages.

import cv2

import numpy as np

from typing import Tuple, Optional

from PIL import Image

class ImagePreprocessor:

"""

Handles preprocessing of hand-drawn diagram photographs.

Corrects perspective, removes noise, and normalizes images.

"""

def __init__(self, target_size: Tuple[int, int] = (1024, 1024)):

"""

Initialize image preprocessor.

Args:

target_size: Target dimensions for normalized output (width, height)

"""

self.target_size = target_size

self.logger = logging.getLogger(__name__)

def preprocess(self, image: Union[Image.Image, np.ndarray]) -> Image.Image:

"""

Complete preprocessing pipeline for hand-drawn diagram images.

Args:

image: Input image as PIL Image or numpy array

Returns:

Preprocessed PIL Image

"""

# Convert to numpy array if needed

if isinstance(image, Image.Image):

img_array = np.array(image)

else:

img_array = image

# Convert to grayscale if needed

if len(img_array.shape) == 3:

gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)

else:

gray = img_array

# Detect and correct perspective

corrected = self._correct_perspective(gray)

# Apply adaptive thresholding

binary = self._adaptive_threshold(corrected)

# Remove noise

denoised = self._remove_noise(binary)

# Normalize size

normalized = self._normalize_size(denoised)

# Enhance contrast

enhanced = self._enhance_contrast(normalized)

# Convert back to PIL Image

return Image.fromarray(enhanced)

def _correct_perspective(self, image: np.ndarray) -> np.ndarray:

"""

Detect document boundaries and correct perspective distortion.

Args:

image: Grayscale input image

Returns:

Perspective-corrected image

"""

# Apply edge detection

edges = cv2.Canny(image, 50, 150, apertureSize=3)

# Dilate edges to close gaps

kernel = np.ones((5, 5), np.uint8)

dilated = cv2.dilate(edges, kernel, iterations=1)

# Find contours

contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL,

cv2.CHAIN_APPROX_SIMPLE)

if not contours:

self.logger.warning("No contours found for perspective correction")

return image

# Find largest contour (assumed to be document boundary)

largest_contour = max(contours, key=cv2.contourArea)

# Approximate contour to polygon

epsilon = 0.02 * cv2.arcLength(largest_contour, True)

approx = cv2.approxPolyDP(largest_contour, epsilon, True)

# If we found a quadrilateral, apply perspective transform

if len(approx) == 4:

return self._apply_perspective_transform(image, approx)

else:

self.logger.info("Document boundary not quadrilateral, skipping perspective correction")

return image

def _apply_perspective_transform(self, image: np.ndarray,

corners: np.ndarray) -> np.ndarray:

"""

Apply perspective transformation to correct distortion.

Args:

image: Input image

corners: Four corner points of the document

Returns:

Transformed image

"""

# Reshape corners

corners = corners.reshape(4, 2)

# Order corners: top-left, top-right, bottom-right, bottom-left

rect = self._order_points(corners)

# Calculate dimensions of corrected image

width_a = np.linalg.norm(rect[2] - rect[3])

width_b = np.linalg.norm(rect[1] - rect[0])

max_width = max(int(width_a), int(width_b))

height_a = np.linalg.norm(rect[1] - rect[2])

height_b = np.linalg.norm(rect[0] - rect[3])

max_height = max(int(height_a), int(height_b))

# Define destination points

dst = np.array([

[0, 0],

[max_width - 1, 0],

[max_width - 1, max_height - 1],

[0, max_height - 1]

], dtype=np.float32)

# Calculate perspective transform matrix

matrix = cv2.getPerspectiveTransform(rect.astype(np.float32), dst)

# Apply transformation

warped = cv2.warpPerspective(image, matrix, (max_width, max_height))

return warped

def _order_points(self, points: np.ndarray) -> np.ndarray:

"""

Order points in clockwise order starting from top-left.

Args:

points: Array of 4 points

Returns:

Ordered points array

"""

rect = np.zeros((4, 2), dtype=np.float32)

# Sum and difference to find corners

s = points.sum(axis=1)

diff = np.diff(points, axis=1)

rect[0] = points[np.argmin(s)] # Top-left has smallest sum

rect[2] = points[np.argmax(s)] # Bottom-right has largest sum

rect[1] = points[np.argmin(diff)] # Top-right has smallest difference

rect[3] = points[np.argmax(diff)] # Bottom-left has largest difference

return rect

def _adaptive_threshold(self, image: np.ndarray) -> np.ndarray:

"""

Apply adaptive thresholding to separate foreground from background.

Args:

image: Grayscale input image

Returns:

Binary image

"""

# Apply Gaussian blur to reduce noise before thresholding

blurred = cv2.GaussianBlur(image, (5, 5), 0)

# Apply adaptive threshold

binary = cv2.adaptiveThreshold(

blurred,

255,

cv2.ADAPTIVE_THRESH_GAUSSIAN_C,

cv2.THRESH_BINARY,

11,

)

return binary

def _remove_noise(self, image: np.ndarray) -> np.ndarray:

"""

Remove noise using morphological operations.

Args:

image: Binary input image

Returns:

Denoised image

"""

# Define morphological kernels

small_kernel = np.ones((3, 3), np.uint8)

# Remove small noise with opening

opened = cv2.morphologyEx(image, cv2.MORPH_OPEN, small_kernel, iterations=1)

# Close small gaps with closing

closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, small_kernel, iterations=1)

return closed

def _normalize_size(self, image: np.ndarray) -> np.ndarray:

"""

Normalize image to target size while preserving aspect ratio.

Args:

image: Input image

Returns:

Resized image

"""

height, width = image.shape[:2]

target_width, target_height = self.target_size

# Calculate scaling factor to fit within target size

scale = min(target_width / width, target_height / height)

# Calculate new dimensions

new_width = int(width * scale)

new_height = int(height * scale)

# Resize image

resized = cv2.resize(image, (new_width, new_height),

interpolation=cv2.INTER_AREA)

# Create canvas with target size

canvas = np.ones((target_height, target_width), dtype=np.uint8) * 255

# Calculate position to center the image

y_offset = (target_height - new_height) // 2

x_offset = (target_width - new_width) // 2

# Place resized image on canvas

canvas[y_offset:y_offset+new_height, x_offset:x_offset+new_width] = resized

return canvas

def _enhance_contrast(self, image: np.ndarray) -> np.ndarray:

"""

Enhance image contrast using CLAHE.

Args:

image: Input image

Returns:

Contrast-enhanced image

"""

# Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)

clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))

enhanced = clahe.apply(image)

return enhanced

The preprocessing pipeline transforms raw photographs into clean, normalized images suitable for analysis. Perspective correction ensures that diagrams are viewed from a frontal orientation, eliminating distortion caused by camera angle. Adaptive thresholding handles varying lighting conditions better than global thresholding. Morphological operations remove isolated noise pixels while preserving the structure of drawn elements. Size normalization ensures consistent processing regardless of input resolution.

OPTICAL CHARACTER RECOGNITION AND TEXT EXTRACTION

Hand-drawn diagrams frequently contain text labels, annotations, and descriptions. Accurate text extraction is essential for understanding diagram semantics and preserving information in the digital output. The system must handle various handwriting styles, text orientations, and font sizes.

The text extraction component combines multiple OCR engines to maximize accuracy. It uses both Tesseract OCR for printed-style text and EasyOCR for handwritten text recognition. The component detects text regions, extracts text content, and associates text with spatial locations for later use in diagram reconstruction.

import pytesseract

import easyocr

from typing import List, Dict, Tuple

import numpy as np

from dataclasses import dataclass

@dataclass

class TextRegion:

"""Represents a detected text region with content and location."""

text: str

confidence: float

bbox: Tuple[int, int, int, int] # (x, y, width, height)

orientation: float # Angle in degrees

class TextExtractor:

"""

Extracts text from hand-drawn diagrams using multiple OCR engines.

Combines Tesseract and EasyOCR for robust text detection.

"""

def __init__(self, languages: List[str] = ['en']):

"""

Initialize text extractor.

Args:

languages: List of language codes for OCR

"""

self.languages = languages

self.logger = logging.getLogger(__name__)

# Initialize EasyOCR reader

self.logger.info("Initializing EasyOCR reader...")

self.easyocr_reader = easyocr.Reader(languages, gpu=True)

self.logger.info("EasyOCR reader initialized")

def extract_text(self, image: Union[Image.Image, np.ndarray]) -> List[TextRegion]:

"""

Extract all text regions from an image.

Args:

image: Input image as PIL Image or numpy array

Returns:

List of TextRegion objects containing detected text

"""

# Convert to numpy array if needed

if isinstance(image, Image.Image):

img_array = np.array(image)

else:

img_array = image

# Extract text using both engines

tesseract_results = self._extract_with_tesseract(img_array)

easyocr_results = self._extract_with_easyocr(img_array)

# Merge results, preferring higher confidence detections

merged_results = self._merge_text_regions(tesseract_results, easyocr_results)

self.logger.info(f"Extracted {len(merged_results)} text regions")

return merged_results

def _extract_with_tesseract(self, image: np.ndarray) -> List[TextRegion]:

"""

Extract text using Tesseract OCR.

Args:

image: Input image as numpy array

Returns:

List of TextRegion objects

"""

results = []

try:

# Get detailed OCR data

data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)

# Process each detected word

n_boxes = len(data['text'])

for i in range(n_boxes):

text = data['text'][i].strip()

if text: # Only process non-empty text

confidence = float(data['conf'][i])

if confidence > 0: # Filter out low-confidence detections

x = data['left'][i]

y = data['top'][i]

w = data['width'][i]

h = data['height'][i]

results.append(TextRegion(

text=text,

confidence=confidence / 100.0, # Normalize to 0-1

bbox=(x, y, w, h),

orientation=0.0

))

except Exception as e:

self.logger.error(f"Tesseract OCR error: {e}")

return results

def _extract_with_easyocr(self, image: np.ndarray) -> List[TextRegion]:

"""

Extract text using EasyOCR.

Args:

image: Input image as numpy array

Returns:

List of TextRegion objects

"""

results = []

try:

# Perform OCR

detections = self.easyocr_reader.readtext(image)

# Process each detection

for detection in detections:

bbox_points, text, confidence = detection

# Convert bbox points to x, y, w, h format

bbox_array = np.array(bbox_points)

x = int(bbox_array[:, 0].min())

y = int(bbox_array[:, 1].min())

w = int(bbox_array[:, 0].max() - x)

h = int(bbox_array[:, 1].max() - y)

# Calculate orientation from bbox points

orientation = self._calculate_text_orientation(bbox_points)

results.append(TextRegion(

text=text,

confidence=confidence,

bbox=(x, y, w, h),

orientation=orientation

))

except Exception as e:

self.logger.error(f"EasyOCR error: {e}")

return results

def _calculate_text_orientation(self, bbox_points: List[List[float]]) -> float:

"""

Calculate text orientation angle from bounding box points.

Args:

bbox_points: List of 4 corner points

Returns:

Orientation angle in degrees

"""

# Calculate angle from top-left to top-right point

p1 = np.array(bbox_points[0])

p2 = np.array(bbox_points[1])

dx = p2[0] - p1[0]

dy = p2[1] - p1[1]

angle = np.degrees(np.arctan2(dy, dx))

return angle

def _merge_text_regions(self, regions1: List[TextRegion],

regions2: List[TextRegion]) -> List[TextRegion]:

"""

Merge text regions from multiple OCR engines, removing duplicates.

Args:

regions1: First list of text regions

regions2: Second list of text regions

Returns:

Merged list with duplicates removed

"""

merged = []

used_indices = set()

# Start with all regions from first list

for r1 in regions1:

merged.append(r1)

# Add regions from second list that don't overlap significantly

for r2 in regions2:

overlaps = False

for r1 in regions1:

if self._regions_overlap(r1, r2):

overlaps = True

break

if not overlaps:

merged.append(r2)

# Sort by position (top to bottom, left to right)

merged.sort(key=lambda r: (r.bbox[1], r.bbox[0]))

return merged

def _regions_overlap(self, r1: TextRegion, r2: TextRegion,

threshold: float = 0.5) -> bool:

"""

Check if two text regions overlap significantly.

Args:

r1: First text region

r2: Second text region

threshold: Minimum IoU to consider regions overlapping

Returns:

True if regions overlap above threshold

"""

x1, y1, w1, h1 = r1.bbox

x2, y2, w2, h2 = r2.bbox

# Calculate intersection

x_left = max(x1, x2)

y_top = max(y1, y2)

x_right = min(x1 + w1, x2 + w2)

y_bottom = min(y1 + h1, y2 + h2)

if x_right < x_left or y_bottom < y_top:

return False

intersection_area = (x_right - x_left) * (y_bottom - y_top)

# Calculate union

area1 = w1 * h1

area2 = w2 * h2

union_area = area1 + area2 - intersection_area

# Calculate IoU

iou = intersection_area / union_area if union_area > 0 else 0

return iou > threshold

The text extraction component provides robust text detection by combining complementary OCR engines. Tesseract excels at printed-style text with clear letterforms, while EasyOCR handles handwritten text more effectively. The merging algorithm eliminates duplicate detections by calculating intersection-over-union between bounding boxes. Orientation detection enables proper handling of rotated text labels.

SHAPE DETECTION AND CLASSIFICATION

Geometric shapes form the fundamental building blocks of most diagrams. Circles represent states or nodes, rectangles indicate processes or components, arrows show relationships and flow, and various specialized shapes convey domain-specific semantics. Accurate shape detection and classification is critical for understanding diagram structure.

The shape detection component uses computer vision techniques to identify geometric primitives in the preprocessed image. It detects lines, circles, rectangles, polygons, and curves using contour analysis and Hough transforms. Each detected shape is classified and parameterized for later use in digital reconstruction.

from typing import List, Dict, Any, Optional

from dataclasses import dataclass

from enum import Enum

import cv2

import numpy as np

class ShapeType(Enum):

"""Enumeration of detectable shape types."""

LINE = "line"

ARROW = "arrow"

RECTANGLE = "rectangle"

CIRCLE = "circle"

ELLIPSE = "ellipse"

TRIANGLE = "triangle"

DIAMOND = "diamond"

POLYGON = "polygon"

CURVE = "curve"

UNKNOWN = "unknown"

@dataclass

class DetectedShape:

"""Represents a detected geometric shape."""

shape_type: ShapeType

confidence: float

contour: np.ndarray

parameters: Dict[str, Any] # Shape-specific parameters

bbox: Tuple[int, int, int, int] # (x, y, width, height)

class ShapeDetector:

"""

Detects and classifies geometric shapes in hand-drawn diagrams.

Uses contour analysis and Hough transforms for robust detection.

"""

def __init__(self, min_area: int = 100, max_area: Optional[int] = None):

"""

Initialize shape detector.

Args:

min_area: Minimum contour area to consider

max_area: Maximum contour area to consider (None for no limit)

"""

self.min_area = min_area

self.max_area = max_area

self.logger = logging.getLogger(__name__)

def detect_shapes(self, image: Union[Image.Image, np.ndarray]) -> List[DetectedShape]:

"""

Detect all shapes in an image.

Args:

image: Input image as PIL Image or numpy array

Returns:

List of DetectedShape objects

"""

# Convert to numpy array if needed

if isinstance(image, Image.Image):

img_array = np.array(image)

else:

img_array = image

# Ensure binary image

if len(img_array.shape) == 3:

gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)

else:

gray = img_array

_, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)

# Find contours

contours, hierarchy = cv2.findContours(

binary,

cv2.RETR_TREE,

cv2.CHAIN_APPROX_SIMPLE

)

shapes = []

# Process each contour

for i, contour in enumerate(contours):

area = cv2.contourArea(contour)

# Filter by area

if area < self.min_area:

continue

if self.max_area and area > self.max_area:

continue

# Classify shape

shape = self._classify_shape(contour, hierarchy[0][i] if hierarchy is not None else None)

if shape:

shapes.append(shape)

# Detect lines and arrows separately using Hough transform

lines = self._detect_lines(binary)

shapes.extend(lines)

self.logger.info(f"Detected {len(shapes)} shapes")

return shapes

def _classify_shape(self, contour: np.ndarray,

hierarchy_info: Optional[np.ndarray]) -> Optional[DetectedShape]:

"""

Classify a contour as a specific shape type.

Args:

contour: Contour points

hierarchy_info: Hierarchy information for this contour

Returns:

DetectedShape object or None if classification fails

"""

# Calculate contour properties

area = cv2.contourArea(contour)

perimeter = cv2.arcLength(contour, True)

if perimeter == 0:

return None

# Approximate contour to polygon

epsilon = 0.04 * perimeter

approx = cv2.approxPolyDP(contour, epsilon, True)

num_vertices = len(approx)

# Get bounding box

x, y, w, h = cv2.boundingRect(contour)

bbox = (x, y, w, h)

# Calculate circularity

circularity = 4 * np.pi * area / (perimeter * perimeter)

# Classify based on properties

if num_vertices == 3:

return self._create_triangle(contour, approx, bbox)

elif num_vertices == 4:

return self._classify_quadrilateral(contour, approx, bbox)

elif circularity > 0.85:

return self._classify_circular(contour, bbox)

elif num_vertices > 4:

if num_vertices < 8 and circularity < 0.7:

return DetectedShape(

shape_type=ShapeType.POLYGON,

confidence=0.8,

contour=contour,

parameters={'vertices': approx.reshape(-1, 2).tolist()},

bbox=bbox

)

else:

return self._classify_curve(contour, bbox)

return None

def _classify_quadrilateral(self, contour: np.ndarray, approx: np.ndarray,

bbox: Tuple[int, int, int, int]) -> DetectedShape:

"""

Classify a 4-sided polygon as rectangle or diamond.

Args:

contour: Original contour

approx: Approximated polygon

bbox: Bounding box

Returns:

DetectedShape object

"""

x, y, w, h = bbox

aspect_ratio = float(w) / h if h > 0 else 0

# Calculate angles between edges

vertices = approx.reshape(-1, 2)

angles = []

for i in range(4):

p1 = vertices[i]

p2 = vertices[(i + 1) % 4]

p3 = vertices[(i + 2) % 4]

v1 = p1 - p2

v2 = p3 - p2

angle = np.abs(np.arccos(

np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)

))

angles.append(np.degrees(angle))

# Check if angles are close to 90 degrees (rectangle)

right_angles = sum(1 for angle in angles if 80 < angle < 100)

if right_angles >= 3:

return DetectedShape(

shape_type=ShapeType.RECTANGLE,

confidence=0.9,

contour=contour,

parameters={

'x': x, 'y': y, 'width': w, 'height': h,

'aspect_ratio': aspect_ratio

bbox=bbox

)

else:

# Check if it's a diamond (rotated square)

center_x = x + w // 2

center_y = y + h // 2

return DetectedShape(

shape_type=ShapeType.DIAMOND,

confidence=0.85,

contour=contour,

parameters={

'center_x': center_x,

'center_y': center_y,

'width': w,

'height': h,

'vertices': vertices.tolist()

bbox=bbox

)

def _classify_circular(self, contour: np.ndarray,

bbox: Tuple[int, int, int, int]) -> DetectedShape:

"""

Classify a circular shape as circle or ellipse.

Args:

contour: Contour points

bbox: Bounding box

Returns:

DetectedShape object

"""

x, y, w, h = bbox

aspect_ratio = float(w) / h if h > 0 else 0

# Fit ellipse if contour has enough points

if len(contour) >= 5:

ellipse = cv2.fitEllipse(contour)

center, axes, angle = ellipse

# Check if it's a circle (axes approximately equal)

axis_ratio = min(axes) / max(axes) if max(axes) > 0 else 0

if axis_ratio > 0.9:

radius = (axes[0] + axes[1]) / 4 # Average radius

return DetectedShape(

shape_type=ShapeType.CIRCLE,

confidence=0.95,

contour=contour,

parameters={

'center_x': int(center[0]),

'center_y': int(center[1]),

'radius': int(radius)

bbox=bbox

)

else:

return DetectedShape(

shape_type=ShapeType.ELLIPSE,

confidence=0.9,

contour=contour,

parameters={

'center_x': int(center[0]),

'center_y': int(center[1]),

'major_axis': int(max(axes)),

'minor_axis': int(min(axes)),

'angle': angle

bbox=bbox

)

# Fallback to circle approximation

(cx, cy), radius = cv2.minEnclosingCircle(contour)

return DetectedShape(

shape_type=ShapeType.CIRCLE,

confidence=0.8,

contour=contour,

parameters={

'center_x': int(cx),

'center_y': int(cy),

'radius': int(radius)

bbox=bbox

)

def _create_triangle(self, contour: np.ndarray, approx: np.ndarray,

bbox: Tuple[int, int, int, int]) -> DetectedShape:

"""

Create triangle shape object.

Args:

contour: Original contour

approx: Approximated triangle vertices

bbox: Bounding box

Returns:

DetectedShape object

"""

vertices = approx.reshape(-1, 2)

return DetectedShape(

shape_type=ShapeType.TRIANGLE,

confidence=0.9,

contour=contour,

parameters={'vertices': vertices.tolist()},

bbox=bbox

)

def _classify_curve(self, contour: np.ndarray,

bbox: Tuple[int, int, int, int]) -> DetectedShape:

"""

Classify a contour as a curve.

Args:

contour: Contour points

bbox: Bounding box

Returns:

DetectedShape object

"""

# Sample points along the curve

num_samples = min(50, len(contour))

indices = np.linspace(0, len(contour) - 1, num_samples, dtype=int)

sampled_points = contour[indices].reshape(-1, 2)

return DetectedShape(

shape_type=ShapeType.CURVE,

confidence=0.75,

contour=contour,

parameters={'points': sampled_points.tolist()},

bbox=bbox

)

def _detect_lines(self, binary_image: np.ndarray) -> List[DetectedShape]:

"""

Detect straight lines using Hough transform.

Args:

binary_image: Binary input image

Returns:

List of DetectedShape objects for lines

"""

lines_shapes = []

# Apply Hough Line Transform

lines = cv2.HoughLinesP(

binary_image,

rho=1,

theta=np.pi / 180,

threshold=50,

minLineLength=30,

maxLineGap=10

)

if lines is not None:

for line in lines:

x1, y1, x2, y2 = line[0]

# Check if line has an arrowhead

is_arrow, arrow_params = self._detect_arrowhead(

binary_image, x1, y1, x2, y2

)

# Calculate bounding box

x_min = min(x1, x2)

y_min = min(y1, y2)

x_max = max(x1, x2)

y_max = max(y1, y2)

bbox = (x_min, y_min, x_max - x_min, y_max - y_min)

if is_arrow:

lines_shapes.append(DetectedShape(

shape_type=ShapeType.ARROW,

confidence=0.85,

contour=np.array([[x1, y1], [x2, y2]]),

parameters={

'start_x': x1, 'start_y': y1,

'end_x': x2, 'end_y': y2,

**arrow_params

bbox=bbox

))

else:

lines_shapes.append(DetectedShape(

shape_type=ShapeType.LINE,

confidence=0.9,

contour=np.array([[x1, y1], [x2, y2]]),

parameters={

'start_x': x1, 'start_y': y1,

'end_x': x2, 'end_y': y2

bbox=bbox

))

return lines_shapes

def _detect_arrowhead(self, image: np.ndarray, x1: int, y1: int,

x2: int, y2: int) -> Tuple[bool, Dict[str, Any]]:

"""

Detect if a line has an arrowhead at the end.

Args:

image: Binary image

x1, y1: Line start point

x2, y2: Line end point

Returns:

Tuple of (is_arrow, arrow_parameters)

"""

# Define search region at line end

search_radius = 20

# Extract region around line end

y_min = max(0, y2 - search_radius)

y_max = min(image.shape[0], y2 + search_radius)

x_min = max(0, x2 - search_radius)

x_max = min(image.shape[1], x2 + search_radius)

region = image[y_min:y_max, x_min:x_max]

# Find contours in region

contours, _ = cv2.findContours(region, cv2.RETR_EXTERNAL,

cv2.CHAIN_APPROX_SIMPLE)

# Look for triangular shapes near line end

for contour in contours:

area = cv2.contourArea(contour)

if 10 < area < 200: # Reasonable arrowhead size

epsilon = 0.1 * cv2.arcLength(contour, True)

approx = cv2.approxPolyDP(contour, epsilon, True)

if len(approx) == 3: # Triangle

return True, {'arrowhead_type': 'triangle'}

return False, {}

The shape detector provides comprehensive geometric analysis of hand-drawn diagrams. Contour approximation reduces complex shapes to their essential vertices, enabling robust classification. The circularity metric distinguishes between polygons and circular shapes. Hough line detection finds straight lines that might be missed by contour analysis. Arrowhead detection identifies directional connectors, which carry important semantic meaning in flowcharts and other diagram types.

DIAGRAM TYPE RECOGNITION

Different diagram types follow distinct visual conventions and semantic rules. UML class diagrams use rectangles divided into compartments. Flowcharts employ specific shapes to represent different operation types. Circuit diagrams use standardized symbols for electronic components. Recognizing the diagram type enables the system to apply domain-specific beautification rules and ensure semantic correctness.

The diagram classifier uses a vision-language model to analyze the overall structure and identify the diagram type. It examines shape distributions, spatial relationships, and detected symbols to classify the input as a specific diagram type or generic drawing.

from typing import List, Dict, Any, Optional

from enum import Enum

from dataclasses import dataclass

class DiagramType(Enum):

"""Enumeration of recognizable diagram types."""

FLOWCHART = "flowchart"

UML_CLASS = "uml_class"

UML_SEQUENCE = "uml_sequence"

UML_ACTIVITY = "uml_activity"

CIRCUIT = "circuit"

NETWORK = "network"

ER_DIAGRAM = "er_diagram"

MINDMAP = "mindmap"

ORGANIZATIONAL = "organizational"

GENERIC = "generic"

@dataclass

class DiagramAnalysis:

"""Results of diagram type recognition."""

diagram_type: DiagramType

confidence: float

characteristics: Dict[str, Any]

suggested_rules: Dict[str, Any]

class DiagramClassifier:

"""

Classifies hand-drawn diagrams into specific types.

Uses vision-language models and heuristic analysis.

"""

def __init__(self, vision_backend: VisionLanguageBackend):

"""

Initialize diagram classifier.

Args:

vision_backend: Vision-language model backend for analysis

"""

self.vision_backend = vision_backend

self.logger = logging.getLogger(__name__)

def classify_diagram(self, image: Image.Image,

shapes: List[DetectedShape],

text_regions: List[TextRegion]) -> DiagramAnalysis:

"""

Classify a diagram into a specific type.

Args:

image: Input image

shapes: List of detected shapes

text_regions: List of detected text regions

Returns:

DiagramAnalysis object with classification results

"""

# First, try heuristic classification based on shapes

heuristic_result = self._heuristic_classification(shapes, text_regions)

# If heuristic classification is confident, use it

if heuristic_result.confidence > 0.8:

self.logger.info(f"Heuristic classification: {heuristic_result.diagram_type.value} "

f"(confidence: {heuristic_result.confidence:.2f})")

return heuristic_result

# Otherwise, use vision-language model for deeper analysis

vlm_result = self._vlm_classification(image, shapes, text_regions)

# Combine results, preferring VLM if available

if vlm_result and vlm_result.confidence > heuristic_result.confidence:

self.logger.info(f"VLM classification: {vlm_result.diagram_type.value} "

f"(confidence: {vlm_result.confidence:.2f})")

return vlm_result

else:

return heuristic_result

def _heuristic_classification(self, shapes: List[DetectedShape],

text_regions: List[TextRegion]) -> DiagramAnalysis:

"""

Classify diagram using heuristic rules based on shape patterns.

Args:

shapes: List of detected shapes

text_regions: List of detected text regions

Returns:

DiagramAnalysis object

"""

# Count shape types

shape_counts = {}

for shape in shapes:

shape_type = shape.shape_type

shape_counts[shape_type] = shape_counts.get(shape_type, 0) + 1

total_shapes = len(shapes)

if total_shapes == 0:

return DiagramAnalysis(

diagram_type=DiagramType.GENERIC,

confidence=0.5,

characteristics={},

suggested_rules={}

)

# Calculate shape proportions

arrow_ratio = shape_counts.get(ShapeType.ARROW, 0) / total_shapes

rectangle_ratio = shape_counts.get(ShapeType.RECTANGLE, 0) / total_shapes

diamond_ratio = shape_counts.get(ShapeType.DIAMOND, 0) / total_shapes

circle_ratio = shape_counts.get(ShapeType.CIRCLE, 0) / total_shapes

# Flowchart detection: high arrow ratio, mix of rectangles and diamonds

if arrow_ratio > 0.3 and rectangle_ratio > 0.2 and diamond_ratio > 0.1:

return DiagramAnalysis(

diagram_type=DiagramType.FLOWCHART,

confidence=0.85,

characteristics={

'arrow_ratio': arrow_ratio,

'rectangle_ratio': rectangle_ratio,

'diamond_ratio': diamond_ratio

suggested_rules={

'align_shapes': True,

'standardize_sizes': True,

'arrow_routing': 'orthogonal'

}

)

# UML Class Diagram: rectangles with internal divisions

if rectangle_ratio > 0.6 and arrow_ratio < 0.4:

# Check for compartmentalized rectangles

compartmentalized = self._check_compartmentalized_rectangles(

shapes, text_regions

)

if compartmentalized:

return DiagramAnalysis(

diagram_type=DiagramType.UML_CLASS,

confidence=0.9,

characteristics={

'rectangle_ratio': rectangle_ratio,

'compartmentalized': True

suggested_rules={

'align_shapes': True,

'standardize_widths': True,

'compartment_lines': True

}

)

# Circuit diagram: high symbol density, specific shapes

if circle_ratio > 0.3 or self._has_circuit_symbols(shapes):

return DiagramAnalysis(

diagram_type=DiagramType.CIRCUIT,

confidence=0.8,

characteristics={

'circle_ratio': circle_ratio,

'has_circuit_symbols': True

suggested_rules={

'use_standard_symbols': True,

'wire_routing': 'orthogonal'

}

)

# Network diagram: circles/ellipses connected by lines

if (circle_ratio > 0.4 or shape_counts.get(ShapeType.ELLIPSE, 0) / total_shapes > 0.4) \

and arrow_ratio > 0.2:

return DiagramAnalysis(

diagram_type=DiagramType.NETWORK,

confidence=0.75,

characteristics={

'circle_ratio': circle_ratio,

'connection_ratio': arrow_ratio

suggested_rules={

'node_layout': 'force_directed',

'edge_routing': 'curved'

}

)

# Mind map: tree structure radiating from center

if self._is_tree_structure(shapes):

return DiagramAnalysis(

diagram_type=DiagramType.MINDMAP,

confidence=0.7,

characteristics={

'tree_structure': True

suggested_rules={

'layout': 'radial',

'edge_style': 'curved'

}

)

# Default to generic

return DiagramAnalysis(

diagram_type=DiagramType.GENERIC,

confidence=0.6,

characteristics=shape_counts,

suggested_rules={'preserve_layout': True}

)

def _check_compartmentalized_rectangles(self, shapes: List[DetectedShape],

text_regions: List[TextRegion]) -> bool:

"""

Check if rectangles contain internal horizontal divisions.

Args:

shapes: List of detected shapes

text_regions: List of detected text regions

Returns:

True if compartmentalized rectangles detected

"""

rectangles = [s for s in shapes if s.shape_type == ShapeType.RECTANGLE]

for rect in rectangles:

x, y, w, h = rect.bbox

# Look for horizontal lines inside rectangle

for shape in shapes:

if shape.shape_type == ShapeType.LINE:

line_x1 = shape.parameters['start_x']

line_x2 = shape.parameters['end_x']

line_y1 = shape.parameters['start_y']

line_y2 = shape.parameters['end_y']

# Check if line is horizontal and inside rectangle

if abs(line_y1 - line_y2) < 5: # Approximately horizontal

if x < line_x1 < x + w and x < line_x2 < x + w:

if y < line_y1 < y + h:

return True

return False

def _has_circuit_symbols(self, shapes: List[DetectedShape]) -> bool:

"""

Check for presence of circuit-specific symbols.

Args:

shapes: List of detected shapes

Returns:

True if circuit symbols detected

"""

# Look for patterns characteristic of circuit symbols

# Resistors: rectangles with specific aspect ratio

# Capacitors: parallel lines

# Batteries: combination of long and short parallel lines

for shape in shapes:

if shape.shape_type == ShapeType.RECTANGLE:

aspect_ratio = shape.parameters.get('aspect_ratio', 0)

# Resistor typically has aspect ratio around 3:1

if 2.5 < aspect_ratio < 4.0:

return True

return False

def _is_tree_structure(self, shapes: List[DetectedShape]) -> bool:

"""

Determine if shapes form a tree structure.

Args:

shapes: List of detected shapes

Returns:

True if tree structure detected

"""

# Build adjacency graph from arrows/lines

nodes = [s for s in shapes if s.shape_type in

[ShapeType.RECTANGLE, ShapeType.CIRCLE, ShapeType.ELLIPSE]]

edges = [s for s in shapes if s.shape_type in

[ShapeType.ARROW, ShapeType.LINE]]

if len(nodes) < 3 or len(edges) < 2:

return False

# Simple heuristic: check if number of edges is approximately nodes - 1

# (characteristic of trees)

return abs(len(edges) - (len(nodes) - 1)) <= 2

def _vlm_classification(self, image: Image.Image,

shapes: List[DetectedShape],

text_regions: List[TextRegion]) -> Optional[DiagramAnalysis]:

"""

Use vision-language model for diagram classification.

Args:

image: Input image

shapes: List of detected shapes

text_regions: List of detected text regions

Returns:

DiagramAnalysis object or None if classification fails

"""

try:

# Construct detailed prompt for diagram analysis

prompt = f"""Analyze this hand-drawn diagram and classify it into one of these types:

- flowchart: Process flow diagram with decision points

- uml_class: UML class diagram showing classes and relationships

- uml_sequence: UML sequence diagram showing interactions over time

- uml_activity: UML activity diagram showing workflow

- circuit: Electronic circuit schematic

- network: Network topology diagram

- er_diagram: Entity-relationship database diagram

- mindmap: Mind map or concept map

- organizational: Organizational chart or hierarchy

- generic: Generic drawing not matching above types

The diagram contains {len(shapes)} shapes and {len(text_regions)} text regions.

Respond with ONLY the diagram type (one of the values above) followed by a confidence score (0-1) and a brief explanation.

Format: TYPE|CONFIDENCE|EXPLANATION"""

# Get model response

response = self.vision_backend.analyze_image(image, prompt, max_tokens=200)

# Parse response

parts = response.strip().split('|')

if len(parts) >= 3:

type_str = parts[0].strip().lower()

confidence = float(parts[1].strip())

explanation = parts[2].strip()

# Map string to enum

try:

diagram_type = DiagramType(type_str)

except ValueError:

diagram_type = DiagramType.GENERIC

return DiagramAnalysis(

diagram_type=diagram_type,

confidence=confidence,

characteristics={'vlm_explanation': explanation},

suggested_rules=self._get_default_rules(diagram_type)

)

except Exception as e:

self.logger.error(f"VLM classification error: {e}")

return None

def _get_default_rules(self, diagram_type: DiagramType) -> Dict[str, Any]:

"""

Get default beautification rules for a diagram type.

Args:

diagram_type: Classified diagram type

Returns:

Dictionary of suggested rules

"""

rules_map = {

DiagramType.FLOWCHART: {

'align_shapes': True,

'standardize_sizes': True,

'arrow_routing': 'orthogonal'

DiagramType.UML_CLASS: {

'align_shapes': True,

'standardize_widths': True,

'compartment_lines': True

DiagramType.CIRCUIT: {

'use_standard_symbols': True,

'wire_routing': 'orthogonal'

DiagramType.NETWORK: {

'node_layout': 'force_directed',

'edge_routing': 'curved'

DiagramType.MINDMAP: {

'layout': 'radial',

'edge_style': 'curved'

}

return rules_map.get(diagram_type, {'preserve_layout': True})

The diagram classifier combines heuristic pattern matching with deep learning analysis. Heuristic rules provide fast classification for common patterns, while the vision-language model handles ambiguous cases and complex diagrams. The confidence scoring allows the system to fall back to generic handling when classification is uncertain. Suggested rules guide the beautification process to respect diagram-specific conventions.

SEMANTIC ANALYSIS AND GRAPH CONSTRUCTION

Beyond recognizing individual shapes and text, the system must understand the semantic relationships between elements. Arrows connect nodes, text labels describe components, and spatial proximity indicates grouping. The semantic analyzer constructs a graph representation capturing these relationships.

The graph construction process associates text with nearby shapes, identifies connections between elements, and builds a structured representation of diagram semantics. This graph serves as the foundation for digital reconstruction.

from typing import List, Dict, Any, Set, Tuple, Optional

from dataclasses import dataclass, field

import networkx as nx

import numpy as np

@dataclass

class DiagramNode:

"""Represents a node in the diagram graph."""

node_id: str

shape: DetectedShape

labels: List[str] = field(default_factory=list)

properties: Dict[str, Any] = field(default_factory=dict)

@dataclass

class DiagramEdge:

"""Represents an edge in the diagram graph."""

source_id: str

target_id: str

edge_shape: Optional[DetectedShape] = None

labels: List[str] = field(default_factory=list)

properties: Dict[str, Any] = field(default_factory=dict)

class SemanticAnalyzer:

"""

Analyzes diagram semantics and constructs graph representation.

Associates text with shapes and identifies relationships.

"""

def __init__(self, proximity_threshold: int = 50):

"""

Initialize semantic analyzer.

Args:

proximity_threshold: Maximum distance for text-shape association

"""

self.proximity_threshold = proximity_threshold

self.logger = logging.getLogger(__name__)

def analyze(self, shapes: List[DetectedShape],

text_regions: List[TextRegion],

diagram_type: DiagramType) -> nx.DiGraph:

"""

Analyze diagram semantics and build graph representation.

Args:

shapes: List of detected shapes

text_regions: List of detected text regions

diagram_type: Classified diagram type

Returns:

NetworkX directed graph representing diagram structure

"""

# Create graph

graph = nx.DiGraph()

# Separate shapes into nodes and edges

node_shapes, edge_shapes = self._separate_shapes(shapes)

# Create nodes

nodes = self._create_nodes(node_shapes, text_regions)

# Add nodes to graph

for node in nodes:

graph.add_node(

node.node_id,

shape=node.shape,

labels=node.labels,

properties=node.properties

)

# Create edges based on connections

edges = self._create_edges(edge_shapes, nodes, diagram_type)

# Add edges to graph

for edge in edges:

graph.add_edge(

edge.source_id,

edge.target_id,

shape=edge.edge_shape,

labels=edge.labels,

properties=edge.properties

)

self.logger.info(f"Constructed graph with {graph.number_of_nodes()} nodes "

f"and {graph.number_of_edges()} edges")

return graph

def _separate_shapes(self, shapes: List[DetectedShape]) -> Tuple[List[DetectedShape],

List[DetectedShape]]:

"""

Separate shapes into nodes (rectangles, circles, etc.) and edges (arrows, lines).

Args:

shapes: List of all detected shapes

Returns:

Tuple of (node_shapes, edge_shapes)

"""

node_types = {ShapeType.RECTANGLE, ShapeType.CIRCLE, ShapeType.ELLIPSE,

ShapeType.TRIANGLE, ShapeType.DIAMOND, ShapeType.POLYGON}

edge_types = {ShapeType.ARROW, ShapeType.LINE}

node_shapes = [s for s in shapes if s.shape_type in node_types]

edge_shapes = [s for s in shapes if s.shape_type in edge_types]

return node_shapes, edge_shapes

def _create_nodes(self, node_shapes: List[DetectedShape],

text_regions: List[TextRegion]) -> List[DiagramNode]:

"""

Create diagram nodes from shapes and associate text labels.

Args:

node_shapes: List of shapes representing nodes

text_regions: List of detected text regions

Returns:

List of DiagramNode objects

"""

nodes = []

used_text_indices = set()

for i, shape in enumerate(node_shapes):

node_id = f"node_{i}"

# Find text regions near this shape

associated_text = []

for j, text_region in enumerate(text_regions):

if j in used_text_indices:

continue

if self._is_text_inside_shape(text_region, shape):

associated_text.append(text_region.text)

used_text_indices.add(j)

elif self._is_text_near_shape(text_region, shape):

associated_text.append(text_region.text)

used_text_indices.add(j)

# Create node

node = DiagramNode(

node_id=node_id,

shape=shape,

labels=associated_text,

properties={

'shape_type': shape.shape_type.value,

'bbox': shape.bbox

}

)

nodes.append(node)

return nodes

def _is_text_inside_shape(self, text_region: TextRegion,

shape: DetectedShape) -> bool:

"""

Check if text region is inside a shape.

Args:

text_region: Text region to check

shape: Shape to check against

Returns:

True if text is inside shape

"""

text_x, text_y, text_w, text_h = text_region.bbox

text_center_x = text_x + text_w // 2

text_center_y = text_y + text_h // 2

shape_x, shape_y, shape_w, shape_h = shape.bbox

# Check if text center is inside shape bounding box

if (shape_x <= text_center_x <= shape_x + shape_w and

shape_y <= text_center_y <= shape_y + shape_h):

# For more accurate check, use contour

point = np.array([[text_center_x, text_center_y]], dtype=np.float32)

result = cv2.pointPolygonTest(shape.contour,

(float(text_center_x), float(text_center_y)),

False)

return result >= 0

return False

def _is_text_near_shape(self, text_region: TextRegion,

shape: DetectedShape) -> bool:

"""

Check if text region is near a shape.

Args:

text_region: Text region to check

shape: Shape to check against

Returns:

True if text is near shape

"""

text_x, text_y, text_w, text_h = text_region.bbox

text_center_x = text_x + text_w // 2

text_center_y = text_y + text_h // 2

shape_x, shape_y, shape_w, shape_h = shape.bbox

shape_center_x = shape_x + shape_w // 2

shape_center_y = shape_y + shape_h // 2

# Calculate distance between centers

distance = np.sqrt((text_center_x - shape_center_x) ** 2 +

(text_center_y - shape_center_y) ** 2)

return distance < self.proximity_threshold

def _create_edges(self, edge_shapes: List[DetectedShape],

nodes: List[DiagramNode],

diagram_type: DiagramType) -> List[DiagramEdge]:

"""

Create edges by connecting nodes based on arrows and lines.

Args:

edge_shapes: List of shapes representing edges (arrows, lines)

nodes: List of diagram nodes

diagram_type: Type of diagram

Returns:

List of DiagramEdge objects

"""

edges = []

for edge_shape in edge_shapes:

# Get edge endpoints

if edge_shape.shape_type in {ShapeType.ARROW, ShapeType.LINE}:

start_x = edge_shape.parameters['start_x']

start_y = edge_shape.parameters['start_y']

end_x = edge_shape.parameters['end_x']

end_y = edge_shape.parameters['end_y']

# Find nodes at endpoints

source_node = self._find_node_at_point(nodes, start_x, start_y)

target_node = self._find_node_at_point(nodes, end_x, end_y)

if source_node and target_node:

edge = DiagramEdge(

source_id=source_node.node_id,

target_id=target_node.node_id,

edge_shape=edge_shape,

properties={

'edge_type': edge_shape.shape_type.value

}

)

edges.append(edge)

return edges

def _find_node_at_point(self, nodes: List[DiagramNode],

x: int, y: int) -> Optional[DiagramNode]:

"""

Find node at or near a specific point.

Args:

nodes: List of nodes to search

x, y: Point coordinates

Returns:

DiagramNode if found, None otherwise

"""

search_radius = 30

for node in nodes:

shape_x, shape_y, shape_w, shape_h = node.shape.bbox

# Check if point is inside or near shape bounding box

if (shape_x - search_radius <= x <= shape_x + shape_w + search_radius and

shape_y - search_radius <= y <= shape_y + shape_h + search_radius):

# More precise check using contour

result = cv2.pointPolygonTest(

node.shape.contour,

(float(x), float(y)),

True # Measure distance

)

if result >= -search_radius:

return node

return None

The semantic analyzer transforms low-level shape detections into a high-level graph structure. Text association uses both containment and proximity heuristics to handle labels both inside and adjacent to shapes. Edge creation identifies connections by finding nodes at arrow endpoints. The resulting graph captures the essential semantic structure of the diagram, independent of visual presentation details.

DIGITAL RECONSTRUCTION AND BEAUTIFICATION

The final stage transforms the semantic graph into a clean digital representation. This involves layout optimization, shape regularization, text formatting, and rendering to vector graphics. The beautification process applies diagram-type-specific rules while preserving semantic equivalence to the original drawing.

The reconstruction engine uses the classified diagram type to select appropriate layout algorithms and styling rules. It generates SVG output that can be further edited or converted to other formats.

import svgwrite

from svgwrite import cm, mm

import networkx as nx

from typing import Dict, Any, Tuple, List

import math

class DiagramReconstructor:

"""

Reconstructs clean digital diagrams from semantic graph representation.

Applies beautification and layout optimization.

"""

def __init__(self, canvas_size: Tuple[int, int] = (800, 600)):

"""

Initialize diagram reconstructor.

Args:

canvas_size: Output canvas dimensions (width, height)

"""

self.canvas_size = canvas_size

self.logger = logging.getLogger(__name__)

def reconstruct(self, graph: nx.DiGraph,

diagram_analysis: DiagramAnalysis,

output_path: str) -> str:

"""

Reconstruct diagram as clean digital SVG.

Args:

graph: Semantic graph representation

diagram_analysis: Diagram type and characteristics

output_path: Path for output SVG file

Returns:

Path to generated SVG file

"""

# Create SVG drawing

dwg = svgwrite.Drawing(output_path, size=self.canvas_size)

# Add definitions for reusable elements

self._add_definitions(dwg, diagram_analysis.diagram_type)

# Optimize layout based on diagram type

layout = self._compute_layout(graph, diagram_analysis)

# Apply beautification rules

beautified_graph = self._beautify_graph(graph, diagram_analysis)

# Render nodes

self._render_nodes(dwg, beautified_graph, layout, diagram_analysis)

# Render edges

self._render_edges(dwg, beautified_graph, layout, diagram_analysis)

# Save SVG

dwg.save()

self.logger.info(f"Diagram reconstructed and saved to {output_path}")

return output_path

def _add_definitions(self, dwg: svgwrite.Drawing, diagram_type: DiagramType):

"""

Add SVG definitions for markers and reusable elements.

Args:

dwg: SVG drawing object

diagram_type: Type of diagram

"""

# Add arrowhead marker

marker = dwg.marker(

id='arrowhead',

insert=(10, 5),

size=(10, 10),

orient='auto'

)

marker.add(dwg.path(d='M 0 0 L 10 5 L 0 10 z', fill='black'))

dwg.defs.add(marker)

# Add diamond marker for UML aggregation

if diagram_type == DiagramType.UML_CLASS:

diamond = dwg.marker(

id='diamond',

insert=(10, 5),

size=(10, 10),

orient='auto'

)

diamond.add(dwg.path(d='M 0 5 L 5 0 L 10 5 L 5 10 z',

fill='white', stroke='black'))

dwg.defs.add(diamond)

def _compute_layout(self, graph: nx.DiGraph,

diagram_analysis: DiagramAnalysis) -> Dict[str, Tuple[float, float]]:

"""

Compute optimal layout for diagram nodes.

Args:

graph: Semantic graph

diagram_analysis: Diagram analysis results

Returns:

Dictionary mapping node IDs to (x, y) positions

"""

suggested_rules = diagram_analysis.suggested_rules

# Choose layout algorithm based on diagram type

if suggested_rules.get('layout') == 'radial':

return self._radial_layout(graph)

elif suggested_rules.get('node_layout') == 'force_directed':

return self._force_directed_layout(graph)

elif diagram_analysis.diagram_type == DiagramType.FLOWCHART:

return self._hierarchical_layout(graph)

elif diagram_analysis.diagram_type == DiagramType.UML_CLASS:

return self._grid_layout(graph)

else:

return self._preserve_layout(graph)

def _hierarchical_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:

"""

Compute hierarchical layout suitable for flowcharts.

Args:

graph: Semantic graph

Returns:

Node positions dictionary

"""

try:

# Use Sugiyama layout (hierarchical)

pos = nx.spring_layout(graph, k=2, iterations=50)

# Scale to canvas size with margins

margin = 50

width = self.canvas_size[0] - 2 * margin

height = self.canvas_size[1] - 2 * margin

scaled_pos = {}

for node_id, (x, y) in pos.items():

scaled_x = margin + (x + 1) * width / 2

scaled_y = margin + (y + 1) * height / 2

scaled_pos[node_id] = (scaled_x, scaled_y)

return scaled_pos

except Exception as e:

self.logger.error(f"Layout computation error: {e}")

return self._preserve_layout(graph)

def _grid_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:

"""

Compute grid layout suitable for UML class diagrams.

Args:

graph: Semantic graph

Returns:

Node positions dictionary

"""

nodes = list(graph.nodes())

n_nodes = len(nodes)

# Calculate grid dimensions

cols = math.ceil(math.sqrt(n_nodes))

rows = math.ceil(n_nodes / cols)

# Calculate spacing

margin = 50

h_spacing = (self.canvas_size[0] - 2 * margin) / cols

v_spacing = (self.canvas_size[1] - 2 * margin) / rows

# Position nodes

pos = {}

for i, node_id in enumerate(nodes):

col = i % cols

row = i // cols

x = margin + col * h_spacing + h_spacing / 2

y = margin + row * v_spacing + v_spacing / 2

pos[node_id] = (x, y)

return pos

def _radial_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:

"""

Compute radial layout suitable for mind maps.

Args:

graph: Semantic graph

Returns:

Node positions dictionary

"""

nodes = list(graph.nodes())

n_nodes = len(nodes)

if n_nodes == 0:

return {}

# Place first node at center

center_x = self.canvas_size[0] / 2

center_y = self.canvas_size[1] / 2

pos = {nodes[0]: (center_x, center_y)}

# Place remaining nodes in circle

radius = min(self.canvas_size) / 3

for i, node_id in enumerate(nodes[1:], 1):

angle = 2 * math.pi * i / (n_nodes - 1)

x = center_x + radius * math.cos(angle)

y = center_y + radius * math.sin(angle)

pos[node_id] = (x, y)

return pos

def _force_directed_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:

"""

Compute force-directed layout.

Args:

graph: Semantic graph

Returns:

Node positions dictionary

"""

pos = nx.spring_layout(graph, k=1.5, iterations=100)

# Scale to canvas

margin = 50

width = self.canvas_size[0] - 2 * margin

height = self.canvas_size[1] - 2 * margin

scaled_pos = {}

for node_id, (x, y) in pos.items():

scaled_x = margin + (x + 1) * width / 2

scaled_y = margin + (y + 1) * height / 2

scaled_pos[node_id] = (scaled_x, scaled_y)

return scaled_pos

def _preserve_layout(self, graph: nx.DiGraph) -> Dict[str, Tuple[float, float]]:

"""

Preserve original layout from hand-drawn diagram.

Args:

graph: Semantic graph

Returns:

Node positions dictionary

"""

pos = {}

for node_id in graph.nodes():

shape = graph.nodes[node_id]['shape']

x, y, w, h = shape.bbox

center_x = x + w / 2

center_y = y + h / 2

pos[node_id] = (center_x, center_y)

return pos

def _beautify_graph(self, graph: nx.DiGraph,

diagram_analysis: DiagramAnalysis) -> nx.DiGraph:

"""

Apply beautification rules to graph.

Args:

graph: Original semantic graph

diagram_analysis: Diagram analysis with suggested rules

Returns:

Beautified graph

"""

beautified = graph.copy()

rules = diagram_analysis.suggested_rules

# Standardize node sizes if requested

if rules.get('standardize_sizes'):

self._standardize_node_sizes(beautified)

# Align shapes if requested

if rules.get('align_shapes'):

self._align_shapes(beautified)

return beautified

def _standardize_node_sizes(self, graph: nx.DiGraph):

"""

Standardize sizes of similar node types.

Args:

graph: Graph to modify in place

"""

# Group nodes by shape type

shape_groups = {}

for node_id in graph.nodes():

shape_type = graph.nodes[node_id]['properties']['shape_type']

if shape_type not in shape_groups:

shape_groups[shape_type] = []

shape_groups[shape_type].append(node_id)

# Standardize each group

for shape_type, node_ids in shape_groups.items():

if len(node_ids) < 2:

continue

# Calculate average size

total_w = 0

total_h = 0

for node_id in node_ids:

_, _, w, h = graph.nodes[node_id]['shape'].bbox

total_w += w

total_h += h

avg_w = total_w / len(node_ids)

avg_h = total_h / len(node_ids)

# Apply average size

for node_id in node_ids:

shape = graph.nodes[node_id]['shape']

x, y, _, _ = shape.bbox

shape.bbox = (x, y, int(avg_w), int(avg_h))

def _align_shapes(self, graph: nx.DiGraph):

"""

Align shapes to grid.

Args:

graph: Graph to modify in place

"""

grid_size = 20

for node_id in graph.nodes():

shape = graph.nodes[node_id]['shape']

x, y, w, h = shape.bbox

# Snap to grid

aligned_x = round(x / grid_size) * grid_size

aligned_y = round(y / grid_size) * grid_size

shape.bbox = (aligned_x, aligned_y, w, h)

def _render_nodes(self, dwg: svgwrite.Drawing, graph: nx.DiGraph,

diagram_analysis: DiagramAnalysis):

"""

Render diagram nodes to SVG.

Args:

dwg: SVG drawing object

graph: Semantic graph

diagram_analysis: Diagram analysis

"""

for node_id in graph.nodes():

node_data = graph.nodes[node_id]

shape = node_data['shape']

labels = node_data['labels']

x, y = layout[node_id]

# Render based on shape type

if shape.shape_type == ShapeType.RECTANGLE:

self._render_rectangle(dwg, x, y, shape, labels, diagram_analysis)

elif shape.shape_type == ShapeType.CIRCLE:

self._render_circle(dwg, x, y, shape, labels)

elif shape.shape_type == ShapeType.DIAMOND:

self._render_diamond(dwg, x, y, shape, labels)

elif shape.shape_type == ShapeType.ELLIPSE:

self._render_ellipse(dwg, x, y, shape, labels)

def _render_rectangle(self, dwg: svgwrite.Drawing, x: float, y: float,

shape: DetectedShape, labels: List[str],

diagram_analysis: DiagramAnalysis):

"""Render a rectangle node."""

width = shape.parameters.get('width', 100)

height = shape.parameters.get('height', 60)

rect = dwg.rect(

insert=(x - width/2, y - height/2),

size=(width, height),

fill='white',

stroke='black',

stroke_width=2

)

dwg.add(rect)

# Add text labels

if labels:

text_y = y - height/2 + 20

for label in labels:

text = dwg.text(

label,

insert=(x, text_y),

text_anchor='middle',

font_size='14px',

font_family='Arial'

)

dwg.add(text)

text_y += 20

def _render_circle(self, dwg: svgwrite.Drawing, x: float, y: float,

shape: DetectedShape, labels: List[str]):

"""Render a circle node."""

radius = shape.parameters.get('radius', 30)

circle = dwg.circle(

center=(x, y),

r=radius,

fill='white',

stroke='black',

stroke_width=2

)

dwg.add(circle)

# Add text label

if labels:

text = dwg.text(

labels[0],

insert=(x, y + 5),

text_anchor='middle',

font_size='14px',

font_family='Arial'

)

dwg.add(text)

def _render_diamond(self, dwg: svgwrite.Drawing, x: float, y: float,

shape: DetectedShape, labels: List[str]):

"""Render a diamond node."""

width = shape.parameters.get('width', 80)

height = shape.parameters.get('height', 60)

points = [

(x, y - height/2),

(x + width/2, y),

(x, y + height/2),

(x - width/2, y)

]

polygon = dwg.polygon(

points=points,

fill='white',

stroke='black',

stroke_width=2

)

dwg.add(polygon)

# Add text label

if labels:

text = dwg.text(

labels[0],

insert=(x, y + 5),

text_anchor='middle',

font_size='14px',

font_family='Arial'

)

dwg.add(text)

def _render_ellipse(self, dwg: svgwrite.Drawing, x: float, y: float,

shape: DetectedShape, labels: List[str]):

"""Render an ellipse node."""

major_axis = shape.parameters.get('major_axis', 60)

minor_axis = shape.parameters.get('minor_axis', 40)

ellipse = dwg.ellipse(

center=(x, y),

r=(major_axis/2, minor_axis/2),

fill='white',

stroke='black',

stroke_width=2

)

dwg.add(ellipse)

# Add text label

if labels:

text = dwg.text(

labels[0],

insert=(x, y + 5),

text_anchor='middle',

font_size='14px',

font_family='Arial'

)

dwg.add(text)

def _render_edges(self, dwg: svgwrite.Drawing, graph: nx.DiGraph,

diagram_analysis: DiagramAnalysis):

"""

Render diagram edges to SVG.

Args:

dwg: SVG drawing object

graph: Semantic graph

diagram_analysis: Diagram analysis

"""

for source, target in graph.edges():

x1, y1 = layout[source]

x2, y2 = layout[target]

edge_data = graph.edges[source, target]

edge_shape = edge_data.get('shape')

# Determine if arrow or plain line

has_arrow = edge_shape and edge_shape.shape_type == ShapeType.ARROW

# Create path

if diagram_analysis.suggested_rules.get('arrow_routing') == 'orthogonal':

path_d = self._create_orthogonal_path(x1, y1, x2, y2)

else:

path_d = f'M {x1} {y1} L {x2} {y2}'

# Add path with optional arrowhead

path_attrs = {

'stroke': 'black',

'stroke_width': 2,

'fill': 'none'

}

if has_arrow:

path_attrs['marker_end'] = 'url(#arrowhead)'

path = dwg.path(d=path_d, **path_attrs)

dwg.add(path)

def _create_orthogonal_path(self, x1: float, y1: float,

x2: float, y2: float) -> str:

"""

Create orthogonal (right-angle) path between two points.

Args:

x1, y1: Start point

x2, y2: End point

Returns:

SVG path data string

"""

# Simple orthogonal routing with midpoint

mid_x = (x1 + x2) / 2

return f'M {x1} {y1} L {mid_x} {y1} L {mid_x} {y2} L {x2} {y2}'

The reconstruction engine produces publication-quality diagrams from rough hand-drawn input. Layout algorithms respect diagram-type conventions while optimizing visual clarity. Shape rendering uses clean geometric primitives with consistent styling. Text placement ensures readability without overlapping elements. The SVG output format enables further editing and scaling without quality loss.

ALTERNATIVE TOOLS AND APPROACHES

Vectorization Tools: Services like Kittl, Linearity Curve, Vectorizer.AI, Adobe Illustrator (Auto Trace), and Vectr scan and convert raster sketches to scalable, clean vector graphics, including shape and text recognition.

AI Sketch Enhancement: Platforms such as Fotor, Deep-image.ai, and Playform.io offer AI-powered beautification, enhancing rough sketches and even generating digital art or realistic renders in a chosen style.

Diagram and Structure Recognition: Systems like Sketch2Scheme focus on schematic diagrams and flowcharts, recognizing drawn elements and transforming them into clean, digital diagrams.

Customization and Style: Some tools provide style transfer and editing capabilities, letting you apply specific artistic or technical styles to your improved sketch.

Supported Elements:

Shapes and Lines: AI can detect, clarify, and vectorize freehand shapes and lines for a clean, professional result.
Text: For hand-written labels or annotations, certain AI tools can enhance, straighten, or even convert the handwriting to editable or beautified digital text.
Artistic and Technical Styles: Choose between technical (blueprint, diagram, architectural) and artistic (cartoon, painting, minimalist) output according to your final use case.

This technology lets anyone—from engineers and teachers to artists and hobbyists—transform a scanned sketch into a refined digital asset quickly, without needing advanced design or image editing skills.

COMPLETE PRODUCTION SYSTEM

The following complete implementation integrates all components into a production-ready system. This code provides a full pipeline from image input to SVG output, supporting all discussed features including multi-GPU acceleration, multiple LLM backends, and comprehensive diagram type recognition.

#!/usr/bin/env python3

"""

Hand-Drawn Diagram Digitization System

A complete production-ready system for converting hand-drawn diagrams

to clean digital representations with semantic preservation.

Supports:

- Multiple GPU architectures (NVIDIA CUDA, AMD ROCm, Apple MPS, Intel XPU)

- Local and remote vision-language models

- Multiple diagram types (flowcharts, UML, circuits, etc.)

- Comprehensive shape and text detection

- Semantic analysis and graph construction

- Beautification and layout optimization

"""

import logging

import argparse

import sys

from pathlib import Path

from typing import Optional

from PIL import Image

# Configure logging

logging.basicConfig(

level=logging.INFO,

format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'

)

class DiagramDigitizationPipeline:

"""

Complete pipeline for digitizing hand-drawn diagrams.

Integrates all components from preprocessing to SVG generation.

"""

def __init__(self, backend_type: str = 'local_llama',

backend_config: Optional[dict] = None):

"""

Initialize the digitization pipeline.

Args:

backend_type: Type of vision-language backend to use

backend_config: Configuration dictionary for the backend

"""

self.logger = logging.getLogger(__name__)

# Initialize device manager

self.logger.info("Initializing device manager...")

self.device_manager = DeviceManager()

self.device_manager.optimize_for_inference()

# Initialize vision-language backend

self.logger.info(f"Initializing {backend_type} backend...")

backend_config = backend_config or {}

self.vision_backend = VisionBackendFactory.create_backend(

backend_type,

device_manager=self.device_manager,

**backend_config

)

# Initialize pipeline components

self.logger.info("Initializing pipeline components...")

self.preprocessor = ImagePreprocessor(target_size=(1024, 1024))

self.text_extractor = TextExtractor(languages=['en'])

self.shape_detector = ShapeDetector(min_area=100)

self.diagram_classifier = DiagramClassifier(self.vision_backend)

self.semantic_analyzer = SemanticAnalyzer(proximity_threshold=50)

self.reconstructor = DiagramReconstructor(canvas_size=(800, 600))

self.logger.info("Pipeline initialization complete")

def process_image(self, input_path: str, output_path: str) -> str:

"""

Process a hand-drawn diagram image and generate digital SVG.

Args:

input_path: Path to input image file

output_path: Path for output SVG file

Returns:

Path to generated SVG file

"""

self.logger.info(f"Processing image: {input_path}")

# Load image

image = Image.open(input_path)

self.logger.info(f"Loaded image: {image.size}")

# Preprocess image

self.logger.info("Preprocessing image...")

preprocessed = self.preprocessor.preprocess(image)

# Extract text

self.logger.info("Extracting text...")

text_regions = self.text_extractor.extract_text(preprocessed)

# Detect shapes

self.logger.info("Detecting shapes...")

shapes = self.shape_detector.detect_shapes(preprocessed)

# Classify diagram type

self.logger.info("Classifying diagram type...")

diagram_analysis = self.diagram_classifier.classify_diagram(

preprocessed, shapes, text_regions

)

self.logger.info(f"Diagram type: {diagram_analysis.diagram_type.value} "

f"(confidence: {diagram_analysis.confidence:.2f})")

# Perform semantic analysis

self.logger.info("Performing semantic analysis...")

semantic_graph = self.semantic_analyzer.analyze(

shapes, text_regions, diagram_analysis.diagram_type

)

# Reconstruct digital diagram

self.logger.info("Reconstructing digital diagram...")

output_svg = self.reconstructor.reconstruct(

semantic_graph, diagram_analysis, output_path

)

self.logger.info(f"Processing complete. Output saved to: {output_svg}")

return output_svg

def get_system_info(self) -> dict:

"""

Get information about the system configuration.

Returns:

Dictionary with system information

"""

return {

'device': self.device_manager.get_device_info(),

'backend': self.vision_backend.get_backend_info(),

'components': {

'preprocessor': 'ImagePreprocessor',

'text_extractor': 'TextExtractor',

'shape_detector': 'ShapeDetector',

'diagram_classifier': 'DiagramClassifier',

'semantic_analyzer': 'SemanticAnalyzer',

'reconstructor': 'DiagramReconstructor'

}

def main():

"""Main entry point for the diagram digitization system."""

parser = argparse.ArgumentParser(

description='Digitize hand-drawn diagrams to clean SVG representations'

)

parser.add_argument(

'input',

type=str,

help='Path to input image file'

)

parser.add_argument(

'output',

type=str,

help='Path for output SVG file'

)

parser.add_argument(

'--backend',

type=str,

choices=['local_llama', 'openai'],

default='local_llama',

help='Vision-language model backend to use'

)

parser.add_argument(

'--model',

type=str,

default='llava-hf/llava-1.5-7b-hf',

help='Model name for local backend'

)

parser.add_argument(

'--api-key',

type=str,

help='API key for remote backends (e.g., OpenAI)'

)

parser.add_argument(

'--info',

action='store_true',

help='Display system information and exit'

)

args = parser.parse_args()

# Build backend configuration

backend_config = {}

if args.backend == 'local_llama':

backend_config['model_name'] = args.model

elif args.backend == 'openai':

if not args.api_key:

print("Error: --api-key required for OpenAI backend")

sys.exit(1)

backend_config['api_key'] = args.api_key

# Initialize pipeline

try:

pipeline = DiagramDigitizationPipeline(

backend_type=args.backend,

backend_config=backend_config

)

# Display system info if requested

if args.info:

import json

info = pipeline.get_system_info()

print(json.dumps(info, indent=2))

sys.exit(0)

# Process image

output_path = pipeline.process_image(args.input, args.output)

print(f"\nSuccess! Digital diagram saved to: {output_path}")

except Exception as e:

logging.error(f"Pipeline error: {e}", exc_info=True)

sys.exit(1)

if __name__ == '__main__':

main()

This production system provides a complete command-line interface for diagram digitization. Users can specify input images, output paths, and backend configurations through command-line arguments. The system handles errors gracefully and provides detailed logging throughout the processing pipeline. The modular architecture allows easy extension with additional diagram types, backends, or processing stages.

USAGE EXAMPLES AND DEPLOYMENT

To use the system, first ensure all dependencies are installed. The required packages include PyTorch with appropriate GPU support, OpenCV, Tesseract OCR, EasyOCR, NetworkX, and svgwrite. For local vision-language models, install the Transformers library and download the desired model.

For NVIDIA CUDA systems, install PyTorch with CUDA support. For AMD ROCm, use the ROCm-enabled PyTorch build. Apple Silicon users should install PyTorch with MPS support. Intel GPU users need the Intel Extension for PyTorch.

Basic usage with a local model processes an image as follows:

python diagram_digitizer.py input_sketch.jpg output_diagram.svg --backend local_llama --model llava-hf/llava-1.5-7b-hf

For OpenAI GPT-4 Vision, provide an API key:

python diagram_digitizer.py input_sketch.jpg output_diagram.svg --backend openai --api-key YOUR_API_KEY

The system automatically detects available GPU hardware and configures acceleration accordingly. It processes the input image through all pipeline stages and generates a clean SVG representation preserving the semantic content of the original drawing.

PERFORMANCE CONSIDERATIONS AND OPTIMIZATION

The system performance depends on several factors including input image resolution, diagram complexity, and available hardware. GPU acceleration significantly improves processing speed for vision models and image processing operations. Local models provide faster inference than remote APIs but require more memory and computational resources.

Image preprocessing benefits from OpenCV's optimized implementations. The adaptive thresholding and morphological operations execute efficiently on both CPU and GPU. Text extraction with EasyOCR leverages GPU acceleration when available, substantially reducing OCR time for complex diagrams.

Shape detection using contour analysis scales linearly with image resolution. The Hough transform for line detection has higher computational complexity but remains practical for typical diagram sizes. Caching intermediate results between pipeline stages avoids redundant computation.

The semantic analysis and graph construction stages have minimal computational cost compared to vision processing. Layout optimization algorithms vary in complexity, with force-directed layouts requiring more iterations than grid or hierarchical layouts. The SVG rendering stage is lightweight and completes quickly regardless of diagram complexity.

For production deployment, consider implementing batch processing for multiple diagrams, caching preprocessed images, and using model quantization to reduce memory requirements. Distributed processing across multiple GPUs can accelerate throughput for high-volume scenarios.

CONCLUSION

This comprehensive system demonstrates how modern computer vision and artificial intelligence techniques enable robust digitization of hand-drawn diagrams. The modular architecture supports diverse hardware configurations and backend choices while maintaining clean separation of concerns. The integration of multiple OCR engines, shape detection algorithms, and vision-language models provides robust handling of varied input quality and diagram types.

The semantic preservation approach ensures that digitized diagrams maintain the meaning and intent of original drawings while applying beautification and standardization. Support for multiple diagram types with type-specific layout and styling rules produces outputs that respect domain conventions. The extensible design allows easy addition of new diagram types, backends, or processing techniques.

The production-ready implementation provides a solid foundation for building diagram digitization applications. The comprehensive error handling, logging, and configuration options enable deployment in diverse environments. The open-source technology stack ensures accessibility and allows customization for specific requirements.

Future enhancements could include interactive editing of semantic graphs before rendering, support for additional diagram types such as Gantt charts or BPMN diagrams, integration with diagramming tools through API interfaces, and machine learning-based improvement of shape classification accuracy through user feedback.