The proposed Swin Transformer achieves strong performance on the recognition tasks of image classification, object detection, and semantic segmentation.

Rather than going into the experimental details, this post focuses on the core ideas behind the model. In a nutshell, Swin Transformer is a hierarchical transformer whose representations are computed using shifted windows - that’s where its name (“Shifted Window Transformer”) comes from. The hierarchical design and the window-based self-attention with shifted windows are the two key structures of Swin that we will walk through in more detail later.

Motivation

Transformer is the default backbone of NLP applications: the same basic architecture can power language modeling, translation, question answering, and many other tasks. In computer vision, CNNs play a similar “general-purpose backbone” role. The Swin Transformer paper sets out to close this gap:

In this paper, we seek to expand the applicability of Transformer such that it can serve as a general-purpose backbone for computer vision, as it does for NLP and as CNNs do in vision.

However, it’s challenging to transfer Transformer’s high performance in the language domain to the vision domain. The challenge comes from differences between the two modalities: scale and resolution.

Scale challenge

For typical sentence-level NLP tasks, Transformers operate on a single scale of tokens in a 1D sequence and do not rely on explicit multi-scale representations.¹ In contrast, visual elements can vary substantially in scale: small textures, medium-sized objects, large structures. Classical CNN backbones such as ResNet (He et al., 2016), often combined with a feature pyramid network (FPN; Lin et al., 2017), handle this wide range of scales using hierarchical feature maps that gradually reduce spatial resolution while increasing channel depth. ViT (see more in this post), on the other hand, uses fixed-sized patches throughout the network and keeps a single feature resolution (see Figure 1b). Though this works well for image classification, in its original form (Dosovitskiy et al., 2021) it is not well suited as a general-purpose backbone for dense prediction tasks such as object detection or semantic segmentation.²

Resolution challenge

Images are also much higher resolution than typical text sequences. Standard self-attention has \(O(N^2)\) complexity in the number of tokens \(N\), because it computes interactions between all pairs of tokens. As a result, applying global self-attention at high resolution quickly becomes prohibitively expensive, especially for dense tasks that need fine spatial detail at the pixel level.

How Swin addresses these challenges

Swin Transformer is designed around these challenges. It introduces a hierarchical representation to handle scale variation, and window-based self-attention with shifted windows to make attention feasible at high resolution - reducing the computational complexity from quadratic in the number of tokens to approximately linear - while still allowing information to flow across the image.

Core ideas

In this section, we introduce Swin’s two structural ideas: hierarchical representation and shifted windows. We also explain why Swin achieves an approximate linear computational complexity.

Hierarchical representation

The gray grids in Figure 1a represent patch-level feature maps at different spatial resolutions. Shallow layers contain many patch tokens arranged at a higher resolution. As neighboring tokens are merged, deeper layers contain fewer, coarser patch tokens, where each token represents a larger area of the original image and typically has a larger embedding dimension. This progression from high-resolution, fine-grained features to low-resolution, more abstract features is what we mean by a hierarchical feature representation.

The red boxes represent the local windows within which self-attention is computed. Instead of attending over all patches globally, Swin partitions the feature map into these red-outlined local windows, and applies self-attention only inside each window. This greatly reduces the computational cost compared with global attention.

In summary, Swin combines a CNN-like feature hierarchy with local window-based attention, giving it both multi-scale representations and efficient computation. Here, the gray grids show what the representation consists of (i.e. the spatial grid of patch tokens), while the red boxes show where self-attention is computed (i.e. within local windows over those tokens).

Figure 2 provides an alternative visualization of the architectural differences between ViT and Swin Transformer.

Why window-based attention scales linearly

We mentioned above that Swin’s self-attention is much cheaper than global self-attention. This efficiency comes from limiting self-attention to fixed-sized local windows. To understand why, let us introduce some notation:

• \(N\): the total number of tokens (gray patches in Figure 1a) in the feature map at a given stage
• \(M\): the side length of each window
• \(M^2\): number of tokens in each local attention window (red boxes in Figure 1a). For example, the original Swin paper sets the side length of \(M = 7\) by default, so each \(7 \times 7\) window contains \(M^2 = 49\) tokens.
• Note:
  • Each token corresponds to a progressively larger region of the original image at deeper stages due to patch merging.
  • Swin tiles the feature map with \(N/M^2\) non-overlapping windows, with each token belonging to exactly one window.

Global self-attention (ViT-style)

In global self-attention, every token can attend to every other token. This means:

• We construct an \(N \times N\) attention matrix, where each entry represents the interaction between a pair of tokens.
• The pairwise-attention cost therefore scales like \(O(N^2)\).

As image resolution increases while the patch size remains fixed, the image is divided into more patches and therefore produces more tokens (\(N\)). The quadratic cost of global self-attentil thus quickly becomes prohibitive for high-resolution feature maps.

Window-based self-attention (Swin)

Swin changes the picture by restricting attention to local windows:

• The feature map is partitioned into \(N/M^2\) non-overlapping windows, each with \(M^2\) tokens.
• Within each window, we compute self-attention only among those \(M^2\) tokens. The cost per window is: \(O({M^2}^2) = O(M^4)\)
• Across the whole feature map, the total cost becomes: \((N/M^2) \times O(M^4) = O(NM^2)\)

Here, \(M\) (the side length of each window) is a fixed hyperparameter of the model (e.g., always \(7×7\) in the original Swin paper). It does not grow with image size. As we increase image resolution, \(N\) increases (i.e. more gray patches), but \(M^2\) (i.e. the number of tokens per window) stays the same. Therefore, for a fixed window size, the self-attention cost grows approximately linearly with the number of tokens, instead of quadratically as in global self-attention.

Note: The full complexity expressions in Equations (1) and (2) of Liu et al. (2021) also include the linear projections used to compute queries \((Q)\), keys \((K)\), values \((V)\), and attention output. These projection costs are the same for global and window-based attention. The simplified derivation here focuses on the pairwise-attention term, which is where their computational complexities differ.

Shifted window

Though Swin can reduce cost by restricting attention to local windows, that creates an isolation problem: tokens in different windows cannot interact within the same attention block. Swin addresses this in the following Swin Transformer block at the same stage using the shifted window approach, allowing information to flow across window boundaries.

How shifted window enables cross-window communication

Figure 3 illustrates how shifting the window partition between two consecutive blocks enables cross-window communication. In one block (\(L\)), Swin computes self-attention within non-overlapping windows. In the following block (\(L + 1\)) at the same stage, Swin shifts the window partition by half-window size in all directions, resulting in new windows that cross the previous window boundaries. As a result, some tokens that previously belonged to separate windows now appear in the same shifted window and attend to one another.

The window partition alternates between these two configurations across consecutive blocks. Thus, block (\(L+2\)) returns to the regular window partition used in (\(L\)), while block (\(L+3\)) uses the shifted partition again, as in (\(L+1\)). In other words, the shift does not accumulate across blocks. In practice, Swin implements the shifted-window configuration using a cyclic shift and an attention mask, which we will discuss in the Architecture section.

Sliding windows vs shifted windows

How does the shifted-window mechanism differ from sliding-window attention (Ramachandran et al., 2019)? Although both restrict self-attention to local neighborhoods, they differ in how those neighborhoods are defined and organized. In Swin, the feature map is partitioned into non-overlapping windows, and all query tokens within the same window share the same set of key and value tokens. In contrast, in sliding-window self-attention, the local neighborhood is centered on each query token, so different queries attend to different sets of keys and values. Figures 4-5 illustrate this distinction.

Architecture

In this section, we explore how the core ideas are assembled into the actual Swin backbone.

Overall architecture

Let’s first have a high-level overview of the Swin model architecture. Figure 6 (reproduced from Figure 3 of Liu et al. 2021) covers the main pieces, and we will unpack the key components below.

Stages

The original Swin Transformer models consist of four hierarchical stages, with spatial resolution progressively reduced between stages through patch merging (Figure 6a).

Patch size and embedding dimensions

Liu et al. (2021) use a patch size of \(4 \times 4\), each patch therefore contains \((4\times 4\times 3 = 48)\) raw pixel values. After patch partitioning, the image is represented as a grid of \((H/4 \times W/4)\) patches, each with a \(48-\)dimensional feature vector. A linear embedding layer then projects each patch vector into a learned \(C-\)dimensional feature vector, giving the input to Stage 1 a shape of \(H/4 \times W/4 \times C\). Note that here \(C\) plays conceptually the same role as the embedding dimension \(D\) in ViT.

Patch merging between stages

Patch merging reduces spatial resolution while increasing the embedding dimension. At each patch-merging layer, Swin groups neighboring tokens in fixed \(2\times 2\) blocks. This \(2\times 2\) grouping is a design choice that gives a natural \(2\times\) spatial downsampling in both height and width, similar to stride-2 downsampling in a CNN.

If the incoming tokens at a given stage have embedding dimension \(d\), where \(d\) is used here as a generic notation for that stage’s current embedding dimension, the four \(d\)-dimensional token vectors are concatenated into a \(4d\)-dimensional vector and then linearly projected to \(2d\). Thus, each patch-merging step halves the spatial resolution in both directions while doubling the embedding dimension.

Therefore, starting with embedding dimension \(C\) in Stage 1, repeated patch-merging produces the following stage-wise progression:

Swin Transformer blocks

Each stage contains multiple Swin Transformer blocks; the number of blocks varies across Swin variants. Within each stage, successive blocks always alternate between W-MSA (window-based multi-head self-attention) and SW-MSA (shifted-window multi-head self-attention), as illustrated by the two consecutive blocks in Figure 6b.

Cyclic shifting

We have seen conceptually how shifted windows enable cross-window communication (Figure 3). But directly implementing the shifted partition introduces an efficiency problem at the feature-map boundaries: it creates more windows, including smaller partial windows along the edges. For example, in the example shown in Figure 7, the number of windows increases from 4 to 9. Swin addresses this using cyclic shifting and attention masking.

Cyclic shifting works by shifting the feature map toward the top-left by approximately half a window size \(M/2\), where \(M\) denotes the side length of each \(M\times M\) attention window. Tokens that cross the top boundary wrap to the bottom, those that cross the left boundary warp to the right, and those that cross both wrap to the bottom-right (Figure 7). This rearrangement allows the shifted configuration to be processed using the same number of regular \(M\times M\) computational windows as regular window partitioning. However, it also places sub-windows that were not adjacent in the original feature map into the same computational window. Masked MSA therefore ensures tokens only attend to other tokens within the same sub-window. For example, in Figure 8, all tokens within region 5 can attend to one another, whereas regions 8 and 2 share a computational window only because of cyclic wraparound, so attention between regions 8 and 2 is masked out. After Masked MSA, a reverse cyclic shift moves the entire feature map back to its original spatial arrangement. See also my sketch in Figure 8 for a visual explanation.

References

• Arkin, E., Yadikar, N., Xu, X., Aysa, A., & Ubul, K. (2023). A survey: Object detection methods from CNN to transformer. Multimedia Tools and Applications, 82, 21353–21383.

• Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An image is worth 16×16 words: Transformers for image recognition at scale. International Conference on Learning Representations (ICLR).

• Hassani, A., Walton, S., Li, J., Li, S., & Shi, H. (2023). Neighborhood Attention Transformer — CVPR 2023 presentation slides. IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR 2023).

• He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 770–778.

• Lin, T.-Y., Dollár, P., Girshick, R., He, K., Hariharan, B., & Belongie, S. (2017). Feature pyramid networks for object detection. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2117–2125.

• Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., & Guo, B. (2021). Swin Transformer: Hierarchical Vision Transformer using shifted windows. Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 10012–10022.

• Ramachandran, P., Parmar, N., Vaswani, A., Bello, I., Levskaya, A., & Shlens, J. (2019). Stand-alone self-attention in vision models. Advances in Neural Information Processing Systems (NeurIPS), 32.

• Wang, H., Zhang, F., & Yi, R. (2025). Multi-scale deformable transformer with iterative query refinement for hot-rolled steel surface defect detection. Sensors, 25(22), 6890.

Footnote

Before ViT, convolutional architectures remained dominant in computer vision. Classic Convolutional Neural Networks (CNNs) such as VGG and especially ResNet, with its residual skip connections, were state of the art. In the original ViT paper (Dosovitskiy et al., 2021, An image is worth 16 x 16 words: transformers for image recognition at scale), the authors applied transformers to image classification tasks and found that, when trained on large datasets, ViT can match or even surpass the performance of top convolutional networks.

Core idea

From words to patches

ViT processes sequences of image patches just as transformers process sequences of word tokens in natural language processing (NLP). In NLP, a transformer takes a sequence of word embeddings as input and learns contextual relationships among them through self-attention. Similarly, ViT splits an image into fixed-size patches, flattens each patch, linearly embeds the patches into vectors, and feeds the resulting sequence of patch embeddings into a standard transformer encoder.

Besides this words vs patches analogy, ViT also mirrors how transformers are trained in NLP. Dosovitskiy et al. (2021) note that:

The dominant approach [in NLP] is to pre-train on a large text corpus and then fine-tune on a smaller task-specific dataset.

Large Transformer-based models are often pre-trained on large corpora and then fine-tuned for the task at hand: BERT uses a denoising self-supervised pre-training task, while the GPT line of work uses language modeling as its pre-training task.

ViT follows a similar pattern: it excels when first trained on large image datasets and then fine-tuned for smaller, domain-specific tasks. We will get to this point later in the post when discussing why ViT benefits so much from large-scale pretraining. For a deeper look at pretraining in large language models, see my previous post.

CNNs vs. ViT

In machine learning, inductive bias refers to the built-in assumptions a model makes about data before seeing any examples. Dosovitskiy et al. (2021) noted that ViT has much less image-specific inductive bias than CNNs, quote:

… In CNNs, locality, two-dimensional neighborhood structure, and translation equivariance are baked into each layer throughout the whole model. In ViT, only MLP layers are local and translationally equivariant, while the self-attention layers are global.

The table below summarizes how CNNs and ViT differ in their built-in assumptions about images.

Why ViT excels with large datasets

Because ViT does not come with strong image-specific inductive biases, it treats all patches equally and relies on self-attention to decide which regions of an image are important. This explains why ViT performs worse on small datasets but excels on large ones.

When the training dataset is small, CNNs often outperform ViTs because their built-in assumptions help them learn meaningful patterns even with limited data. In contrast, ViT has to learn locality, edges, textures, and other visual structures from scratch, making it more data-hungry and less efficient at small scales. With large datasets, however, ViT can learn spatial and structural patterns directly from data rather than relying on hard-coded assumptions. Its global self-attention becomes a major advantage, enabling the model to capture long-range dependencies across the entire image and ultimately surpass CNNs.

Architecture

In this section, we walk through the ViT model architecture using Figure 1 and Equations (1)-(4) from Dosovitskiy et al. (2021). Figure 1 (below) is an annotated version of the original figure from the paper, showing how an image is divided into patches, embedded, and processed through the transformer encoder. Google Research blog also provides an animated visualization that illustrates the ViT workflow.

Here is how ViT is defined mathematically:

Equation (1): embedding

Equation (1) corresponds to the embedding stage in Figure 1. The input image is represented as \(x \in \mathbb{R}^{H \times W \times C}\), where \(H\), \(W\), and \(C\) denote height, width, and number of channels. The image is divided into \(N\) patches, \(x_p \in \mathbb{R}^{N \times (P^2 \cdot C)}\), where \((P, P)\) represents the resolution of each image patch. Each patch vector is multiplied by a learnable projection matrix \(E\) to get a \(D\)-dimensional embedding. A special classification token \(x{\text{class}}\) is prepended - its embedding will later represent the whole image. A positional embedding \(E_{\text{pos}}\) is also added to retain spatial information. These form the input sequence \(\mathbf{z}_0\) that will be feeded to the transformer.

Understanding constant \(D\)

Transformers are designed to process sequences where every token has the same dimensionality across all layers. This uniformity makes stacking layers simple: each layer takes a sequence of vectors of shape \([N \times D]\) and outputs another sequence \([N \times D]\). Here, \(D\) indicates the embedding dimension (sometimes called the hidden size). Every patch embedding and every intermediate vector inside the transformer has size \(D\).

\(D\) is a hyperparameter. With enough data, a larger \(D\) usually means higher model capacity and better performance — but it also requires more compute and memory. In a nutshell, \(D\) is a design choice controlling how much information each token can represent and how big the transformer is.

Understanding the trainable linear layer \(\mathbf{E}\)

A raw image patch of size \(P \times P \times C\) is flattened into a vector, but the transformer expects tokens of dimension \(D\). So a trainable linear layer (denoted as \(\mathbf{E}\)) maps flattened patches into the \(D\)-dimensional embedding space. After this step, all tokens (patches + [CLS]) live in the same latent space.

Let’s walk through the logic using the ViT-Base configuration. Suppose the image size is \(224 \times 224 \times 3\), patch size is \(16 \times 16\), and embedding dimension \(D = 768\).

1. Split into patches.
   Each patch = \(16 \times 16 \times 3 = 768\) raw pixel values.
   Total number of patches = \((224 / 16)^2 = 196\).

2. Flatten each patch.
   Flatten \(16 \times 16 \times 3\) into a vector of length 768.
   Before projection, each patch is a 768-dimensional vector of raw pixel intensities.

3. Project to the embedding dimension \(D\).
   Note: although the patch length (768) and \(D\) (768) are identical in this example (ViT-Base), they are not the same thing.
   The flattened patch values are raw pixel values, while \(D\) represents the dimension of learned embeddings that the transformer requires in a consistent latent space.
   Even when \(P \times P \times C = D\), we still need a linear projection — it’s not about matching dimensions, it’s about learning a mapping from raw pixels → latent embedding space.

   ViT applies a trainable linear layer \(\mathbf{E}\):

   \[\mathbf{e} = \mathbf{W} \mathbf{x} + \mathbf{b}\]

   where:

   • \(\mathbf{x}\) = flattened patch (size \(P^2 C\), here 768)
   • \(\mathbf{W}\) = learnable weight matrix (size \((P^2 C) \times D\))
   • \(\mathbf{b}\) = bias term (size \(D\))
   • \(\mathbf{e}\) = patch embedding (size \(D\))

4. Build the sequence.
   Each of the 196 patches is now projected to a 768-dimensional embedding vector.
   Add one learnable [CLS] token (also 768-dimensional). The sequence length becomes 197, each of size 768 — so the input to the transformer is a \([197 \times 768]\) matrix.

5. Transformer layers.
   Each transformer layer takes a \([197 \times 768]\) input and outputs another \([197 \times 768]\) matrix.

Equation (2): multi-head self-attention (MSA)

Equation (2) represents the self-attention sub-layer in each transformer block. Inside each transformer block, the input passes through layer normalization (LN) and then multi-head self-attention (MSA): \(\mathrm{MSA}\!\left(\mathrm{LN}(\mathbf{z}_{\ell-1})\right)\). Each token (patch embedding) attends to all others globally, capturing relationships across the entire image. The residual connection \(\mathbf{z}_{\ell-1}\) (skip arrow in Figure 1) helps stabilize gradient flow during training.

Equation (3): feed-forward (MLP)

Equation (3) describes the feed-forward network (or MLP) that follows self-attention. After the attention step, the output is normalized again and passed through a two-layer MLP (with a GELU activation in between): \(\mathrm{MLP}(x) = W_2 \, \mathrm{GELU}(W_1 x + b_1) + b_2\). Another residual connection adds the input back to the output.

Together, Equations (2) and (3) make up a transformer block - stacked \(L\) times as shown in Figure 1.

Equation (4): classification head

Equation (4) corresponds to the output head of the model. After L transformer blocks, the output (i.e., the class token embedding \(\mathbf{z}^{0}_{L}\) that represents the entire image) is extracted. \(\mathbf{z}^{0}_{L}\) is layer-normalized and used for classification through a linear classifier. The linear classifier in ViT can be viewed as a standard logistic regression layer that takes the final class token embedding as input and outputs logits for each class.

Implementation

In this section, we demonstrate a jupyter notebook that maps the core ViT equations to a minimal PyTorch implementation. It follows the paper’s equations (1) to (4) and mirrors the structure of open-source ViT implementations (Google Research, timm, Hugging Face).

References

• Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., & Houlsby, N. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. Paper

• Google Research ViT: Repo

• timm: PyTorch Image Models (Ross Wightman) Vision Transformer source: Code

• Hugging Face Transformers — Vision Transformer (ViT): Documentation, code

Model optimizations for deployment

Distillation

Distillation uses a larger model (the teacher model) to train a smaller model (the student model). Then use the smaller model for inference to lower the storage and compute budget. The student model learns to statistically mimic the behavior of the teacher model, either just in the final prediction layer, or in the model’s hidden layers as well.

Distillation workflow:

• Start with a fine-tuned LLM as the teacher model, and create a smaller LLM for student model.
• Freeze the teacher model’s weights and use it to generate completions for the training data. Meanwhile, also generate completions for the training data using the student model.
• The knowledge distillation between teacher and student model is achieved by minimizing a loss function called the distillation loss. To calculate the loss, distillation compares the probability distribution over tokens produced by the student model’s softmax layer to that of the teacher model’s softmax layer.
• Adding a temperature parameter to the softmax function. Why?
  • The teacher model is already fine-tuned on the training data, so the probability distribution likely closely matches the ground truth data and won’t have much variation in tokens. That’s why Distillation applies a little trick adding a temperature parameter to the softmax function.
  • With a temperature parameter greater than one, the probability distribution becomes broader and less strongly peaked. This softer distribution provides you with a set of tokens that are similar to the ground truth tokens.
• In parallel, you train the student model to generate the correct predictions based on your ground truth training data. Here, you don’t vary the temperature setting, and instead, use the standard softmax function. Distillation refers to the student model outputs as the hard predictions and hard labels. The loss between the hard predictions and hard labels is the student loss.
• The combined distillation and student losses are used to updated the weights of the student model via back propagation.

The key benefit of distillation methods is that the smaller student model can be used for inference in deployment instead of the teacher model. In practice, distillation is not as effective for generative decoder models. It’s typically more effective for encoder only models such as BERT that have a lot of representation redundancy.

Quantization

We introduced the concept of quantization in the LLM pre-training note. After a model is trained, we can perform post training quantization (PTQ) to optimize for deployment. PTQ transforms a model’s weights to a lower precision representation, such as 16-bit floating point for 8-bit integer.

• To reduce the model size and memory footprint, as well as the compute resources needed for model serving, quantization can be applied to just the model weights or to both weights and activation layers. In general, quantization approaches that include activations can have a higher impact on model performance.
• Quantization also requires an extra calibration step to statistically capture the dynamic range of the original parameter values.
• Trade-off: sometimes quantization results in a small percentage reduction in model evaluation metrics. That reduction can often be worth the cost saving and performance gains.

Pruning

Pruning is to reduce model size for inference by eliminating weights that are not contributing much to overall model performance. These are the weights with values very close to or equal to zero.

• Some pruning methods require full retraining of the model, while others fall into the category of parameter efficient fine tuning, such as LoRA.
• There are also methods that focus on post-training pruning.

In theory, pruning reduces model size and improves performance. In practice, however, there may not be much impact on the size and performance if only a small percentage of the model weights are close to zero.

Retrieval augmented generation (RAG)

RAG helps LLMs update their understanding of the world and overcomes the knowledge cutoff issue. While we could retrain the model on new data, this would quickly become expensive, and require repeated retraining to regularly update the model with new knowledge. A more flexible and less expensive way to overcome knowledge cutoff is to give model access to additional external data at inference time.

RAG isn’t a specific set of technologies, but rather a framework for providing LLMs access to data they did not see during training. A number of different implementations exist, and the one you choose will depend on the details of your task, and the format of data to work with.

RAG can be used to integrate multiple types of external information sources. In the following, we briefly introduce using vector store for RAG. Vector store contains vector representation of text, and enables a fast and efficient kind of relevant search based on similarity.

Data preparation for vector store

Two considerations:

1. Data must fit inside context window

• The external data sources can be chopped into many chunks, each of which will fit in the context window. Packages like Langchain can handle this.

1. Data must be in format that allows its relevance to be assessed at inference time: embedding vectors

• RAG methods take the small chunks of external data and process them through embedding models to create embedding vectors for each. These embedding vectors can be stored in structures called vector stores, which allow for fast searching of datasets and efficient identification of semantically related text.

Vector databases are a particular implementation of a vector store, where each vector is also identified by a key. This allows for the text generated by RAG to include a citation for the document from which it was received.

Chain-of-thought prompting

Chain-of-thought prompting works by including a series of intermediate reasoning steps into any examples that you use for one or few-shot inference. By structuring the examples in this way, you’re essentially teaching the model how to reason through the task to reach a solution.

Program-aided language models (PAL)

PAL pairs LLM with an external code interpreter to carry out calculations. The strategy behind PAL is to have the LLM generate completions where reasoning steps are accompanied by computer code. This code is then passed to an interpreter to carry out the calculations necessary to solve the problem. You specify the output format for the model by including examples for one or few shot inferences in the prompt.

To prepare for inference with PAL,

• Format prompt to contain one or more examples; each example should contain a question followed by reasoning steps in lines of Python code that solves the problem
• Append the new question to the prompt template; the resulting PAL formatted prompt now contains both the example and the problem to solve
• Pass the combined prompt to LLM, which generates a completion that is in the form of a Python script having learned how to format the output based on the example in the prompt
• Hand off the script to a Python interpreter, which will be used to run the code and generate an answer
• Append the text containing the answer, forming a prompt that includes the correct answer in the context
• Pass the updated the prompt to LLM, and it generates a completion that contains the correct answer

How to automate the process? The orchestrator is a technical component that manage the flow of information and the initiation of calls to external data sources or applications. It can also decide what actions to take based on the information contained in the output of the LLM. The LLM doesn’t have to decide to run the code, it just has to write the script which the orchestrator then passes to the external interpreter to run.

ReAct: Combining reasoning and action

ReAct is a prompting strategy that combines chain-of-thought reasoning with action planning. It uses structured examples to show a LLM how to reason through a problem and decide on actions to take that move it closer to a solution.

The LangChain framework provides modular pieces that contain the components necessary to work with LLMs. These components include prompt templates for many different use cases that you can use to format both input examples and model completions, and memory that we can use to store interactions with an LLM. The framework also includes pre-built tools that carry out a variety of tasks, including calls to external datasets and various APIs. Connecting a selection of these individual components together results in a chain.

The creators of LangChain have developed a set of predefined chains that have been optimized for different use cases. You can use these off the shelf to quickly get your app up and running. Sometimes your application workflow could take multiple paths depending on the information the user provides. In this case, you can’t use a predetermined chain, but instead we’ll need the flexibility to decide which actions to take as the user moves through the workflow. You can use an agent to interpret the input from the user and determine which tool or tools to use to complete the task. LangChain currently includes agents for both PAL and ReAct, among others. Agents can be incorporated into chains to take an action or plan and execute a series of actions.

LLM application architectures

References

• Chain-of-thought Prompting Elicits Reasoning in Large Language Models
  • Paper by researchers at Google exploring how chain-of-thought prompting improves the ability of LLMs to perform complex reasoning.
• PAL: Program-aided Language Models
  • This paper proposes an approach that uses the LLM to read natural language problems and generate programs as the intermediate reasoning steps.
• ReAct: Synergizing Reasoning and Acting in Language Models
  • This paper presents an advanced prompting technique that allows an LLM to make decisions about how to interact with external applications.
• LangChain Library (GitHub)
  • This library is aimed at assisting in the development of those types of applications, such as Question Answering, Chatbots and other Agents. You can read the documentation here
• Who Owns the Generative AI Platform?
  • The article examines the market dynamics and business models of generative AI.

Why we need RLHF

RLHF helps to avoid

• Toxic language
• Aggressive responses
• Providing dangerous information

One potentially exciting application of RLHF is the personalization of LLMs, where models learn the preference of each individual user through a continuous feedback process. Such as individualized learning plans and personalized AI assistant.

Reinforcement learning

Reinforcement learning is a type of machine learning where an agent learns to make decisions related to a specific goal by taking actions in an environment, with the objective of maximizing some notion of a cumulative reward. In this framework, the agent continually learns from its experiences by taking actions, observing the resulting changes in the environment, and receiving rewards or penalties, based on the outcomes of its actions. By iterating through this process, the agent gradually refines its strategy or policy to make better decisions and increase its chances of success.

How to determine the reward when fine-tuning LLM with reinforcement learning?

• One way is to have a human evaluate all of the completions of the model against some alignment metric, such as determining whether the generated text is toxic or non-toxic. This feedback can be represented as a scalar value, either 0 or 1. The LLM weights are then updated iteratively, to maximize the reward obtained from the human classifier, enabling the model to generate non-toxic completions.
• A scalable and practical alternative is the reward model:
  • Start with a smaller number of human examples to train the reward model
  • Once trained, use the reward model to assess the output of the LLM and assign a reward value, which in turn gets used to update the weights of the LLM and train a new human aligned version.
  • Exactly how the weights get updated as the model completions are assessed depends on the algorithm used to optimize the policy.

RLHF: preparation

The first step is to select a model to work with and prepare a dataset for human feedback.

The model should have some capability to carry out the task of interest, whether it is text summarization, question answering, or something else. In general, recommend to start with an instruct model that has already been fine-tuned across many tasks and has some general capabilities.

Prepare dataset for human feedback

Use the selected instruct LLM with a prompt dataset to generate a number of different responses for each prompt. The prompt dataset is composed of multiple prompts, each of which gets processed by the LLM to produce a set of completions.

Collect human feedback

The next step is to collect feedback from human labelers on the completions generated by the LLM. This is the human feedback portion of the RLHF.

• Decide the criteria for the labelers to assess the completions on. This could be helpfulness, toxicity, etc.
• Ask the labelers to assess each completion in the dataset based on the criteria.
• Assign the same task to multiple labelers to dampen the impact of “poor” labelers who may misunderstood the instruction.

Note that the clarify of instructions can make a big difference on the quality of the human feedback.

Prepare labeled data for training

Before training the reward model, convert the ranking data into a pairwise comparison of completions. Re-order the pair so the preferred response comes first. This is an important step because the reward model expects the preferred completion \(y_{j}\) first. Once completing this data restructuring, the human responses will be in the correct format for training the reward model.

RLHF: reward model

By the time we are done with training the reward model, we won’t need to include any more humans in the loop. Instead, the reward model will effectively take place of the human labeler and automatically choose the preferred completions during the RLHF process.

The reward model is also usually a language model which is trained using the supervised learning on the pairwise comparison that we prepared from the human labelers’ assessment of the prmpts.

The human preferred completion is always the first one, labeled as \(y_{j}\). For a given prompt \(x\), the reward model learns to favor the human-preferred completion \(y_{j}\), while minimizing the log sigmoid of the reward difference, \(r_{j} - r_{k}\).

Once the model has been trained on the human ranked prompt-completion pairs, we can use the reward model as a binary classifier to provide a set of logits (i.e. reward values) across the positive and negative classes. The logits are the unnormalized model outputs before applying any activation function. If we apply the softmax function to the logits, we get probabilities. The examples below (Figure 9) show a good and a bad example for the rewards.

RLHF: fine-tuning

Start with a model that already has good performance on the task of interest, i.e. an instruction fine-tuned LLM. A higher reward value represents a more aligned response, vice versa.

Proceed the fine-tuning with multiple iterations. In the following example (Figure 11), we can see that the completion generated by the RL-updated LLM receives a higher reward score, indicating that the updates to weights have resulted in a more aligned completion. If the process works well, we will see the reward improving after each iteration as the model produces text that is increasingly aligned with human preferences.

We continue the iterative process until the model is aligned based on some evaluation criteria. For example, reaching a threshold value for the helpfulness we defined. We can also define a maximum number of steps, for example, 20, 000 as the stopping criteria.

At this point, we can refer the fine-tuned model as the human-aligned LLM.

Proximal policy optimization

Proximal policy optimization (PPO) is the most widely used alogirthm for the reinforcement learning step in RLHF pipelines. PPO optimizes a policy, in this case the LLM, to be more aligned with human preferences. So the goal is to update the policy to maximize the reward. Over many iterations, PPO makes updates to the LLM. The updates are small and within a bounded region, resulting in an updated LLM that is close to the previous version. Keeping the changes within the small region results in a more stable learning.

At a high level, each cycle of PPO goes over 2 phases.

PPO Phase 1: Create completions

In phase 1, the LLM is used to carry out a number of experiments to complete the given prompts. These experiments allow us to update the LLM against the reward model in phase 2.

Recall that the reward model captures human preference. For example, the reward can define how helpful, harmless, and honest the responses are. The expected reward of a completion is an important quantity used in the PPO objective. We estimate this quantity through a separate head of the LLM called the value function.

How to calculate value loss:

Assume a number of prompts are given, we first generate LLM response to the prompts, then calculate the reward for the prompt completions using the reward model.

• The value function estimates the expected total rewards for a given state s. In other words, as the LLM generates each token of a completion, we want to estimate the total future reward based on the current sequence of tokens. Think of this as a baseline to evaluate the quality of completions agains the alignment criteria. In the following example, with the completion of the first token, the total estimated future reward is 0.34; with the next generated token, the estimated future total reward increases to 1.23.

• The goal is to minimize the value loss, which is the difference between the actual future total reward and the estimation to the value function.

The value loss makes the estimates for future rewards more accurate. The value function is then used in Advantage Estimation in phase 2.

PPO Phase 2: Model update

In phase 2, we make small updates to the model, and evaluate the impact of those updates on the alignment goal. The model weights updates are guided by the prompt completion, losses, and rewards.

PPO ensures to keep the model updates within a certain small region called the trust region. This is where the proximal aspect of PPO comes into play. Ideally, this series of small updates will move the model towards higher rewards.

The policy loss is the main objective that the PPO algorithm tries to optimize during training.

Calculate policy loss

• \(a_{t}\) is the next token, \(s_{t}\) is the completed prompt up to the token \(t\).
• The denominator of the first term is the probability of the next token with the initial version of the LLM which is frozen. The numerator is the probability of the next token, through the updated LLM, which we can change for the better reward. \(\hat{A_{t}}\) is called the estimated advantage of a given choice of action. The advantage term estimates how much better or worse the current action is compared to all possible actions at that state. So we look at the expected future rewards of a completion following the new token, and we estimate how advantageous this completion is compared to the rest.
  • A positive \(\hat{A_{t}}\) means the suggested token is better than the average. Therefore, increasing the probability of the current token seems like a good strategy that leads to higher rewards. This translates to maximizing the expression we have here. If the suggested token is worse than average, the advantage will be negative. Then maximizing the expression will demote the token, which is the correct strategy.
  • So the overall conclusion is that maximizing this expression results in a better aligned LLM.
  • Directly maximizing the expression would lead to problems because our calculation is reliable under the assumption that our advantage estimations are valid. The advantage estimations are valid only when the old and new policies are close to each other. This is where the rest of the terms come into play.
• The second term defines a region where the two policies are near each other. These extra terms are guardrails, and simply define a region in proximity to the LLM, where our estimates have small errors. This is called the trust region. These extra terms ensure that we are unlikely to leave the trust region.
• In summary, optimizing the PPO policy objective results in a better LLM without overshooting to unreliable regions.

Calculate entropy loss

• While the policy loss moves the model toward the alignment goal, entropy loss allows the model to maintain creativity. If we keep entropy low, we may end up always completing the prompt in the same way. This is similar to the temperature setting. The difference is that temperature influences model creativity during the inference time, while the entropy influences the model creativity during training.

Once the model weights are updated, PPO starts a new cycle. For the next iteration, the LLM is replaced with the updated LLM, and a new PPO cycle starts. After many iterations, we arrive at the human-aligned LLM.

RLHF: reward hacking

As the policy tries to maximize the reward, it can diverge too much from the original language model.

To avoid reward hacking, we can use the instruct model as a reference model. The weights of the reference model are frozen, and are not updated during iterations of RLHF. This way, we always maintain a single reference model to compare to.

During training, each prompt is passed to both models, generating a completion by the reference LLM and the intermediate LLM updated model. At this point, we can compare the two completions and calculate a value called the KL divergence. KL divergence is a statistical measure of how different two probability distributions are. We can use it to compare the completions of the two models and determine how much the updated model has diverged from the reference model.

KL divergence is calculated for each generated token across the whole vocabulary of the LLM. This can easily be tens or hundreds or thousands of tokens. Although softmax gives a probability for every token, in practice, most of the probability mass is concentrated on a small number of likely tokens, making the distribution effectively sparse. This is still a relatively expensive process. So we will always benefit from using GPUs.

Once we calculate the KL divergence between the two models, we add it as a term to the reward calculation. This will penalize the RL updated model if it shifts too far from the reference LLM and generates completions that are too different.

Note that we need two full copies of the LLM to calculate KL divergence, the frozen reference LLM, and the RL updated PPO LLM. We can benefit from combining RLHF with PEFT. In this case, we only update the weights of a PEFT adapter, not the full weights of the LLM. This means we can reuse the same underlying LLM for both the reference model and the PPO model, which we update with trained PEFT parameters. This reduces the memory footprint during training by approximately half.

Scaling human feedback

Constitutional AI is a method for training models using a set of rules and principles that govern the model’s behavior. We train the model to self-critique and revise its response to comply with those principles.

Constitutional AI is useful not only for scaling feedback, it can also help to address some unintended consequences of RLHF. For example, depending on how the prompt is structured, an aligned model may end up revealing harmful information as it tries to provide the most helpful response it can. Providing the model with a set of constitutional principles can help the model balance these competing interests and minimize the harm.

References

• Training language models to follow instructions with human feedback
  • Paper by OpenAI introducing a human-in-the-loop process to create a model that is better at following instructions (InstructGPT).
• Learning to summarize from human feedback
  • This paper presents a method for improving language model-generated summaries using a reward-based approach, surpassing human reference summaries.
• Proximal Policy Optimization Algorithms
  • The paper from researchers at OpenAI that first proposed the PPO algorithm. The paper discusses the performance of the algorithm on a number of benchmark tasks including robotic locomotion and game play.
• Direct Preference Optimization: Your Language Model is Secretly a Reward Model
  • This paper presents a simpler and effective method for precise control of large-scale unsupervised language models by aligning them with human preferences.
• Constitutional AI: Harmlessness from AI Feedback
  • This paper introduces a method for training a harmless AI assistant without human labels, allowing better control of AI behavior with minimal human input.

Why Parameter efficient fine-tuning (PEFT)

Full fine-tuning results in a new version of the original LLM for every task trained on. Each of the new version is the same size as the original model, so it can create an expensive storage problem if fine-funing for multiple tasks.

In contrast to full fine-tuning where every model weight is updated during supervised learning, PEFT only updates a small subset of parameters. Some PEFT techniques freeze most of the model weights and focus on fine-tuning a subset of existing model parameters, for example, particular layers or components. Other techniques don’t touch the model weights at all, and instead add a small number of new parameters or layers and fine-tune only the new components.

With PEFT, most if not all of the model weights are kept frozen. As a result, the number of trained parameters is much smaller than the number of parameters in the original LLM. In some cases, just 15-20 % of the original LLM weights. This makes the memory requirements for training much more manageable. In fact, PEFT can often be performed on a single GPU. And because the original LLM is only slightly modified or left unchanged, PEFT is less prone to the catastrophic forgetting problem of full fine-tuning.

PEFT methods

Selective

Select subset of initial LLM parameters to fine-tune. There are several methods you can take to identify which parameters you want to update. But performance of these methods (i.e., train only certain components, layers, or just individual parameter types) are mixed and there are significant trade-offs between parameter efficiency and compute efficiency.

Reparameterization

Reparameterize model weights using a low-rank representation. It works with the original LLM parameters, but reduces the number of parameters to train by creating new low rank transformations of the original network weights.

Additive

Add trainable layers or parameters to model. This fine-tuning method keeps all of the original LLM weights frozen and introducing new trainable components. For example,

• Adapter adds new trainable layers to the architecture of the model, typically inside the encoder or decoder components after the attention or feed-forward layers.
• Soft prompts keep the model architecture fixed and frozen, and focus on manipulating the input to achieve better performance. This can be done by adding trainable parameters to the prompt embeddings or keeping the input fixed and retraining the embedding weights.

LoRA

LoRA == Low-Rank Adaptation of Large Language Models

LoRA workflow

1. Freeze most of the original LLM weights, i.e. weights applied to embedding vectors

2. Inject 2 rank decomposition matrices

   • The dimensions of the smaller matrices are set so their product is a matrix with the same dimensions as the weights they are modifying.

3. Train the weights of the smaller matrices

   • Keep the original weights of the LLM frozen and train the smaller matrices.

4. Inference

   • For inference, the two low-rank matrices are multiplied together to create a matrix with the same dimensions as the frozen weights. Add this to the original weights and replace them in the model with the updated values.

Note:

• Because the LoRA fine-tuned model has the same number of parameters with the original model, there is little to no impact on inferenc latency.
• Researchers have found that applying LoRA to just the self-attention layers of the model is often enough to fine-tune for a task and achieve performance gains. In principle, you can also use LoRA on other components like the feed-forward layers. However, since most of the parameters of LLMs are in the attention layers, you get the biggest savings in trainable parameters by applying LoRA to these weights matrices.
• Because LoRA allows you to significantly reduce the number of trainable parameters, you can often perform this PEFT method with a single GPU and avoid the need for a distributed cluster of GPUs.

LoRA: math behind

Use the base Transformer model presented by Vaswani et al. 2017 as an example, consider LoRA wight rank \(r = 8\).

Base weight:

LoRA low-rank matrices:

Total LoRA parameters trained:

Parameter reduction:

Understanding why the total parameters during the inference stays the same.

During training:

• The original weight matrix \(W_0\) is frozen
• Two small matrices \(A \in \mathbb{R}^{r \times k}\), \(B \in \mathbb{R}^{d \times r}\) are trained
• The effective weight becomes \(W = W_0 + \Delta W = W_0 + B \cdot A\)
• The model learns the low-rank update \(\Delta W\), which is far smaller in size then \(W_0\)

At inference:

• Merge the update \(B \cdot A\) into \(W_0\) ahead of time
• So still using a matrix of size \(d x k\), not adding more layers or extra computations during forward pass
• Thus, inference speed and memory stay the same as the original model

LoRA for different tasks

Since the rank-decomposition matrices are small, you can fine-tune a different set for each task and then switch them out at inference time by updating the weights.

Suppose you train a pair of LoRA matrices for a specific task, Task A. To carry out inference on this task, you would multiply these matrices together and then add the resulting matrix to the original frozen weights. You then take this summed weight matrix and replace the original weights where they appear in your model. You can then use this model to carry out inference on Task A. And you can repeat the same process for task B.

LoRA: other notes

ROUGE metrics from LoRA fine-tuning only slightly suffers compared to that from full fine-tuning.

How to choose the rank?

• The smaller the rank, the smaller number of trainable parameters, and the bigger saving on compute.
• Ranks in the range of 4 to 32 can provide a good trade-off between reducing trainable parameters and preserving performance.
• Optimizing the choice of rank is an ongoing area of research and best practices may evolve.

Soft prompt

Prompt tuning is not prompt engineering. With prompt tuning, you add additional trainable tokens to your prompt and leave it up to the supervised learning process to determine their optimal values. The set of trainable tokens is called a soft prompt, and it gets prepended to embedding vectors that represent your input text. The soft prompt vectors have the same length as the embedding vectors of the language tokens. Usually 20-100 tokens can be sufficient for good performance.

The tokens that represent natural language are hard in the sense that they each correspond to a fixed location in the embedding vector space. However, the soft prompts are not fixed discrete words of natural language. Instead, you can think of them as virtual tokens that can take on any value within the continuous multidimensional embedding space. Through supervised learning, the model learns the values for these virtual tokens that maximize performance for a given task.

In full fine-tuning, the training data set consists of input prompts and output completions or labels. The weights of the LLM are updated during supervised learning. In contrast, with prompt tuning, the weights of the LLM are frozen and the underlying model does not get updated. Instead, the embedding vectors of the soft prompt gets updated over time optimize the model’s completion of the prompt. Prompt tuning is a very parameter efficient strategy because only a few parameters are being trained. Similar to LoRA, you can train a different set of soft prompts for each task and then easily swap them out at inference time.

Prompt tuning doesn’t perform as well as full fine-tuning for smaller LLMs. However, as the model size increases, so does the performance of prompt tuning.

Because the soft prompt tokens can take any value within the continuous embedding vector space, the trained tokens don’t correspond to any known token, word, or phrase in the vocabulary of the LLM. Therefore interpretability of the virtual tokens can be a potential issue. An analysis of the nearest neighbor tokens to the soft prompt location shows that they form tight semantic clusters. That is, the words closest to the soft prompt tokens have similar meanings. The words identified usually have some meaning related to the task, suggesting that the prompts are learning word-like representations.

References

• Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning
  • This paper provides a systematic overview of Parameter-Efficient Fine-tuning (PEFT) Methods in all three categories discussed in the lecture videos.
• On the Effectiveness of Parameter-Efficient Fine-Tuning
  • The paper analyzes sparse fine-tuning methods for pre-trained models in NLP.
• LoRA Low-Rank Adaptation of Large Language Models
  • This paper proposes a parameter-efficient fine-tuning method that makes use of low-rank decomposition matrices to reduce the number of trainable parameters needed for fine-tuning language models.
• QLoRA: Efficient Finetuning of Quantized LLMs
  • This paper introduces an efficient method for fine-tuning large language models on a single GPU, based on quantization, achieving impressive results on benchmark tests.
• The Power of Scale for Parameter-Efficient Prompt Tuning
  • The paper explores “prompt tuning,” a method for conditioning language models with learned soft prompts, achieving competitive performance compared to full fine-tuning and enabling model reuse for many tasks.

LLM fine-tuning at a high level

Why we need LLM fine-tuning

Recall In-context learning: zero-shot, one-shot, and few-shot inference. See this note for more details.

We need LLM fine-tuning because in-context learning has several drawbacks:

• For smaller models, the in-context learning doesn’t always work, even when 5 or 6 examples are included.
• Any examples you include in your prompt take up valuable space in the context window, reducing the amount of room you have to include other useful information.

What is LLM fine-tuning

Fine-tuning is a supervised learning process, where you use a dataset of labeled examples to update the weights of the LLM.

Instruction fine-tuning trains the model using examples demonstrating how it should respond to a specific instruction. The labeled examples are prompt-completion pairs; the fine-tuning process extends the training of the model to improve its ability to generate good completions for a specific task.

Each example in the prompt-completion pairs datasets begins with instructions. For example, if you want to fine-tune your model to improve its summarization ability, you’d build up a dataset of examples that begin with instruction “summarize the following text” or a similar phrase; if you are improving the model’s translation skills, your examples would include instructions like “translate this sentence”. These examples allow the model to learn to generate responses following given responses.

Instruction fine-tuning, where all of the model’s weights are updated, is known as full fine-tuning. The process results in a new version of the model with updated weights. Note that just like pre-training, full fine-tuning requires enough memory and compute budget to store and process all the gradients, optimizers, and other components that are being updated during training.

LLM fine-tuning process

First step of instruction fine-tuning is to prepare your training data. There is much publicly available data that has been used to train earlier generations of LLMs, but not all of them are formatted as instructions. Developers have assembled prompt template libraries that can be used to take existing dataset. For example,turn the large dataset of Amazon product review into instruction prompt datasets for fine-tuning (see Figure 3, source). In each case, you pass in the original review, here called review_body, to the template, where it gets inserted into the text that starts with an instruction. The result is a prompt that now contains both an instruction and the example from the dataset.

Once the instruction dataset is ready, we divide the dataset into training, validation, test split.

During fine-tuning, we select prompts from the training set and pass them to the LLM which then generates completions. Next, we compare LLM completion with the response specified in the training data.

Recall that the output of an LLM is a probability distribution across tokens. So we can compare the distribution of the completion and that of the training label, and use the standard cross entropy function to calculate loss between the two token distributions. Then use the calculated loss to update model weights in standard backpropagation. Do this for many batches of prompt completion pairs and over several epochs, update the weights so that the model’s performance on the task improves. Then get the validation_accuracy using the holdout validation set, and get the test_accuracy once you apply the model to the test set.

Single task vs. Multi-task

Single task fine-tuning

Good results can be achieved with relatively few examples (often 500-1000 examples) for single task fine-tuning. However, fine-tuning on a single task can lead to catastrophic forgetting, i.e. fine-tuning significantly improves performance of the model on a specific task but degrades performance on other tasks.

How to avoid catastrophic forgetting?

• First of all, you might not have to if you only care about the task you fine-tuned for
• Fine-tune on multiple tasks at the same time
• Consider Parameter Efficient Fine-tuning (PEFT)

Multi-task fine-tuning

Over many epochs of training, the calculated losses across examples are used to update the weights of the model, resulting in an instruction fine-tuned model that learned how to be good at many different tasks simultaneously.

Drawback to multi-task fine-tuning: requires a lot of data. However, it can be very worthwhile and worth the effort to assemble this data. The resulting models are often very capable and suitable for use in situations where good performance at many tasks is desirable.

Instruction fine-tuning with FLAN

FLAN == Fine-tuned LAnguage Net. FLAN models refer to a specific set of instructions used to perform instruction fine-tuning. FLAN-T5 (fine-tuned version of pre-trained T5 model) is a great, general purpose, instruct model.

LLM evaluation

Metrics: ROUGE

ROUGE is used for text summarization and text generation tasks. Figures in below demonstrate usage of ROUGE-1, ROUGE-2, ROUGE-L, and ROUGE clipping.

ROUGE-1 measures the overlap of unigrams (single words) between the generated and reference texts.

ROUGE-2 measures the overlap of bigrams (two-word sequences) between the generated and reference texts.

ROUGE-L measures the longest common subsequences (LCS) between the generated and reference texts, capturing setence-level fluency and structure.

ROUGE clipping limits the count of overlapping n-grams to the maximum number that appears in the reference, preventing the generated text from getting extra credit for repeated words.

Metrics: BLEU SCORE

BLEU SCORE is used for text translation. Below is the core BLEU formula, from Papineni et al. 2002.

• BLEU: The final BLEU score (0 to 1), evaluating how closely the generated text matches the reference.
• BP: Brevity Penalty — penalizes candidates that are shorter than the reference.
• \(\exp\): Exponential function, used to combine the log-precisions into a geometric mean.
• \(w_n\): Weight for each n-gram order (e.g., 0.25 when using up to 4-grams).
• \(p_n\): Modified precision for n-grams of size \(n\), with clipping to avoid over-counting.
• \(N\): Maximum n-gram order (usually 4 in practice).

Benchmarks

Selecting an evaluating dataset is vital to an accurate evaluation of model performance. Example evaluation benchmarks include GLUE, SuperGLUE, MMLU (Massive Multitask Language Understanding), HELM, Big-bench, etc.

References

• Scaling Instruction-Finetuned Language Models
  • Scaling fine-tuning with a focus on task, model size and chain-of-thought data.
• Introducing FLAN: More generalizable Language Models with Instruction Fine-Tuning
  • This blog (and article) explores instruction fine-tuning, which aims to make language models better at performing NLP tasks with zero-shot inference.
• HELM - Holistic Evaluation of Language Models
  • HELM is a living benchmark to evaluate Language Models more transparently.
• General Language Understanding Evaluation (GLUE) benchmark
  • This paper introduces GLUE, a benchmark for evaluating models on diverse natural language understanding (NLU) tasks and emphasizing the importance of improved general NLU systems.
• SuperGLUE
  • This paper introduces SuperGLUE, a benchmark designed to evaluate the performance of various NLP models on a range of challenging language understanding tasks.
• ROUGE: A Package for Automatic Evaluation of Summaries
  • This paper introduces and evaluates four different measures (ROUGE-N, ROUGE-L, ROUGE-W, and ROUGE-S) in the ROUGE summarization evaluation package, which assess the quality of summaries by comparing them to ideal human-generated summaries.
• BLEU: a Method for Automatic Evaluation of Machine Translation
• Measuring Massive Multitask Language Understanding (MMLU)
  • This paper presents a new test to measure multitask accuracy in text models, highlighting the need for substantial improvements in achieving expert-level accuracy and addressing lopsided performance and low accuracy on socially important subjects.
• BigBench-Hard - Beyond the Imitation Game: Quantifying and Extrapolating the Capabilities of Language Models
  • The paper introduces BIG-bench, a benchmark for evaluating language models on challenging tasks, providing insights on scale, calibration, and social bias.

Model cards

Model cards are useful for understanding how a model is trained, its use case, and known limitations. See here for a more complete definition from HuggingFace.

Figures 1-2 are two example model cards from Mitchel et al. (2018). Figure 3 is the model card for T5 Large, captured from the lecture slides.

Token filtering

Only 1-3% of original tokens are used to train the LLM. We should consider this when we estimate how much data we need to collect if we decide to pre-train our own model. In other words, data quality control or bias removal is always needed to increase data quality before training. This refers to the token filtering or rejection process during LLM pre-training.

LLM developers all start from massive web-scale corpora, then spend huge effort on filtering, deduplicating, scoring for quality, and weighting data by domain or value. See GPT-4 Technical Report, LLaMA, and Chinchilla for details.

LLMs pre-training

Different LLMs pretrain using different objectives - from denoising (such as BERT), to predicting next tokens (such as GPT), to reconstructing corrupted input using a sequence-to-sequence setup (such as T5 and BART). These design choices affect how models behave. For example, BERT is optimized for understanding tasks, GPT excels at generating fluent and coherent outputs. T5 and BART aim to balance both.

Autoencoding models

Autoencoding models like BERT are pre-trained using masked language modeling. In the context of autoencoding models, the pre-training task is often described as a “denoising” objective because the model learns to reconstruct the original text from a corrupted (noisy) version of it.

This setup provides the model with bi-directional context, meaning it can consider both preceding and following words when predicting masked tokens.

Autoregressive models

Autoregressive models are pretrained using causal language modeling, where the model learns to predict the next token based solely on the preceding sequence of tokens. Note: predicting the next token is sometimes called full language modeling by researchers.

This setup provides unidirectional context, meaning the model has no access to future tokens during training. Autoregressive models mask the input sequence and can only see the input tokens leading up to the token in question. The model has no knowledge of the end of the sequence. The model then iterates over the input sequence one by one to predict the following token.

Large decoder-only models, such as GPT, demonstrate strong ability in zero-shot and few-shot inferences, and can generalize well across a wide range of language tasks.

Seq2Seq models

Exact detail of pre-training varies from model to model.

• A popular Seq2Seq model T5 pretrains the encoder using span corruption, which masks random sequences of input tokens. Those masked sequences are then replaced with a unique sentinel token. Sentinel tokens are special tokens added to the vocabulary, but do not correspond to any actual word from the input text. The decoder is then tasked with reconstructing the masked token sequences auto-regressively. The output is the sentinel token followed by the predicted tokens.
• BART combines the ideas of masked language modeling and causal language modeling within a seq2seq (encoder-decoder) framework. It corrupts text with various noise functions (e.g., setence shuffling, token masking), then uses a decoder to reconstruct the original text.

Model size

The growth of model size is powerd by

• introduction of transformer
• access to massive datasets
• more powerful compute resources

Compute challenge

CUDA, short for Compute Unified Device Architecture, is a collection of libraries and tools developed for NVIDIA GPUs to boost performance on common deep-learning operations, including matrix multiplication, among many others. Deep-learning libraries such as PyTorch and TensorFlow use CUDA extensively to handle the low-level, hardware-specific details, including data movement between CPU and GPU memory.

Quote from Generative AI on AWS: Building Context-Aware Multimodal Reasoning Applications (source):

A single-model parameter, at full 32-bit precision, is represented by 4 bytes. Therefore, a 1-billion-parameter model requires 4GB of GPU RAM just to load the model into GPU RAM at full precision. If you want to also train the model, you need more GPU memory to store the states of the numerical optimizer, gradients, and activations, as well as any temporary variables used by your functions.

Quote from Generative AI on AWS: Building Context-Aware Multimodal Reasoning Applications (source):

When you experiment with training a model, it’s recommended that you start with batch_size=1 to find the memory boundaries of the model with just a single training example. You can then incrementally increase the batch size until you hit the CUDA out-of-memory error. This will determine the maximum batch size for the model and dataset. A larger batch size can often speed up your model training.

These additional components lead to approximately 12–20 extra bytes of GPU memory per model parameter. For example, to train a 1-billion-parameter model, you will need approximately 24 GB of GPU RAM at 32-bit full precision, six times the memory compared to just 4 GB of GPU RAM for loading the model

One technique we can use to reduce memory is called quantization. The idea is to reduce the memory required to store the model weights by reducing their precision from 32-bit floating point numbers to 16-bit (or 8-bit) floating point numbers. Quantization statistically projects the original 32-bit floating point numbers into a lower precision space, using scaling factors calculated based on the range of the original 32-bit floating point numbers.

BFLOAT-16 (BF16, developed at Google Brain) has recently become a popular alternative to FP16.

• BF16 is a hybrid between half precision FP16 and full precision FP32.
• BF16 significantly helps with training stability and is supported by newer GPUs such as NVIDIA’s A100.
• BF16 is often described as truncated 32-bit float, as it captures the full dynamic range of the full 32-bit float that uses only 16-bits.
• BF16 uses the full 8 bits to represent the exponent, but truncates the fraction to just 7 bits.
• Downside: BF16 is not well suited for integer calculations, but these are relatively rare in deep learning.

Scalling law

The Chinchilla scaling law shows that, for a fixed compute budget, language model performance is optimized by using more training data and a smaller model size (Hoffmann et al., 2022). Undertraining is the bottleneck - large models like GPT-3 were not trained on enough tokens to fully benefit from their size.

Domain adaptation

Example: BloombergGPT

References

• Generative AI on AWS: Building Context-Aware, Multimodal Reasoning Applications
  • Deep dive into all phases of the generative AI Lifecycle
• BLOOM: BigScience 176B Model
  • BLOOM is an open-source LLM with 176B parameters trained in an open and transparent way. In this paper, the authors present a detailed discussion of the dataset and process used to train the model. You can also see a high-level overview of the model here
• Vector Space Models
  • Series of lessons from DeepLearning.AI’s Natural Language Processing specialization discussing the basics of vector space models and their use in language modeling.
• Scaling Laws for Neural Language Models
  • Pre-training and scaling laws
• What Language Model Architecture and Pretraining Objective Work Best for Zero-Shot Generalization? Model architectures and pre-training objectives. The paper examines modeling choices in large pre-trained language models and identifies the optimal approach for zero-shot generalization.
• HuggingFace Tasks and Model Hub
  • Collection of resources to tackle varying machine learning tasks using the HuggingFace library.
• LLaMA: Open and Efficient Foundation Language Models
  • Article from Meta AI proposing Efficient LLMs (their model with 13B parameters outperform GPT3 with 175B parameters on most benchmarks)
• Language Models are Few-Shot Learners
  • Scaling laws and compute-optimal models. This paper investigates the potential of few-shot learning in Large Language Models.
• Training Compute-Optimal Large Language Models
  • Study from DeepMind to evaluate the optimal model size and number of tokens for training LLMs. Also known as “Chinchilla Paper”.
• BloombergGPT: A Large Language Model for Finance
  • LLM trained specifically for the finance domain, a good example that tried to follow chinchilla laws.

Transformer variants

Encoder only models

• Also known as Autoencoding models.

• Input sequence and output sequence are the same length.

• Use cases

  • Sentiment analysis
  • Named entity recognition
  • Word classification

• Examples

  • BERT
  • RoBERTa

Encoder Decoder models

• Also known as Seq2Seq models.

• Input sequence and output sequence can be different lengths.

• Use cases

  • Translation
  • Text summarization
  • Question answering

• Examples

  • BART
  • T5

Decoder only models

• Also known as Autoregressive models.

• Widely used today. Can generalize to most tasks.

• Use cases

  • Text generation

• Examples

  • GPTs
  • BLOOM
  • Jurassic
  • LLaMA
  • etc

Prompt engineering

In-context learning

In-context learning: including examples of the task that you want the model to carry out inside the prompt can get the model to produce better outcomes. Providing examples inside the context window is called in-context learning.

Zero-shot, one-shot, few-shot

Zero-shot inference: including zero examples.

One-shot inference: including a single example.

Few-shot inference: including several examples.

Note:

• Large LLMs can do well with zero-shot inference, smaller LLMs struggle with zero-shot inference and benefit by learning examples of the desired behavior.

• You have a limit on the amount of in-context learning that you can pass to the model. If you find your model isn’t performing well when including 5 or 6 examples, you should try fine-tuning your model instead.

Generative configuration - inference parameters

Max new tokens

• Used to limit the number of tokens that the model will generate. You can think of this as putting a cap on the number of times the model will go through the selection process (i.e., from computing probability distribution to sampling or selecting the next token).
• Remember it’s max new tokens, not a hard number of new tokens generated.

Top-k and Top-p

Recall that the output from the transformer’s softmax layer is a probability distribution across the entire dictionary of words that the model uses.

Greedy decoding

• Model chooses the word/token with the highest probability.
• This method can work very well for short generation but is susceptible to repeated words or repeated sequences of words. If you want to generate texts that’s more natural, more creative, and avoid repeating words, you need to use some other controls.

Random (weighted) sampling

• The model chooses the output word at random using the probability distribution to weight the selection.
• Reduce the likelihood that words will be repeated. However, depending on the setting, there is a possibility that the output may be too creative, producing words that cause the generation to wander off into topics or words that just don’t make sense.

Top-k and top-p are sampling methods we can use to limit the random sampling and increase the chance that the output will be sensible.

Top-k

• The model to choose from only the \(k\) tokens with the highest probability. The model then selects from these options using the probability weighting
• This method can help the model have some randomness while preventing the selection of highly improbable completion words. This in turn makes your text generation more likely to sound reasonable and to make sense.

Top-p

Limit the random sampling to the predictions whose combined probabilities do not exceed \(p\).

Temperature

• This parameter influences the shape of the probability distribution that the model calculates for the next token.
• The temperature is a scaling factor that’s applied within the final softmax layer of the model that impacts the shape of the probability distribution of the next token. In contrast to top-k and top-p, changing temperature actually alters the predictions the model will make.
• Cooler temperature (\(<1\)) - strongly peaked distribution.
• Warmer temperature (\(>1\)) - broader, flatter distribution.

GenAI lifecycle

References

• Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., … & Polosukhin, I. (2017). Attention Is All You Need. arXiv preprint arXiv:1706.03762.

  • The Transformer architecture, with the core “self-attention” mechanism.

• Radford, A., Narasimhan, K., Salimans, T., & Sutskever, I. (2018). Improving Language Understanding by Generative Pre-Training. OpenAI Technical Report.

  • BERT

• Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., & Stoyanov, V. (2019). RoBERTa: A Robustly Optimized BERT Pretraining Approach. arXiv preprint arXiv:1907.11692.

  • RoBERTa

• Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., … & Liu, P. J. (2020). Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. arXiv preprint arXiv:1910.10683.

  • T5

• Lewis, M., Liu, Y., Goyal, N., Ghazvininejad, M., Mohamed, A., Levy, O., … & Zettlemoyer, L. (2020). BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension. arXiv preprint arXiv:1910.13461.

  • BART

• Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., … & Amodei, D. (2020). Language Models are Few-Shot Learners. arXiv preprint arXiv:2005.14165.

  • GPT-3

• BigScience Workshop. (2022). BLOOM: A 176B-Parameter Open-Access Multilingual Language Model. arXiv preprint arXiv:2211.05100.

  • BLOOM

• AI21 Labs. (2021). Jurassic-1: Technical Details and Evaluation. AI21 Labs Blog.

  • Jurassic

• OpenAI. (2023). Language models can explain neurons in language models. arXiv preprint arXiv:2302.13971.

  • LLaMA

This post walks through my end-to-end machine learning project on predicting U.S. soybean annual yield based on local climate conditions. We cover all key steps of a machine learning workflow, including ideation, feature selection, cross validation, baseline modeling, hyperparameter tuning, and model deployment. We employed annual soybean yield data and monthly climate variables during the historical period from 1981 to 2016. The model was trained on annual soybean yield data and monthly climate variables from 1981 to 2015, with 2016 used as the test set. We found that XGBRegressor performed the best among the regression algorithms (including linear and tree-based models) tested in this project. Our model achieved an \(R^2\) of ~0.9, demonstrating a strong ability to predict U.S. soybean annual yield based on climate conditions. Please refer to this repo for the project’s code.

Introduction

Soybean is a vital source of protein for both human and animal nutrition, but its yield is highly sensitive to weather and climate variability. High temperatures and low soil moisture, particularly during the summer reproductive period, can significantly reduce soybean yields (Hamed et al., 2021). Early-season excessive precipitation can also negatively impact soybean yields by restricting root development, causing nutrient leaching, and increasing disease susceptibility (Ortiz-Bobea et al., 2019).

This project is motivated by two questions:

1. Can we build a machine learning model that predicts local annual soybean yield in the United States based on local climate conditions?
2. Which climatic factors contribute most to the soybean yield prediction?

Figure 1 (source) illustrates the life cycle of a typical end-to-end machine learning project, which consists of four main components: data preparation, feature selection, model training, and deployment. We closely followed this framework throughout the project.

The reminder of the post is structured as follows. We first introduce the dataset and conduct an exploratory data analysis to examine the temporal and spatial patterns of the historical U.S. soybean yield. We then describe the feature selection process that constructed the training dataset. The model training section covers key aspects of model development. Finally, we discuss the deployment process and conclude with an overview of caveats and potential directions for future work.

Data

We focused on two datasets: one for soybean yield and the other for climate variables. Historical annual soybean yield was provided by the Global Dataset of Historical Yields (GDHY) (lizumi and Sakai, 2020), where crop yield is defined as production per unit harvested area (\(t\) \(ha^{-1}\)). For climate variables, we used monthly model outputs from the North American Land Data Assimilation System (NLDAS) (Mitchell et al. 2004). The NLDAS dataset provides near-surface climate variables, including surface energy fluxes, surface water flux and storage, soil moisture, temperature, and land surface parameters. Both GDHY and NLDAS are land-based datasets, meaning that values over the ocean appear as NaNs. We focus on the period 1981-2016, covering the spatial domain \(235.25^{\circ} - 292.75^{\circ}E, 25.25^{\circ} - 52.75^{\circ}N\), which encompasses the primary continental regions in the U.S.

Exploratory Data Analysis (EDA)

To better understand U.S. soybean production, we conducted an exploratory data analysis (Figures 2-5). The majority of soybean production is concentrated in the eastern and midwestern states, while the Great Plains region has relatively lower production density (Figure 2). To complement this spatial view, Figure 3 presents a boxplot of the interquartile range (IQR) across all soybean production grid points for each year, highlighting spatial variations in yield Both local soybean yield and its spatial spread exhibit year-to-year variability (Figures 2-3).

Statistical analysis reveals that U.S. soybean annual yield follows a bimodal distribution (Figure 4), suggesting that not all soybean-producing regions behave uniformly. We hypothesize that this bimodal pattern arises from differences between high-production areas (Midwestern and Eastern states) and lower-production areas (Great Plains) The left mode of the distribution likely corresponds to the Great Plains, while the right mode represents more densely cultivated regions.

Both climatology and long-term trends display a west-east dipole pattern, further the Great Plains apart from other areas(Figure 5). Climatologically, soybean yield is lower in the Great Plains compared to other regions. The long-term trend, however, highlights the Great Plains as a hotspot area that experienced notable increase in soybean yield from 1981 to 2016. The overall increasing trend in the U.S. Soybean yield aligns with what we observed in Figure 3.

In this project, we train a unified model for all soybean production regions across the U.S. This approach assumes that the causal relationship between local climate conditions and soybean yield is consistent across all soybean production regions. We acknowledge that this is an oversimplification and we discuss its limitations in the Caveats and Future Work section.

We merged the GDHY and NLDAS datasets by matching the lat, lon, and year features (see Figure 6 for a screenshot of the merged dataset; see section 4.1 in this notebook). The resulting DataFrame contains over 600 climate features from NLDAS. We will perform feature selection in the next section.

Feature selection

Feature selection requires critical thinking and iterative analysis, making it the most tedious component of this project. Our goal is to narrow down the climate features to a subset that most likely affect soybean yield. This section outlines our thought process for feature selection. Detailed calculations and reasoning can be found in this notebook and this notebook.

Key considerations for feature selection:

• Feature selection should only be performed using the training dataset to prevent bias in feature and model choice from the test set.
• Feature selection does not always guarantee improved predictions.
• Linear and non-linear feature selection methods may yield conflicting results. When this happens, we rely on domain knowledge and our understanding of the problem to make decisions.

Feature selection process

We break the feature selection process into two phases:

• Phase 1: Initial feature filtering based on correlation analysis and domain knowledge.
• Phase 2: Further refinement using scikit-learn’s feature selection methods. We used features derived from Phase 1 to build baseline models, and used the refined features from Phase 2 for final model selection (see the Model Development section).

Phase 1: Initial Filtering

We followed these key principles in Phase 1:

1. Temporal consideration:

• Most U.S. soybeans are planted in May and early June and harvested in late September and October (Source: USDA). Thus, we only consider climate features during months before October, removing November and December variables from the current year.
• Besides current-year climate, we explored whether past-year climate conditions matter. In particular, we explored past year’s annual total rainfall, past winter’s snowfall, rainfall, and soil moisture.

1. Correlation analysis:

• We examined correlation between annual yield and climate features to identify important features. We kept variables with strong correlation with yield, and discarded the rest.
• Since some relationships between climate and yield are nonlinear, we kept certain features despite weak linear correlations if they were scientifically relevant.

1. Redundant feature removal:

• We aggregated (e.g., seasonal averages) climate variables that are highly auto-correlated across months to avoid redundancy.
• Some climate variables exhibit strong cross-correlations, andi we only kept the most relevant ones.
• We excluded vegetation and land cover features (such as leaf area index) that directly reflect the yield data during the soybean growing and harvest season.

1. Feature engineering: We created new features when necessary. For example, we derived an evaporative fraction using surface energy fluxes and used it instead of latent and sensible heat fluxes to simplify the feature set.

Phase 2: scikit-learn based feature selection methods

In Phase 2, we refined the feature set using various scikit-learn methods: VarianceThreshold, f_regression, mutual_info_regression, Lasso, and SelectFromModel. Although typically fewer methods are needed in practical applications, here we experimented with multiple approaches for the sake of practice. More details can be found in this notebook.

Model Development

Cross-validation

Cross-validation helps reduce overfitting in machine learning models. While scikit-learn’s train_test_split` method allows for random splitting of data, it is not suitable for our problem due to the sequential nature of the dataset. Instead, we implemented an expanding window backtesting procedure for cross-validation (Figure 7). Using this approach, we created a 5-fold split in the training set, ensuring that the test size remained constant across splits for comparability in their performance statistics:

• Split 1: train on 1980-2010, test on 2011
• Split 2: train on 1980-2011, test on 2012
• Split 3: train on 1980-2012, test on 2013
• Split 4: train on 1980-2013, test on 2014
• Split 5: train on 1980-2014, test on 2015

This method ensures that each model is trained on an expanding historical dataset, reflecting real-world forecasting conditions where future predictions are based only on past information.

Model selection

We first explored both linear models (OLS, Ridge, Lasso) and tree-based models (Decision Tree, Random Forest, XGBoost) to build baseline models. We used both \(R^2\) and RMSE to evaluate model performance. Not surprisingly, tree-based models outperformed linear models, with XGBoost achieving the best performance among the three tree-based models (Figure 8).

Hyperparameter tuning

We focused on tuning the following XGBoost hyperparameters to control overfitting:

• n_estimators: Number of rounds for boosting.
• learning_rate: Step size shrinkage used in update to prevent overfitting.
• colsample_bytree: Subsample ratio of columns when constructing each tree
• subsample: Subsample ratio of the training instances.
• gamma: Minimum loss reduction required to split a leaf node of the tree; higher values make the algorithm more conservative.
• min_child_weight: Minimum sum of instance weights needed in a child; higher values make the algorithm more conservative. -max_depth: Maximum depth of a tree; higher values increase model complexity and risk of overfitting.
• reg_alpha: L1 regularization term on weights; higher values make the model more conservative.

Grid Search, Random Search, and Bayesian Optimization are commonly used hyperparameter tuning methods. We tested both Random Search and Bayesian Optimization (via Optuna). We selected the RandomizedSearchCV-tuned parameters as they produced a slightly higher \(R^2\) value.

Feature importance

We used XGBoost’s built-in function to evaluate feature importance (Figure 9). Both longitude and year emerge as key features, reflecting the spatial-temporal nature of soybean yield prediction. Longitude serves as a proxy for geographic differences in soybean production, as observed in the EDA, while year likely captures long-term trends affecting yield. Several variables related to soil moisture content also stand out, including spring and summer soil moisture and summertime evaporative fraction. Springtime leaf area index also contributes to yield prediction, likely because it reflects pre-season vegetation health, which could correlate with soil fertility and growing conditions.

Inference and Deployment

How well does the model perform in 2016 soybean annual yield? Figure 10 compares the true versus predicted values, suggesting that the model does a decent job in capturing both the spatial pattern and local magnitude of soybean yield. However, the model overall tends to underestimate annual yield c, while overestimating yield in some localized areas. We did not investigate the cause of this underestimation, which could be an area for future analysis. We deployed the trained XGBoost model on AWS EC2, and built a web interface using Flask, basic css and javascript.

Caveats and Future Work

In this section, we discuss several limitations of the current approach and potential future improvements.

Limited predictors

Unlike some other crops (e.g. corn) that are strongly influenced by soil type, drainage, population, row width, tillage and other physical and human factors, soybean yield is more dependent on the natural environment (Source: Stack the odds for soybeans this spring). Still, additional factors such as genotype, soil type, and seeding experiments could improve our model’s predictability.

In this project, we focused solely on the predictive power of local monthly climate due to the lack of suitable datasets for other predictors. This limitation may have constrained our model’s predictability. Even for the climate data, we relied on a single source (NLDAS). Future work could involve testing alternative datasets that provide both atmospheric and land-based variables to validate or refine our findings.

Coarse timescale of crop yield

The crop yields data used in this project only provide annual yields, thus we lack information on the timing of soybean planting and seasonal growth cycle at each location and year. Several studies (Colet et al., 2023; Vann and Stokes, 2024) suggest that planting timing affects yield. Future work incorporating sub-annual yield estimates or phenological data could help capture these temporal variations.

Localized climate perspective

This work takes a localized perspective of climate impacts.hat is, we only considered impacts of local climate on local yield. However, large-scale climate patterns can also affect U.S. soybean yield through teleconnections. For example, El Niño-Southern Oscillation (ENSO) events can impact U.S. growing conditions by synchronizing climate risks across major agricultural regions (Anderson et al., 2018). Including remote climate indices as features could improve our understanding of remote climate influences on soybean yield.

Spatial variation

As shown in the EDA, U.S. soybean yield varies across the Great Plains, Midwest, and Eastern states. We hypothesize that different sub-regions are governed by different climate factors. Even a common set of climate features is relevant across all regions, their importance may vary by location. Future work could explore training separate models for different sub-regions to improve predictions for local areas.

LSTM as an alternative model architecture

Long Short-Term Memory (LSTM) networks are well-suited for time-series forecasting tasks. Given the temporal nature of our soybean yield prediction problem, incorporating LSTM in future work could provide valuable insights. Several studies have demonstrated the effectiveness of LSTMs in crop yield prediction (Sun et al, 2019; Bhimavarapu et al., 2023). However, while LSTMs excel at capturing sequential dependencies, they require larger datasets and are more computationally intensive than tree-based models like XGBoost.

References

Journal Articles:

• Anderson, W., Seager, R., Baethgen, W., & Cane, M. (2018). Trans-Pacific ENSO teleconnections pose a correlated risk to agriculture. Agricultural and Forest Meteorology. https://doi.org/10.1016/j.agrformet.2018.07.023
• Bhimavarapu, U., Battineni, G., & Chintalapudi, N. (2023). Improved optimization algorithm in LSTM to predict crop yield. Computers, 12(1), 10. https://doi.org/10.3390/computers12010010
• Colet, F., Lindsey, A. J., & Lindsey, L. E. (2023). Soybean planting date and seeding rate effect on grain yield and profitability. Agronomy Journal. https://doi.org/10.1002/agj2.21434
• Hamed, M. K., Vogel, M. M., Patricola, C. M., & Seneviratne, S. I. (2021). Impacts of compound hot–dry extremes on US soybean yields. Earth System Dynamics, 12(4), 1371–1386. https://doi.org/10.5194/esd-12-1371-2021
• Iizumi, T., & Sakai, T. (2020). The global dataset of historical yields for major crops 1981–2016. Scientific Data, 7(1), 97. https://doi.org/10.1038/s41597-020-0433-7
• Mitchell, K. E., Lohmann, D., Houser, P. R., Wood, E. F., Schaake, J. C., Robock, A., … & Cosgrove, B. A. (2004). The multi-institution North American Land Data Assimilation System (NLDAS): Utilizing multiple GCIP products and partners in a continental distributed hydrological modeling system. Journal of Geophysical Research: Atmospheres, 109(D7), D07S90. https://doi.org/10.1029/2003JD003823
• Ortiz-Bobea, A., Wang, H., Carrillo, C. M., & Ault, T. R. (2019). Unpacking the climatic drivers of U.S. agricultural yields. Environmental Research Letters, 14(6), 064003. https://doi.org/10.1088/1748-9326/ab1e75
• Sun, J., Di, L., Sun, Z., Shen, Y., & Lai, Z. (2019). County-level soybean yield prediction using deep CNN-LSTM model. Sensors, 19(20), 4363. https://doi.org/10.3390/s19204363

Online Tutorials & Blog Posts:

• Vann, R., & Stokes, D. J. (2024, January). How does soybean planting date impact plant height and soybean yield? NC State Extension. https://soybeans.ces.ncsu.edu/2024/01/how-does-soybean-planting-date-impact-plant-height-and-soybean-yield/
• AgWeb. (2024, February). Stack the odds for soybeans this spring. https://www.agweb.com/news/crops/soybeans/stack-odds-soybeans-spring
• XGBoost. (n.d.). XGBoost hyperparameter tuning guide. XGBoost Documentation. https://xgboost.readthedocs.io/en/stable/tutorials/param_tuning.html
• XGBoost. (n.d.). XGBoost parameters. XGBoost Documentation. https://xgboost.readthedocs.io/en/stable/parameter.html
• Amazon Web Services. (n.d.). Tune an XGBoost model in Amazon SageMaker. AWS Documentation. https://docs.aws.amazon.com/sagemaker/latest/dg/xgboost-tuning.html
• Uber. (n.d.). Omphalos: Uber’s point of interest engine for better user experiences. Uber Blog. https://www.uber.com/blog/omphalos/
• Towards Data Science. (2020, April). Grid search vs. random search vs. Bayesian optimization. https://towardsdatascience.com/grid-search-vs-random-search-vs-bayesian-optimization-2e68f57c3c46
• Practical Data Science. (n.d.). How to tune an XGBRegressor model with Optuna. https://practicaldatascience.co.uk/machine-learning/how-to-tune-an-xgbregressor-model-with-optuna

GitHub Repository:

• Kong, W. (2022). U.S. soybean yield prediction [GitHub repository]. https://github.com/wenwenkong/data-science-portfolio/tree/main/US_soybean_yield

Topics

1. Fundamentals

2. Linear Regression

3. Lasso, Ridge, and ElasticNet

4. Logistic Regression

5. Tree-based models (Decision Trees, Random Forests, Bagging, Boosting)

6. XGBoost

7. Regression and Classification Metrics

8. Support Vector Machine

9. Unsupervised Learning

… …