Generating Pixels One by One

Your First Autoregressive Image Generation Model

We’ll build a basic autoregressive model using a simple MLP to generate images ofhandwritten digits, focusing on understanding the core concept of predicting the next pixel based on its predecessors. It’s a hands-on exploration of fundamental generative AI, showing some of the core concepts using a pretty simple model. The model we will train will be far away from the state-of-the-art, but it will be a good starting point to understand the core concepts of autoregressive models.

Welcome, I am glad you are here!

I’m Tuna. My world is pretty much all about image and video generation. It is what I focus on in my PhD and during my internships at places like Adobe (working on Firefly!) and Amazon AGI. For a while, I have been working with diffusion-based models, and I know that they are incredibly powerful.

But the landscape of generative modeling is always growing, and I want to explore other types of generative models. Right now, I am diving into autoregressive models. I always find the best way to learn a topic is by trying to teach it to others. So, this blog post series is an attempt to teach myself the basics of autoregressive models, hoping you can learn something from it, too. I’ll start with the basics and try to understand how these models work piece by piece.

What Makes a Model “Autoregressive”?

Alright, “Autoregressive”. Let’s break it down with some mathematical intuition.

You have already seen “auto-regressive” models in action even if you didn’t call them that. At its heart, it basically means predicting the next outcome based on all the things that came before it.

Think about how you type on your phone. When you write “the weather is …”, the keyboard will suggest completions based on the words you entered such as “sunny”, “rainy”, “perfect for research” (maybe not that last one). That’s an auto-regressive model in action for language.

Mathematically, for a sequence (x_1, x_2, …, x_T), an autoregressive model learns:

\[P(x_1, x_2, ..., x_T) = P(x_1) \cdot P(x_2|x_1) \cdot P(x_3|x_1, x_2) \cdot ... \cdot P(x_T|x_1, ..., x_{T-1})\]

This is just the chain rule of probability. Each new element depends on all the previous ones.

For images, we can think of each pixel as an element in our sequence. So instead of predicting the next word, we’re predicting the next pixel value based on all the pixels we’ve seen so far. Cool, right?

Let’s Start with Data: Meet Our MNIST Friends

Before we dive into the math, let’s get acquainted with our data. We’ll use MNIST digits - they’re simple, well-understood, and good for learning the basics.

Global Configuration and Hyperparameters

Define all key hyperparameters for the models, data, and training in one place. This allows for easy modification and consistency across the different model versions.

The “Pixel as Token” Approach: Quantizing Intensities

We treat image generation more like language modeling. Each pixel’s intensity will be quantized into one of a discrete number of bins. Each bin then gets an integer label, turning our pixels into “tokens” or “words” from a fixed vocabulary.

1. Quantization

• We take the continuous grayscale pixel values (0.0 to 1.0).
• We divide this range into K discrete bins (e.g., K=16 or K=256).
• Each pixel’s original intensity is mapped to the integer label of the bin it falls into.
• Example: If K=4, intensities 0.0-0.25 map to token 0, 0.25-0.5 to token 1, etc.

2. Prediction as Classification

• The model’s task is now to predict the integer label (token) of the next pixel, given the tokens of previous pixels.
• This is a K-class classification problem for each pixel.

Benefits:

• Allows using powerful classification machinery (like Cross-Entropy loss).
• Can be combined with techniques like embedding layers for tokens, similar to NLP.

Trade-offs:

• Information Loss: Quantization inherently loses some precision from the original grayscale values. More bins reduce this loss but increase model complexity.
• Vocabulary Size: The number of bins (K) becomes our vocabulary size.

Building Our First Token Predictor: A Basic MLP

Now that we have quantized our pixels, or tokens, (from 0 to NUM_QUANTIZATION_BINS-1), let’s build our first autoregressive model. The goal is simple: given a token sequence of length CONTEXT_LENGTH, predict the next token.

For this first version, we will try a fairly direct approach:

1. Representing Tokens: Our tokens are integers from 0 to NUM_QUANTIZATION_BINS-1. To feed them into the neural network, a common approach is to use one-hot encoding. For example, if a token is represented as 3, the one-hot encoding is a list of length (NUM_QUANTIZATION_BINS + 1) with a 1 at index 3 and 0s elsewhere. The +1 is for the start token and it will have its own encoding, too. This way each token has its own unique representation.
2. Model Architecture: We’ll use a simple MLP to predict the next token, the input will be the one-hot encoding of the previous tokens, and the MLP will learn to map the one-hot encodings to the next token.
3. Output: The MLP will output scores (logits) for each of the NUM_QUANTIZATION_BINS possible pixel tokens, telling us how likely each token is to be the next one.

Let’s implement this and see what kind of outputs we can get. Our focus here is to create a basic auto-regressive loop and see how our model can generate an image pixel-by-pixel using direct representation of the pixels.

The Heart of Autoregression: Context Windows

Now for the fun part! How do we actually implement “predicting the next pixel based on previous pixels”?

The key insight is context windows. Instead of using all previous pixels, we use a sliding window of the last k pixels as our context.

Mathematically, instead of modeling

\[P(x_i | x_1, x_2, ..., x_{i-1})\]

we approximate it as:

\[P(x_i | x_{i-1}, x_{i-2}, ..., x_{i-k})\]

This is called a k-th order Markov assumption - we assume the immediate past is most informative for predicting the future.

What about the first few pixels that don’t have enough history? That’s where our start tokens come in - they’re like saying “this is the beginning of an image” to our model.

Building Our Neural Network: Predicting Pixel Tokens

Our neural network will now predict a distribution over the possible pixel tokens (integer labels) for the next pixel.

Training Our OneHotPixelPredictor (Model V1)

With our model_one_hot_pixel_predictor (Model V1) defined and the training data (train_contexts, train_targets) prepared, we’re ready to train.

• Loss Function: Since our model outputs logits for NUM_QUANTIZATION_BINS possible pixel token classes, and our target is a single integer class label, we’ll use CrossEntropyLoss. This loss function is ideal for multi-class classification tasks as it combines a LogSoftmax layer and Negative Log-Likelihood Loss in one step.
• Optimizer: We’ll use AdamW, a common and effective optimizer known for good performance and weight decay.

Let’s kick off the training for Model V1!

Generating Images with Model V1

After training our OneHotPixelPredictor, let’s see what kind of images it can generate. The autoregressive generation process involves:

1. Initializing a context window, filled with our special START_TOKEN_VALUE_INT.
2. Feeding this current context to the trained model to obtain logits (scores) for each possible next pixel token.
3. Converting these logits into a probability distribution using the softmax function. We can also apply a “temperature” parameter here to control the randomness of sampling:
   • Temperature = 1.0: Standard sampling according to the model’s learned probabilities.
   • Temperature < 1.0 (e.g., 0.7): Makes the model’s choices “sharper” or more deterministic, favoring higher-probability tokens (more greedy).
   • Temperature > 1.0 (e.g., 1.2): Makes the sampling more random, increasing diversity but potentially at the cost of coherence.
4. Sampling the next pixel token from this probability distribution.
5. Appending this newly sampled token to our sequence of generated pixels.
6. Updating the context window by shifting it (removing the oldest token) and adding the newly generated token.
7. Repeating steps 2-6 until all N_PIXELS for a full image have been generated.
8. Finally, de-quantizing the sequence of generated integer tokens back into visualizable grayscale values.

Let’s see how our Model V1, which uses one-hot encoded tokens and no explicit positional information, performs at this task.

Model V1: Initial Results and Observations

Alright, let’s take a look at the first batch of images generated by our OneHotPixelPredictor (Model V1)!

Looking at these generated samples, here are a few key observations:

• Recognizable Shapes? Unfortunately, not really. The images don’t form coherent, recognizable MNIST digits. Instead, we see patterns that are quite abstract and noisy.
• Dominant Patterns: Horizontal Streaks: A very prominent feature is the presence of horizontal streaks or bands of white (or lighter) pixels against a predominantly dark background. It seems the model has learned some very local, short-range correlations, possibly related to the raster scan order (e.g., “if the last few pixels were white, the next one is also likely white for a short stretch”).
• Repetitive Textures: The generations have a somewhat repetitive, textural quality due to these streaks. There isn’t much variation in the types of structures being formed beyond these horizontal elements.
• Lack of Global Coherence: Crucially, there’s an almost complete absence of global structure. The horizontal streaks appear somewhat randomly placed across the image canvas and don’t coordinate to form larger, meaningful shapes like the curves and loops of digits. The model doesn’t seem to have a plan for the overall image.
• Impact of Quantization Visible: The blocky nature of the 16 quantization bins is apparent, which is expected.

Why are we seeing these kinds of results? Two key aspects of our very basic Model V1 are likely major contributors:

No Positional Information: Our current MLP processes a flat window of context tokens. It has no explicit information about where in the 2D image it is currently trying to predict a pixel. Is it at the top-left, the center, the bottom-right? Without this spatial awareness, it’s incredibly hard for the model to learn to, for example, start a stroke at a particular location, curve it appropriately, and end it correctly to form part of a digit. The horizontal streaks might be a consequence of learning simple local rules that get repeated because the model doesn’t know when or where to change its behavior based on image location.

These initial results, while not producing digits, are incredibly valuable. They highlight a key limitation and clearly motivate our next step. What if we could give our model a sense of ‘location’ for each pixel it predicts? This is precisely what we’ll explore by introducing positional encodings in Model V2. We’ll keep the one-hot encoding for pixel tokens for now and see just how much impact adding spatial information can have.

Model V2: Giving Our Predictor a Sense of Location with Positional Encodings

Our Model V1 struggled to create coherent images, and we hypothesized that a major reason was its lack of spatial awareness – it didn’t know where in the image it was predicting a pixel.

For Model V2, we’ll directly address this by introducing positional encodings. The idea is to provide the model with explicit information about the (row, column) coordinates of the pixel it’s currently trying to predict.

How will we do this?

1. Learnable Positional Embeddings: We’ll create two separate embedding layers, one for row positions (0 to IMG_SIZE-1) and one for column positions (0 to IMG_SIZE-1). Each position (e.g., row 5, column 10) will be mapped to a learnable vector (its embedding).
2. Concatenation: For each prediction step, we’ll determine the row and column of the target pixel. We’ll fetch their respective embeddings. These two positional embedding vectors will then be concatenated with the (still one-hot encoded) context window of previous pixel tokens.
3. MLP Input: This richer, combined representation (one-hot context + row embedding + column embedding) will be fed into our MLP.

Our pixel tokens themselves will still be one-hot encoded for now. The main change in Model V2 is adding this crucial positional signal. The hypothesis is that by knowing where it’s operating, the model can learn location-dependent rules and generate more structured images. Let’s see!

Preparing Training Data for Model V2 (with Positional Information)

Our Model V2, OneHotPixelPredictorWithPosition, now expects not only the context window of previous tokens but also the absolute position of the pixel it is trying to predict.

This means we need to use our data preparation function to output three things for each training sample:

1. The context window of integer tokens (x_context_tokens).
2. The target integer token (target_token).
3. The absolute “flat” position of the target token in the image (e.g., an integer from 0 to N_PIXELS-1).

Our model will then internally convert this flat position into row and column indices to fetch the appropriate positional embeddings.

Let’s ensure our data generation function provides this. We’ll use the version that generates context, target, and position.

Training Model V2 (with Positional Encodings)

Now that we have our model_onehot_with_pos and the corresponding training data which includes positional information for each target pixel, we can proceed with training.

The training setup will be very similar to Model V1:

• Loss Function: CrossEntropyLoss, as we’re still predicting one of K pixel tokens.
• Optimizer: AdamW.
• Process: We’ll iterate for a set number of epochs, shuffling our data and processing it in mini-batches. The key difference is that during the forward pass, we will now also provide the pixel_positions_flat to the model.

Let’s see if providing this spatial awareness helps the model learn to generate more structured images, even with the one-hot token representation.

Generating Images with Model V2 (with Positional Encodings)

Now that Model V2, OneHotPixelPredictorWithPosition, has been trained with positional information, we can generate images. The autoregressive process remains largely the same as with Model V1, but with one crucial addition:

1. Initialize a context window with START_TOKEN_VALUE_INT.
2. For each pixel we want to generate (from pixel 0 to N_PIXELS-1): a. Determine the current pixel’s absolute flat position. b. Feed the current context window and this current pixel position to the model. c. Obtain logits, apply temperature, convert to probabilities via softmax. d. Sample the next pixel token. e. Append the sampled token to our generated sequence. f. Update the context window.
3. Repeat until the image is complete.
4. De-quantize the generated tokens.

Let’s see if the added spatial awareness from positional encodings helps Model V2 produce more structured or recognizable images compared to Model V1.

Model V2: Results with Positional Encodings

Now, let’s examine the images generated by Model V2 (OneHotPixelPredictorWithPosition), which incorporated positional encodings while still using one-hot representations for the pixel tokens.

Comparing these to the outputs from Model V1 (which had no positional information), we can make some significant observations:

• Improved Structure - Emergence of Verticality: This is the most striking difference! The dominant horizontal streaks from Model V1 are largely gone. Instead, we see a clear tendency towards vertical structures or alignments. The model now seems to understand that pixels “above” and “below” each other are related in a way that pixels far apart horizontally might not be (or at least, it’s learned to generate patterns consistent with the vertical strokes common in digits).
• Hints of Digit-like Forms (but still abstract): While still not producing clear, recognizable digits, some of these vertical structures have a more “digit-like” feel than the random noise of Model V1. You can almost squint and see hints of ‘1’s, or parts of ‘7’s, ‘4’s, or other vertical segments. There’s a sense that the model is trying to place “ink” in a more constrained, column-oriented fashion.
• Centering Tendency: Many of the generated patterns appear somewhat centered within the 28x28 canvas, which is typical for MNIST digits. This suggests the positional encodings are helping the model learn where to place “active” pixels.
• Still “Noisy” and “Fragmented”: Despite the move towards verticality, the generations are still quite noisy and fragmented. The “strokes” are often broken, and there’s a lack of smooth curves or connected components that would form complete digits.
• Coarseness from One-Hot Encoding Persists: The blocky appearance due to the 16 quantization bins and the one-hot encoding of tokens is still evident. The model isn’t (and can’t easily with this representation) learning smooth transitions between pixel intensities.

The Impact of Positional Information:

Adding positional encoding has clearly made a substantial difference. By knowing the (row, column) coordinate of the pixel it’s predicting, Model V2 has been able to:

1. Move away from the undirected, horizontal streak patterns of Model V1.
2. Learn to generate patterns with a strong vertical bias, which is a key characteristic of many handwritten digits.
3. Roughly position these patterns within the typical area of a digit on the canvas.

This demonstrates the critical importance of spatial awareness for image generation tasks. Even a simple MLP, when given location information, can start to learn rudimentary structural properties.

However, we’re not quite there yet. The generated images still lack clarity and fine detail. A remaining bottleneck is likely our one-hot encoding of pixel tokens. This representation treats each of our 16 pixel intensity levels as a completely distinct, unrelated category. The model has no inherent understanding that, for example, token 3 (a dark gray) is semantically “closer” to token 4 (a slightly different dark gray) than it is to token 15 (a very light gray). It must learn these relationships purely from co-occurrence, which can be inefficient and limit its ability to model subtle intensity variations.

This naturally leads us to our next and final refinement for this MLP-based autoregressive model: What if we replace the one-hot encoding of pixel tokens with learned, dense vector representations (embeddings)? This naturally leads us to our next refinement: What if we replace the one-hot encoding of pixel tokens with learned, dense vector representations (embeddings)? Furthermore, instead of just generating any digit, can we guide the model to generate a specific digit? This is what we’ll explore in Model V3. We will combine the benefits of learned token embeddings and positional encodings, and crucially, we will introduce conditional generation by providing the model with the desired class label (i.e., the digit we want it to generate)

Model V3: Enhancing Token Understanding with Learned Embeddings

Our journey so far has shown progress. Model V1, with no spatial awareness, produced noisy streaks. Model V2, by incorporating positional encodings, started to generate more structured, vertically-oriented patterns, hinting at digit-like forms. However, the results are still far from clear MNIST digits, and a key limitation we identified was the one-hot encoding of pixel tokens. This representation forces the model to learn the relationships between different pixel intensities (e.g., that “dark gray” is similar to “slightly darker gray”) from scratch, without any inherent notion of their similarity.

For Model V3, we’ll address this by introducing learned token embeddings for our pixel intensity values. This is a technique widely used in Natural Language Processing (NLP) with great success.

How it Works:

1. Embedding Layer (nn.Embedding): Instead of one-hot encoding, each integer pixel token (from 0 to NUM_QUANTIZATION_BINS-1) and our special START_TOKEN_VALUE_INT will be mapped to a dense, lower-dimensional vector. These vectors are learnable parameters of the model.
2. Learning Relationships: During training, the model will adjust these embedding vectors. If certain pixel intensities often appear in similar contexts or lead to similar predictions, their embedding vectors will tend to become similar. This allows the model to capture the semantic “closeness” of different pixel values.
3. Combined with Positional Encoding: We will retain the absolute positional encodings from Model V2. The input to our MLP will now be a concatenation of:
   1. The learned embedding vectors for each token in the context window.
   2. The learned positional embedding for the target pixel’s row.
   3. 1. The learned positional embedding for the target pixel’s column.

By providing the model with both a richer, learned representation for pixel intensities and spatial awareness, we hope to see a significant improvement in the quality and coherence of the generated images. This Model V3 represents our most sophisticated MLP-based autoregressive generator in this tutorial.

Preparing Training Data for Model V3 (with Class Labels)

Our Model V3, CategoricalPixelPredictor, is designed for conditional generation.

This means it now expects not only the context window and the target pixel’s position but also the class label of the image from which the sample is drawn. We’ll adapt our data preparation to include these class labels:

1. Context Windows (train_contexts_v3): Sequences of integer tokens.
2. Target Tokens (train_targets_v3): The integer token of the actual next pixel.
3. Target Pixel Positions (train_positions_v3): The absolute flat position of the target pixel.
4. Class Labels (train_labels_v3): The integer label (0-9 for MNIST) of the image corresponding to the current training sample. The model will use an embedding of this label as part of its input.”

Training Model V3 (with Token Embeddings and Positional Encodings)

We’re now ready to train our latest MLP-based model for this tutorial, model_cat_predictor. This model combines:

• Learned dense embeddings for the pixel intensity tokens (and the start token) in the context window.
• Learned positional embeddings for the row and column of the target pixel.

The training setup (loss function, optimizer) will be the same as for Model V2. We expect that the richer input representation (learned token embeddings for context, positional embeddings for location, and class embeddings for the desired digit) will allow the model to learn more effectively and generate images conditioned on the provided class label.

Generating Images with Model V3 (Token Embeddings + Positional Encodings)

With our CategoricalPixelPredictor (Model V3) trained, it’s time for the moment of truth! This model uses learned embeddings for pixel intensity tokens and positional encodings, representing our most sophisticated MLP-based model in this tutorial, now capable of conditional image generation.

The generation process will be similar to Model V2, but with a key difference: we will now provide a target class label (e.g., “generate a 7”) to the model at each step. The model will use an embedding of this class label, along with the context and current pixel position, to predict the next pixel token.

1. Initialize a context window with START_TOKEN_VALUE_INT.
2. Provide the desired class_label for the image to be generated.
3. For each pixel i (from 0 to N_PIXELS-1): a. Feed the current context window, the current pixel’s position (i), and the target class_label to the model. b. Obtain logits, apply temperature, sample the next pixel token. c. Update context and append the token to the generated sequence.
4. Repeat until the image is complete.
5. De-quantize.

We are eager to see if the model can now generate specific digits when prompted, leveraging learned token embeddings, positional encodings, and class conditioning.

— Analyzing Generation for Digit 7 (Model V3) —

Our exploration of autoregressive image generation has taken us on an iterative journey, starting from a simple pixel predictor to a a conditional digit generator. The progression through Models V1, V2 and V3 has been a step by step improvement, highlighing key principles in building these systems.

Model V1, our simplest MLP using only one hot encoded pizel valus without any spatial awarenes, prodiced results,that, as seen in the generated images, were largely abstract and dominated by local, repetitive patterns like horizontal streaks. It was a good example of why we need to encode spatial information in our model.

Model V2 introduced learned positional embeddings. The generated images immediately began to exhibit more global structure, with a clear shift towards recognizable (albeit nor perfectly formed) digits. While not perfect, this confirmed our hypothesis: spatial information is crucial for the model to learn more meaningful image patterns. eyond local correlations.

Model V3 was the culmination of these insights. By replacing one-hot encoding with learnable token embeddings for pixel intensities and, most importantly, introducing class conditioning, we achieved a breakthrough. As the generated images for Model V3 demonstrate, the model can now generate images that are not only more coherent but also represent specific, requested digits.

Limitations and Future Directions:

It’s important to note that our MLP-based models are incredibly simple. The fixed CONTEXT_LENGTH limits the long-range dependencies the model can capture. Quantization bins also lead to a somewhat blocky appearance.

The generations don’t match the quality of state-of-the-art generative models. However, this was never the primary goal; the aim was to understand the foundational concepts of autoregressive generation.

I hope this step-by-step progression has been as enlightening for you to read as it was for me to build and explore. Generating images pixel by pixel, even simple ones, truly demystifies some of the magic behind generative AI. The core idea of predicting one piece at a time, conditioned on what came before, remains a powerful and versatile paradigm in the world of artificial intelligence.

Happy generating!

You can find the code for this tutorial here.

Note: I hope I will continue this series with more complex models and more interesting applications.