TL;DR

We reduced our Next.js Docker image by 91% (2.1GB to 180MB) using multi-stage builds, Alpine base, standalone mode, and layer caching. Results: 8x faster deploys, 80%+ lower registry costs, pod startup dropped from 90s to 12s.

How We Reduced Our Next.js Docker Image from 2.1GB to 180MB

We deploy our Next.js frontend to Kubernetes multiple times per day. A few months ago, each deploy was painful. Our Docker image had ballooned to 2.1GB. Pull times exceeded 6 minutes. Registry costs were climbing. Security scans flagged 47 vulnerabilities from packages we didn’t even need.

We got it down to 180MB. Here’s the complete step-by-step journey.

This article originally appeared on the Zingaro AI Engineering Blog. Click here to see how we are building the future of human-like voice agents.

How We Reduced Our Next.js Docker Image from 2.1GB to 180MB

We deploy our Next.js frontend to Kubernetes multiple times per day. A few months ago, each deploy was painful. Our Docker image had ballooned to 2.1GB. Pull times exceeded 6 minutes. Registry costs were climbing. Security scans flagged 47 vulnerabilities from packages we didn’t even need.

We got it down to 180MB. Here’s the complete step-by-step journey.

The Starting Point: What Went Wrong

Our original Dockerfile looked like every tutorial on the internet:

FROM node:20
WORKDIR /app
COPY . .
RUN npm install
RUN npm run build
EXPOSE 3000
CMD ["npm", "start"]

Simple. Elegant. And absolutely terrible for production.

This creates a 2GB+ disaster:

ProblemImpactFull Node.js Debian image~1GB base before your codeCOPY . . before installBusts cache on every file changenpm install includes devDepsHundreds of MB of build toolsNo .dockerignoreCopies .git, node_modules, testsBuild artifacts left in image.next/cache, source mapsRoot userSecurity vulnerability

Result: 2.1GB image, 47 security vulnerabilities, 6+ minute pulls.

We couldn’t keep shipping like this.

Step 1: Switch to Alpine Base Image

The single biggest win. One line change that cuts 860MB.

Alpine Linux is a security-focused, lightweight distribution. The Node.js Alpine image is ~140MB vs ~1GB for Debian-based images.

# Before: ~1GB base
FROM node:20# After: ~140MB base
FROM node:24-alpine

Alpine uses musl instead of glibc, which can cause issues with some native modules. Add compatibility layer to be safe:

FROM node:24-alpine
RUN apk add --no-cache libc6-compat

Impact: -860MB immediately.

Step 2: Multi-Stage Builds

This is the architecture that makes everything else possible. We split the build into three distinct stages:

+------------------+     +------------------+     +------------------+
|     Stage 1      |     |     Stage 2      |     |     Stage 3      |
|      deps        | --> |     builder      | --> |     runner       |
|                  |     |                  |     |                  |
|  Install prod    |     |  Install all     |     |  Copy only       |
|  dependencies    |     |  deps + build    |     |  what's needed   |
+------------------+     +------------------+     +------------------+
      ~200MB                   ~900MB                   ~180MB

Each stage starts fresh. Only the final stage goes to production. All the build cruft stays behind.

Stage 1: Dependencies

FROM node:24-alpine AS deps
RUN apk add --no-cache libc6-compat
RUN corepack enable && corepack prepare pnpm@10.22.0 --activate
WORKDIR /app
COPY package.json pnpm-lock.yaml ./
RUN pnpm install --prod --frozen-lockfile

Key points:

• --prod flag: Only production dependencies
• --frozen-lockfile: Fails if lockfile is outdated (reproducible builds)
• Copy only package.json and lockfile first for layer caching

Stage 2: Builder

FROM node:24-alpine AS builder
RUN corepack enable && corepack prepare pnpm@10.22.0 --activate
WORKDIR /app
COPY --from=deps /app/node_modules ./node_modules
COPY . .
COPY .env.build ./.env
RUN pnpm install --frozen-lockfile
ENV NODE_ENV=production
RUN pnpm build

This stage:

• Copies prod deps from stage 1 (layer cache hit)
• Installs ALL deps (including dev) for the build
• Runs the actual Next.js build
• Everything here stays in this stage

Stage 3: Runner (Production)

FROM node:24-alpine AS runner
WORKDIR /app
RUN addgroup --system --gid 1001 nodejs && \
    adduser --system --uid 1001 nextjs
ENV NODE_ENV=production
ENV PORT=3000
COPY --from=builder --chown=nextjs:nodejs /app/.next/standalone ./
COPY --from=builder --chown=nextjs:nodejs /app/.next/static ./.next/static
COPY --from=builder --chown=nextjs:nodejs /app/public ./public
USER nextjs
EXPOSE 3000
CMD ["node", "server.js"]

This stage starts completely fresh:

• New Alpine image (no build cruft)
• Non-root user for security
• Only three directories copied from builder

Step 3: Enable Next.js Standalone Mode

This is the secret weapon. In next.config.js:

module.exports = {
  output: 'standalone',
}

Next.js traces your actual imports and bundles only what’s needed into a self-contained server:

.next/standalone/
├── server.js           # Self-contained server
├── node_modules/       # Only traced dependencies (~30MB)
└── [your app code]

Before standalone: Copy 500MB+ of node_modules After standalone: Copy ~30MB of traced dependencies

No full node_modules. No source files. No devDependencies.

Step 4: Optimize Layer Caching

Docker caches layers. If a layer hasn’t changed, Docker reuses it. Order matters.

# BAD: Any file change busts the cache
COPY . .
RUN npm install
# GOOD: Dependencies cached unless package.json changes
COPY package.json pnpm-lock.yaml ./
RUN pnpm install
COPY . .

Dependencies change rarely. Source code changes constantly. Put stable things first.

Build time impact:

• First build: ~4 minutes
• Subsequent builds (no dep changes): ~45 seconds

Step 5: Configure .dockerignore

Don’t copy garbage into the build context:

# Version control
.git
.gitignore
# Dependencies (each stage installs fresh)
node_modules
# Build outputs (we build fresh)
.next
out
build
# Environment files (secrets!)
.env
.env.local
.env.development
.env.production
.env.prod
# Allow only build-time env vars
!.env.build
# Development files
.vscode
.idea
*.log
*.md
tests
__tests__
coverage
# Docker files
Dockerfile
docker-compose.yml

Impact: Build context reduced from 1.2GB to 180MB.

Step 6: Separate Build-Time vs Runtime Environment Variables

This trips up many teams. Next.js has two types of env vars:

TypeWhen EmbeddedStorageNEXT_PUBLIC_*Build time (baked into JS).env.build in DockerEverything elseRuntimeContainer orchestrator secrets

Wrong approach (security risk):

COPY .env.prod ./.env  # Secrets baked into your image!

Correct approach:

# Build stage - only public vars
COPY .env.build ./.env

.env.build (safe to embed):

NEXT_PUBLIC_APP_NAME=MyApp
NEXT_PUBLIC_API_URL=https://api.example.com

.env.prod (never in Docker):

DATABASE_URL=postgres://...
API_SECRET_KEY=sk_...
PAYMENT_API_KEY=...

Runtime secrets injected via orchestrator:

# Kubernetes example
envFrom:
  - secretRef:
      name: app-secrets

Step 7: Use pnpm Instead of npm/yarn

pnpm is faster and more disk-efficient:

Enable via Corepack (built into Node 16+):

RUN corepack enable && corepack prepare pnpm@10.22.0 --activate

Step 8: Run as Non-Root User

Security best practice. If your container is compromised, the attacker has limited permissions:

RUN addgroup --system --gid 1001 nodejs && \
    adduser --system --uid 1001 nextjs
# Copy files with proper ownership
COPY --from=builder --chown=nextjs:nodejs /app/.next/standalone ./
USER nextjs

Step 9: Server External Packages

Some packages are only used on the server (API routes, server components) but get bundled anyway. Tell Next.js to exclude them:

// next.config.js
module.exports = {
  output: 'standalone',
  serverExternalPackages: [
    'firebase-admin',  // ~20MB, server-only SDK
    'nodemailer',      // Email sending
    'pg',              // PostgreSQL client
  ],
}

These packages use native Node.js require() instead of being webpack-bundled. Saves 10-30MB depending on your stack.

Step 10: Optimize Package Imports (Tree Shaking)

Many popular packages use “barrel exports” — one index file that re-exports everything. Without optimization, importing one icon imports the entire library.

// next.config.js
module.exports = {
  experimental: {
    optimizePackageImports: [
      'lucide-react',     // Icon library
      'date-fns',         // Date utilities
      'lodash',           // Utility functions
      'framer-motion',    // Animations
      'recharts',         // Charts
    ],
  },
}

Next.js transforms barrel imports into direct imports, enabling proper tree-shaking.

Before: import { format } from 'date-fns' pulls in entire library After: Only the format function is bundled

Step 11: Audit and Remove Unused Packages

Dead dependencies are dead weight. Find them:

# Check what's actually imported
npx depcheck
# Or manually grep for imports
grep -r "from 'package-name'" src/

We found several packages installed but never used. Removing them saved ~15MB from node_modules before the build even starts.

Step 12: node-prune (Remove Junk from node_modules)

Even after standalone tracing, node_modules contains junk: README files, changelogs, TypeScript definitions, test files, source maps.

Add node-prune to the builder stage:

# After pnpm build
RUN apk add --no-cache curl && \
    curl -sf https://gobinaries.com/tj/node-prune | sh && \
    node-prune .next/standalone/node_modules

node-prune removes:

• *.md files
• *.map source maps
• __tests__ directories
• *.d.ts type definitions
• Documentation folders

Impact: 10–20MB savings.

Step 13: Distroless Base Image (Advanced)

For maximum security and minimum size, replace Alpine with Google’s Distroless:

# Instead of Alpine (~50MB)
FROM node:24-alpine AS runner
# Use Distroless (~30MB)
FROM gcr.io/distroless/nodejs22-debian12 AS runner

Distroless contains only the Node.js runtime. No shell, no package manager, no utilities.

FROM gcr.io/distroless/nodejs22-debian12 AS runner
WORKDIR /app
ENV NODE_ENV=production
ENV PORT=3000
# Distroless runs as nonroot (uid 65532) by default
COPY --from=builder --chown=65532:65532 /app/.next/standalone ./
COPY --from=builder --chown=65532:65532 /app/.next/static ./.next/static
COPY --from=builder --chown=65532:65532 /app/public ./public
USER nonroot
EXPOSE 3000
CMD ["server.js"]

Trade-off: You can’t kubectl exec into the container for debugging. Rely on logs instead.

Impact: 20MB smaller, near-zero attack surface.

Step 14: Optimize Static Assets

Your public/ folder adds directly to image size. Optimize images:

# Convert PNGs to WebP (70-90% smaller)
for f in public/*.png; do
  cwebp -q 85 "$f" -o "${f%.png}.webp"
  rm "$f"
done

Update your code to use .webp instead of .png. Next.js Image component handles this automatically with formats: ['image/webp'].

We reduced our public folder from 11MB to 3.7MB.

The Complete Dockerfile

With all optimizations applied (including node-prune and distroless):

# =============================================================================
# Stage 1: Dependencies
# =============================================================================
FROM node:24-alpine AS deps
RUN apk add --no-cache libc6-compat
RUN corepack enable && corepack prepare pnpm@10.22.0 --activate
WORKDIR /app
COPY package.json pnpm-lock.yaml ./
RUN pnpm install --prod --frozen-lockfile
# =============================================================================
# Stage 2: Builder
# =============================================================================
FROM node:24-alpine AS builder
RUN corepack enable && corepack prepare pnpm@10.22.0 --activate
WORKDIR /app
COPY --from=deps /app/node_modules ./node_modules
COPY . .
COPY .env.build ./.env
RUN pnpm install --frozen-lockfile
ENV NODE_ENV=production
RUN pnpm build
# Remove junk from traced node_modules
RUN apk add --no-cache curl && \
    curl -sf https://gobinaries.com/tj/node-prune | sh && \
    node-prune .next/standalone/node_modules

# =============================================================================
# Stage 3: Runner (Production) - Distroless
# =============================================================================
FROM gcr.io/distroless/nodejs22-debian12 AS runner
WORKDIR /app
ENV NODE_ENV=production
ENV PORT=3000
# Distroless runs as nonroot (uid 65532) by default
COPY --from=builder --chown=65532:65532 /app/.next/standalone ./
COPY --from=builder --chown=65532:65532 /app/.next/static ./.next/static
COPY --from=builder --chown=65532:65532 /app/public ./public
USER nonroot
EXPOSE 3000
CMD ["server.js"]

The Results

Quick Verification Commands

# Build the image
docker build -t myapp:latest .
# Check the size
docker images myapp:latest
# Inspect layers
docker history myapp:latest
# Run locally
docker run -p 3000:3000 myapp:latest
# Security scan
docker scout cves myapp:latest

Optimization Checklist

Before shipping to production:

Docker Basics:

• Multi-stage build (deps → builder → runner)
• .dockerignore excludes node_modules, .next, .git, .env files
• Only .env.build (public vars) copied to image
• Runtime secrets via orchestrator, not baked in
• Non-root user in production stage
• --frozen-lockfile for reproducible builds
• Layer caching optimized (copy deps first)

Next.js Configuration:

• output: 'standalone' in next.config.js
• serverExternalPackages for server-only deps
• optimizePackageImports for barrel-export packages

Advanced Optimizations:

• Alpine or Distroless base image
• node-prune to remove junk from node_modules
• Audit and remove unused packages
• Convert images to WebP format
• pnpm for faster, smaller installs

Final Thoughts

Every megabyte in your Docker image costs real money and time. In Kubernetes, it compounds:

• Larger images = slower pulls = slower scaling
• Slower scaling = missed SLAs during traffic spikes
• More storage = higher registry and node costs

We went from dreading deploys to shipping 15+ times per day. The 2-hour investment in optimizing our Dockerfile pays dividends on every single deploy.

These techniques aren’t Next.js-specific. Multi-stage builds, Alpine base images, and layer caching apply to any Docker workflow. Start with the biggest wins (Alpine + multi-stage), then iterate.

Your CI/CD pipeline will thank you.

Ready to Build the Future of Voice?

We are building the most responsive and natural-sounding voice agents at Zingaro AI. Whether you are a developer looking to experiment or a business ready to scale, we have the infrastructure you need.

🚀 Start building today: Check out Zingaro.ai and sign up for an account.

💼 For Enterprise Inquiries: If you need custom solutions or high-volume processing, reach out to our team directly at sales@zingaro.ai.

Imagine creating dynamic forms in seconds, tailored to your specific needs, with just a simple text prompt. No more manually coding each field, setting up validation, or worrying about the layout. With the power of OpenAI and Python, you can now generate complex forms dynamically, straight from a text prompt. In this blog post, I’ll guide you through how to build this powerful tool that transforms natural language into fully functional HTML forms.

This blog post is inspired by my recent project, which you can check out here.

The Building Blocks

Here’s the stack and libraries we’re using:

• OpenAI: The AI model that parses your text prompt and returns structured data for form generation.
• Pydantic: For defining the schema and ensuring the data returned by the AI fits the expected format.
• FastHTML: FastHTML is a cutting-edge web framework that allows us to create forms entirely in Python, without the need for JavaScript or HTML.

Step-by-Step Guide to Building Your Dynamic Form Generator

Step 1: Define the Form Structure

We start by defining the data structure for our forms using Pydantic models. This structure includes various field types (text, number, date, etc.) and allows for the specification of options for fields like dropdowns or checkboxes.

from pydantic import BaseModel, Field
from enum import Enum
from typing import Optional, List

class TypeEnum(str, Enum):
    text = "text"
    number = "number"
    date = "date"
    radio = "radio"
    checkbox = "checkbox"
    select = "select"
    textarea = "textarea"

class Options(BaseModel):
    label: str = Field(description="Label for the option")
    value: str = Field(description="Value for the option")

class FormField(BaseModel):
    label: str = Field(description="Field label")
    type: TypeEnum = Field(description="Field type")
    name: str = Field(description="Unique field name")
    required: bool = Field(description="Is the field required?")
    placeholder: Optional[str] = Field(description="Placeholder text for input fields")
    options: Optional[List[Options]] = Field(description="Options for radio, checkbox, and select fields")

class DynamicForm(BaseModel):
    title: str = Field(description="Form title")
    fields: List[FormField] = Field(description="List of form fields")

As you can see from the Pydantic structure above, it covers nearly all use cases. If you think I’ve missed anything, feel free to let me know in the comments!

Step 2: Create the Dynamic Form

The structured output from OpenAI is in the above Pydantic format, which we use to generate the actual HTML form. We iterate through the fields, creating the appropriate HTML elements for each field type. The exciting part is that with FastHTML, you can write everything in pure Python to render the form — no JavaScript or HTML is required. In fact, the entire code is contained within a single file.

def create_dynamic_form(form_data: DynamicForm | None):
    if dynamic_form_data is None:
        return None

    fields = []

    for field in form_data.fields:
        field_type = field.type
        if field_type in ["text", "number", "date"]:
            fields.append(
                Div(
                    Label(field.label, _for=field.name),
                    Input(
                        type=field_type,
                        name=field.name,
                        placeholder=field.placeholder,
                        required=field.required,
                    ),
                    style="margin-bottom: 15px;",
                )
            )
        elif field_type == "select":
            options = [
                Option(value=opt.value, label=opt.label) for opt in field.options
            ]
            fields.append(
                Div(
                    Label(field.label, _for=field.label),
                    Select(*options, name=field.name, required=field.required),
                    style="margin-bottom: 15px;",
                )
            )
        elif field_type == "checkbox":
            checkboxes = [
                Div(
                    Input(
                        type="checkbox",
                        name=field.name,
                        value=opt.value,
                        id=f"{field.name}_{opt.value}",
                    ),
                    Label(opt.label, _for=f"{field.name}_{opt.value}"),
                )
                for opt in field.options
            ]
            fields.append(
                Div(Label(field.label), Div(*checkboxes), style="margin-bottom: 15px;")
            )
        elif field_type == "radio":
            radios = [
                Div(
                    Input(
                        type="radio",
                        name=field.name,
                        value=opt.value,
                        id=f"{field.name}_{opt.value}",
                        required=field.required,
                    ),
                    Label(opt.label, _for=f"{field.name}_{opt.label}"),
                )
                for opt in field.options
            ]
            fields.append(
                Div(Label(field.label), Div(*radios), style="margin-bottom: 15px;")
            )
        elif field_type == "textarea":
            fields.append(
                Div(
                    Label(field.label, _for=field.name),
                    Textarea(
                        name=field.name,
                        placeholder=field.placeholder,
                        required=field.required,
                    ),
                    style="margin-bottom: 15px;",
                )
            )

    form = Form(
        *fields,
        Button("Submit", type="submit", style="margin-top: 20px;"),
        method="post",
        action="/submit",
    )

    return Container(H1(form_data.title), form)

Don’t worry if the code above seems confusing — it’s just the FastHTML approach to rendering HTML. You can learn more about FastHTML by visiting their website at fastht.ml.

Step 3: Generate Form Data with OpenAI

Next, we use OpenAI to generate the form structure based on a text prompt. The AI reads the prompt, understands the form requirements, and outputs the corresponding form structure in the above Pydantic format.

from openai import OpenAI
from dotenv import load_dotenv
from textwrap import dedent

load_dotenv()
client = OpenAI()

def get_form_response(prompt):
    system_prompt = """
        You are an AI that generates HTML forms based on user prompts.
        Your task is to take a description of a form and output the corresponding form fields and structure.
    """
    completion = client.beta.chat.completions.parse(
        model="gpt-4o-2024-08-06",
        messages=[
            {"role": "system", "content": dedent(system_prompt)},
            {"role": "user", "content": prompt},
        ],
        response_format=DynamicForm,
    )
    return completion.choices[0].message

The key point to note here is the response format, which we’re defining using the Above Pydantic class. This explicitly instructs OpenAI to generate output in this specific format.

Step 4: Deploy and Run the Web Application

Finally, we wrap everything in a web application using FastHTML. The app takes user input, passes it to OpenAI to generate the form, and then dynamically displays the generated form.

@app.get("/")
def get():
    prompt_area = Form(H2("Enter your prompt here to create a dynamic form"), Textarea(name="prompt"), Button("Create Form", hx_post="/update-prompt"))
    form_area = Div(id="form_area")
    return Container(prompt_area, form_area)

@app.post("/update-prompt")
async def update_prompt(request):
    form_data = await request.form()
    form_prompt = form_data.get("prompt", "")
    if form_prompt:
        dynamic_form_data = get_form_response(form_prompt)
        return Div(create_dynamic_form(dynamic_form_data), id="form_area")
    return Div("Error generating form. Please try again.", id="form_area")

Bringing It All Together

This tool can be invaluable in any scenario where forms are frequently customized or generated on the fly, from admin panels to user surveys. The beauty lies in the seamless integration of AI to automate the creation process, freeing developers from repetitive tasks and allowing them to focus on more creative aspects of their work.

The full code is available here, and I encourage you to experiment with it and adapt it to your projects.

If you have any questions or want to share your experiences, feel free to drop a comment below or connect with me on LinkedIn. Happy coding!

If you’re already working with large language models (LLMs), chances are you’re using natural language generation (NLG) in some way. Maybe you’re using Hugging Face models, either with the easy-to-use Hugging Face pipeline or the model.generate function. But today, we’re going to look closer and really get into how to create text [tokens] right from the beginning.

Large language models are autoregressive, meaning they generate each word based on the previous ones and continue this process until certain criteria are met. This article aims to explore and code these criteria from scratch. While we’ll still be utilizing Hugging Face models, we’ll craft our own generate method. This approach will help us gain insights into what occurs behind the scenes.

I’ve also shared a fully functional example on Colab. You can find the link below.

Why you need to understand this ?

1. deeper understanding of how the large language models work
2. you can customize the generation process
3. can lead to new generation methods
4. also helps you how to prompt these models better

If you’ve ever built a large language model (LLM) from scratch, you probably know what the final result is like. If you haven’t built an LLM model from the beginning, I’ll start with the basic ideas first.

In any deep learning model, there are two main components: the input and the output. For large language models (LLMs), the input typically consists of a batch size and a sequence of words, often referred to as the ‘prompt.’ The length of this prompt sequence can vary based on our choice — we decide how many words to use as input. To show how this works in practice, I’ll use the Zephyr 7B model from Hugging Face as an example for demonstration.

Tokenization:-

Let’s import the model from huggingface !

# Define the model name to be used from Hugging Face
model_name = "HuggingFaceH4/zephyr-7b-beta"

# Initialize tokenizer for converting text to model-understandable format
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Load the causal language model with auto device mapping and custom settings
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map='auto',  # Automatically assign model to best available device
    trust_remote_code=True,  # Allow use of custom layers from Hugging Face Hub
    torch_dtype=torch.float16,  # Set model parameters to float16 for efficiency
)

# Set the model to evaluation mode
model.eval()

# Define the prompt text to generate language from
prompt = "write a story about a cat"

Let’s tokenize the prompt. Tokenizing is a process where we break down text into smaller parts. In simple terms, this can mean dividing the text into individual words or sometimes splitting a word into several parts. For our purposes, let’s think of it as simply splitting the text into words. Now, let’s see how the text looks once it’s tokenized.

tokens = tokenizer.tokenize(prompt)
print(tokens)

# ['▁write', '▁a', '▁story', '▁about', '▁a', '▁cat']

Our prompt consists of six words. Let’s proceed to encode it. Encoding refers to assigning a unique number to each word.

input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
print(input_ids)
print(input_ids.shape)

# tensor([[   1, 3324,  264, 2838,  684,  264, 5255]], device='cuda:0')
# torch.Size([1, 7])

Our prompt contains just six words, yet the encoded prompt shows a length of seven. Why is that? It’s because we’re adding a start token at the beginning of the sequence.

Let’s use the model’s generate function, a method provided by Hugging Face, to produce the output. This will be the only instance where we utilize the predefined generate function. Moving forward, we’ll switch to our custom generate function that we will develop.

output = model.generate(input_ids, max_length=50)
generated_text = tokenizer.decode(output[0])
print(generated_text)

# <s> write a story about a cat who solves crimes\n\nTitle: The Purrfect Detective\n\nChapter One: The Case of the Missing Mouse\n\nIn the quiet neighborhood of Maple Street, a cat named Oscar was'

In the output, we notice that the model repeats our prompt for the first seven words. Starting with the eighth word, it begins to generate new content. This is where the model starts creating its own text, moving beyond just repeating what we gave it.

Now we understand how the model produces output: it uses our prompt, applies internal attention mechanisms, and generates the next token, continuing this process. Next, we’ll skip the Hugging Face generate method. Instead, we’ll directly input our data into the model for a forward pass and analyze the resulting output.

Forward Pass:-

In a forward pass for large language models (LLMs), the model receives input text, such as a sentence or phrase, processes it systematically, and predicts the next word. This forward pass needs to be repeated multiple times until we either encounter the end-of-sentence token (</s>) or reach the predefined maximum length for the text sequence.

output = model(input_ids)
print(output.keys())

# odict_keys(['logits', 'past_key_values'])

what we just did ?

We feed the model our start token along with six words, making a total of seven tokens including start token, to predict the eighth token. The result we get is called logits. Now, let’s examine the shape of these logits to understand them better.

The shape of logits will give us insight into the dimensions of the output, such as how many tokens were processed and the range of possible word predictions for each token.

output.logits.shape

# torch.Size([1, 7, 32000])

The logits shape `torch.Size([1, 7, 32000])` means:

• 1 sequence is being processed or you can think of a batch size.
• 7 tokens in the sequence, including your input.
• 32,000 possible words or tokens that the model predicts for each of these 7 tokens. which is also called vocabulary size

The logits shape torch.Size([1, 7, 32000]) indicates one sequence with seven tokens and 32,000 word predictions for each token. Next, we’ll explore how to pick the next token from these 32,000 options.

Decoding Strategies

1. Greedy Decoding
2. Beam Search
3. Top-K Sampling
4. Nucleus Sampling

Let’s discuss all these strategies and implement them from scratch using pytorch

1. Greedy Decoding

Greedy decoding is probably the easiest decoding strategy.At each time step, when we input our sequence into the model, we receive a probability distribution across the entire vocabulary(30K tokens in our case). Here, we consistently select the token with the highest probability at every step.

Let’s code this in pytorch !

def greedy_decoding(input_ids, max_tokens=300):
    with torch.inference_mode():
        for _ in range(max_tokens):
            # input_ids shape: [1, current_sequence_length]
            outputs = model(input_ids)

            # outputs.logits shape: [1, current_sequence_length, vocab_size]
            next_token_logits = outputs.logits[:, -1, :]
            # next_token_logits shape: [1, vocab_size]

            next_token = torch.argmax(next_token_logits, dim=-1)
            # next_token shape: [1] (the most probable next token ID)
            
            # stop generation if the model produces the end of sentence </s> token 
            if next_token == tokenizer.eos_token_id:
                break

            # rearrange(next_token, 'c -> 1 c'): changes shape to [1, 1] for concatenation
            input_ids = torch.cat([input_ids, rearrange(next_token, 'c -> 1 c')], dim=-1)
            # input_ids shape after concatenation: [1, current_sequence_length + 1]

        generated_text = tokenizer.decode(input_ids[0])
        # input_ids[0] shape for decoding: [current_sequence_length]

    return generated_text

# Call the function with initial input_ids
print(greedy_decoding(input_ids))

# <s> write a story about a cat 
# who solves crimes\n\nTitle: The Purrfect Detective\n\nChapter One: The Case of the Missing Mouse\n\nIn the quiet neighborhood of Maple Street, a cat named Oscar was known for his sharp senses and quick reflexes. Oscar was not an ordinary cat; he was a detective. He had a reputation for solving mysteries that left the other cats in awe.\n\nOne fine day, Oscar received a call from his friend, the mouse, Timmy. Timmy was missing, and Oscar was the only one who could help find him. Oscar knew that this was his chance to'

The line “next_token = torch.argmax(next_token_logits, dim=-1)” means that we’re selecting the token with highest probability.

Then, in the next line, we take this chosen token and add it to our input sentence to keep generating text. So, we’re building our text step by step, always choosing the word that seems the most probable according to the model.

For a straightforward explanation, please note that this code is designed to work with a batch size of one. If you intend to use it for batch decoding, some modifications may be necessary.

Why Greedy Decoding is Not best choice ?

1. Not the Best Choice: Greedy decoding often doesn’t produce the best overall text. It only picks the most likely next word each time, which might not make the best sentence when you look at it as a whole.
2. Repetitive and Predictable: It tends to create text that’s not very varied or interesting. Since it always goes for the most obvious next word, the results can be quite generic.
3. Can’t Fix Mistakes: Once greedy decoding makes a mistake, it can’t go back and fix it. This means that one wrong word choice can lead to more errors down the line.
4. Limited Exploration: Greedy decoding doesn’t consider many different options or paths. It just sticks to a narrow set of choices, which means it might miss out on better or more suitable ways to complete the text.

In summary, greedy decoding is a quick and straightforward method, but it’s limited in its ability to create diverse, high-quality text, especially for complex or longer content.

2. Beam Search:

The fundamental idea of beam search is exploring multiple paths instead of just single one.: Unlike greedy decoding, which only keeps the single best path (the most probable next word) at each step, beam search keeps track of the k most probable paths at each step. This means at every step in the sequence, it doesn’t just consider the single most likely next word, but the k most likely next words.

The diagram illustrates how beam search operates with a beam size of 2. Let’s walk through the steps depicted in the diagram, which show the progression of the search through different timesteps:

Timestep 1:

• Start: The process begins with the <START> token.
• First Word Choices: Two paths are considered based on the first word choice, “the” and “a”. These are the two most probable words following the <START> token.

Timestep 2:

• Second Word Choices: For each of the initial two paths, the model then considers the next set of words. For the path starting with “the”, it explores “poor” and “people”. For the path starting with “a”, it explores “poor” and “person”.
• Pruning to Beam Size: After considering the probabilities of these expansions, the algorithm keeps the two most probable paths. Let’s assume these are “the poor” and “a poor”.

Timestep 3:

• Third Word Choices: From “the poor”, the model explores further expansions and similarly from “a person”. For “the poor”, the next words considered are “are” and “don’t”. For “a poor”, the model looks at “person” and “but”. However, “a poor” path would have a lower probability and be pruned.

Timestep 4:

• Fourth Word Choices: Continuing from the remaining high-probability paths, “the poor are” leads to “always” and “not”, while “the poor don’t” leads to “have” and “take”. The top two paths with the highest cumulative probability are kept, and the rest are pruned.

Timestep 5:

• Fifth Word Choices: Now, the model expands “the poor are not” and “the poor don’t have”. For “not”, it considers “in” and “with”, and for “have”, it considers “any” and “enough”.

At each step, only the two most probable paths are kept (as per the beam size of 2), and all other possibilities are discarded. This process continues until a stopping criterion is met, which can be the model generating an end-of-sequence token or reaching a maximum sequence length. The final output would be the sequence with the highest overall probability from the last set of paths.

First, let’s implement beam search in PyTorch, and then I’ll delve into the specifics.

from einops import rearrange
import torch.nn.functional as F

def beam_search(input_ids, max_tokens=100, beam_size=2):
    # Initialize the scores for each beam with zeros. Shape: [beam_size]
    beam_scores = torch.zeros(beam_size).to(device)
    
    # Duplicate the initial sequence for each beam. Shape: [beam_size, seq_length]
    beam_sequences = input_ids.clone()
    
    # Create a boolean mask to keep track of active beams. Shape: [beam_size]
    active_beams = torch.ones(beam_size, dtype=torch.bool)
   
    
    for step in range(max_tokens):
        # Generate model outputs for the current sequences. 
        # The model is expected to handle batched input, hence the shape of beam_sequences is [beam_size, current_seq_length].
        outputs = model(beam_sequences)
        
        # Extract the last logits from the output to get the probabilities for the next token. Shape: [beam_size, vocab_size]
        logits = outputs.logits[:, -1, :]
        probs = F.softmax(logits, dim=-1)
        
        # Calculate the score for each beam and token by flattening the probabilities and selecting the top ones.
        # The flattened shape is [beam_size * vocab_size], from which we select the top beam_size scores.
        top_scores, top_indices = torch.topk(probs.flatten(), k=beam_size, sorted=False)
       
        # Map flat indices back to beam and token indices.
        # beam_indices is the index in the beam, shape: [beam_size]
        # token_indices is the index of the token in the vocabulary, shape: [beam_size]
        beam_indices = top_indices // probs.shape[-1]
        token_indices = top_indices % probs.shape[-1]
        
        # Update the sequences with the new tokens at the end. Shape after update: [beam_size, current_seq_length + 1]
        # This concatenates the best token for each beam to the end of the sequences.
        beam_sequences = torch.cat([
            beam_sequences[beam_indices],
            token_indices.unsqueeze(-1)
        ], dim=-1)

        # Update the beam scores with the top scores. Shape: [beam_size]
        beam_scores = top_scores
        
        # Check for the end-of-sequence tokens and update the active beams.
        # If a beam produces an EOS token, it is marked as inactive.
        active_beams = ~(token_indices == tokenizer.eos_token_id)
        
        # If all beams are inactive, exit the loop.
        if not active_beams.any():
            print("no active beams")
            break
    
    # Select the beam with the highest score as the best sequence. Shape: [best_seq_length]
    best_beam = beam_scores.argmax()
    best_sequence = beam_sequences[best_beam]
    
    # Decode the best sequence to generate the final text.
    generated_text = tokenizer.decode(best_sequence)
    
    return generated_text

If the above code seems unclear, here’s a simplified version of the beam search algorithm.

Starting Point

Let’s begin with two sequences of token IDs:

input_ids = torch.tensor([[1, 2, 3, 0], [1, 2, 3, 2]])

The beam size determines how many sequences we’ll keep track of at each step:

beam_size = 2

In this case, we’ll explore the two most promising continuations for our sequences.

Vocabulary and Probabilities

We have a small example vocabulary and corresponding probabilities for the next token, given by a model:

vocab = [0, 1, 2, 3]
probs = torch.tensor([
    [0.2, 0.3, 0.5, 0.1], [0.2, 0.2, 0.2, 0.2]
])

Here, each sublist in probs represents the model's estimated probability of each vocabulary token being the next in the sequence. For instance, the model believes there's a 50% chance that the next token after [1, 2, 3, 0] is 2.

Selecting the Top Candidates

Next, we flatten the probability tensor and pick the top beam_size probabilities:

top_scores, top_indices = torch.topk(probs.flatten(), k=beam_size, sorted=False)

The flatten function turns our 2D probabilities array into a 1D array, and torch.topk retrieves the highest values and their indices.

Mapping Back to Beams and Tokens

We then determine which beam each top score came from and the token associated with that score:

beam_indices = top_indices // len(vocab)
token_indices = top_indices % len(vocab)

The division and modulus operations help us map the flat index back to a specific beam and token ID.

Updating the Sequences

Finally, we update our sequences with the new tokens:

beam_sequences = torch.cat([
    input_ids[beam_indices],
    token_indices.unsqueeze(-1)
], dim=-1)

Let’s put all the code together so you can run it as a whole in one cell to understand how the above algorithm works.

input_ids = torch.tensor([[1, 2, 3, 0], [1,2, 3, 2]])
beam_size = 2
vocab = [0, 1 , 2, 3]
probs = torch.tensor([
    [0.2, 0.3, 0.5, 0.1], [0.2, 0.2, 0.2, 0.2]
])
top_scores, top_indices = torch.topk(probs.flatten(), k=beam_size, sorted=False)
print(top_scores, top_indices) # tensor([0.5000, 0.3000]) tensor([2, 1])
beam_indices = top_indices // len(vocab)
token_indices = top_indices % len(vocab)
print(beam_indices,token_indices) # tensor([0, 0]) tensor([2, 1])
beam_sequences = torch.cat([
    input_ids[beam_indices],
    token_indices.unsqueeze(-1)
], dim=-1)
print(beam_sequences) # tensor([[1, 2, 3, 0, 2], [1, 2, 3, 0, 1]])

After running this piece of code, beam_sequences contains the updated sequences with the new tokens added at the end. This process would be repeated for each new token we want to generate, gradually building up more complete sequences.

3. Top-K Sampling

The concept of Top-K sampling involves sampling only from the top K tokens in the probability distribution. Let’s say we set K to 50; this means that when we predict the next token, we sort the entire 32,000-word vocabulary according to the softmax scores of our logits. Then, we retain only the top 50 options from the vocabulary and remove the rest of the items. Subsequently, we randomly select a token from the previously filtered list.

def top_k_sampling(input_ids, max_tokens=100, top_k=50, temperature=1.0):
   for _ in range(max_tokens):
        # Temporarily disables gradient calculation to improve performance and reduce memory usage
        with torch.inference_mode():
            # Obtain the model's outputs based on the input IDs
            outputs = model(input_ids)
            
            # Extract logits of the next token
            next_token_logits = outputs.logits[:, -1, :]
            
            # Select the top K tokens from the probability distribution
            top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k)
            
            # Apply softmax to convert logits to probabilities, with optional temperature scaling
            top_k_probs = F.softmax(top_k_logits / temperature, dim=-1)
            
            # Sample from the top K tokens to determine the next token
            next_token_index = torch.multinomial(top_k_probs, num_samples=1)
            
            # Map the sampled token back to its original index in the logits tensor
            next_token = top_k_indices.gather(-1, next_token_index)
            
            # Concatenate the new token to the sequence of input IDs
            input_ids = torch.cat([input_ids, next_token], dim=-1)
        
    # Decode the tensor of input IDs to a string of text
    generated_text = tokenizer.decode(input_ids[0])
    
    return generated_text

What is the problem with Top-K sampling ?

1. Dynamic Probability Distributions: When we use Top-K sampling, the list of possible next words (tokens) and their chances of being picked can change a lot each time we choose a new word. Sometimes, many words have a similar chance of being chosen, and other times, one or a few words are much more likely than the others.
2. When the List is Flat (Many Words Have Similar Chances): Imagine we have a big basket of fruits, and you’re almost equally likely to pick any fruit because you like them all pretty much the same. If you’re only allowed to choose from the top 5 fruits in the basket, you might miss out on a lot of other fruits you also like just as much. This is like when the probability distribution is flat, and limiting ourselves to the top K options (like the top 5 fruits) means we ignore many other good choices.
3. When the List is Peaky (One or a Few Words are Much More Likely): Imagine you’re at a buffet with one dish you absolutely love and many others that are just okay. If you’re allowed to choose from a very large selection, including all the okay dishes alongside your favorite, the presence of so many less appealing options could lead you to pick something less satisfying by chance. In a peaky distribution, where one or a few words are much more likely to be the next correct choice, setting a high K allows many less probable (less appealing) options to be considered. This increases the chances of selecting a word that isn’t the best fit, simply because it’s among the many options available, diluting the likelihood of choosing the most appropriate or expected word based on the context.

This issue can be addressed by the subsequent technique we’ll explore, known as Top-P (Nucleus Sampling).

4. Top-P (Nucleus) Sampling:

Instead of selecting a fixed number of top tokens (like the top 50 in Top-K sampling), Top-P sampling chooses tokens based on their cumulative probability. It looks at the list of all possible next tokens, ranked by their probability, and adds them up from the most likely to the least likely until the total probability reaches a threshold P (like 0.9 or 90%).

• Cumulative Probability Mass: Imagine you have a pie chart representing the probability of each token being the next word, with the whole pie being 100%. Top-P sampling slices this pie from the most significant piece (most probable token) onwards, stopping when the combined slice reaches the P threshold. This cumulative slice represents the “mass” where the probabilities are concentrated.
• Adapting to the Distribution (Varies K): The key advantage of Top-P sampling is its flexibility. The number of tokens selected (equivalent to K in Top-K sampling) automatically adjusts based on how spread out or concentrated the probabilities are. If a few tokens hold most of the probability mass (a “peaky” distribution), only those will be considered, resulting in a smaller “K”. If the probability is more evenly distributed (“flatter”), more tokens are included to reach the P threshold, resulting in a larger “K”.
• Uniformity of Pt: “Pt” refers to the probability distribution at a given time step (t) for selecting the next token. Top-P sampling dynamically adjusts the number of options considered by directly responding to the shape of this distribution. When Pt is uniform (flatter), indicating many plausible next tokens, more tokens are included to reach the cumulative probability threshold P. When Pt is less uniform (peakier), fewer tokens are needed.

def top_p_sampling(input_ids, max_tokens=100, top_p=0.95):
    # Temporarily disables gradient calculation to improve performance and reduce memory usage
    with torch.inference_mode():
        for _ in range(max_tokens):
                # Obtain the model's outputs based on the input IDs
                outputs = model(input_ids)
                next_token_logits = outputs.logits[:, -1, :]
                
                # Sort the logits in descending order and apply softmax to get probabilities
                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                sorted_probabilities = F.softmax(sorted_logits, dim=-1) # Apply temperature
                
                # Calculate cumulative probabilities
                cumulative_probs = torch.cumsum(sorted_probabilities, dim=-1)
                
                # Identify and remove tokens with cumulative probability above the threshold (top_p)
                # Ensuring the first token is always selected by setting the first position to False
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[..., 0] = False  # Keep the most probable token
                
                # Get the actual indices to remove from the original logits tensor
                indices_to_remove = sorted_indices[sorted_indices_to_remove]
                
                # Set the logits of removed tokens to -infinity to exclude them from sampling
                next_token_logits.scatter_(-1, indices_to_remove[None, :], float('-inf'))
                
                # Re-calculate probabilities after filtering and sample from this distribution
                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                
                # Append the sampled token to the input IDs for the next iteration
                input_ids = torch.cat([input_ids, next_token], dim=-1)
        
        # Decode the tensor of input IDs to a string of text
        generated_text = tokenizer.decode(input_ids[0])
    
    return generated_text

Decoding in natural language generation (NLG) remains a tough puzzle, but even small changes to how we choose words can make a big difference. These tweaks help us make the generated text more natural and are key to the latest breakthroughs in making computers talk more like humans.

Colab Link:-

https://colab.research.google.com/drive/13L_PE7QBMCkIWFD9OzthD-oOGgGyz027?usp=sharing

Conclusion

I hope you found this article insightful and valuable. If you enjoyed the content and wish to stay updated with future posts, please consider following me on LinkedIn. Your support motivates me to continue sharing knowledge. Thank you for reading!

If you’ve delved into the realm of deep learning, you’re likely aware of the impact that transformer architectures have had on the field of artificial intelligence. These architectures stand at the core of numerous groundbreaking advancements in AI. In this Article, we will embark on an in-depth exploration, guiding you through the process of building Vision Transformers from the ground up.

This article is the first in a four-part series. The next one will show how to build ‘LoRa’ from scratch, for the Vision Transformer we are building here.

I’ve also shared a fully functional example on Colab. You can find the link below.

Lets Define a configuration for the Vision Transformer model using a data class

@dataclass
class ModelArgs:
    dim: int = 256          # Dimension of the model embeddings
    hidden_dim: int = 512   # Dimension of the hidden layers
    n_heads: int = 8        # Number of attention heads
    n_layers: int = 6       # Number of layers in the transformer
    patch_size: int = 4     # Size of the patches (typically square)
    n_channels: int = 3     # Number of input channels (e.g., 3 for RGB images)
    n_patches: int = 64     # Number of patches in the input
    n_classes: int = 10     # Number of target classes
    dropout: float = 0.2    # Dropout rate for regularization

MultiHead Attention

The MultiHeadAttention module in the provided code is an implementation of the multi-head self-attention mechanism, which stands as a fundamental component in transformer architectures. This self-attention mechanism empowers the model to weigh input elements differently, offering the capability to focus more intently on certain parts of the input when generating the output.

First, let’s code the multi-head attention block. Afterward, I’ll break down the key components in detail

class MultiHeadAttention(nn.Module):
    def __init__(self, args:ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        
        # Linear projections for Q, K, and V
        self.wq = nn.Linear(self.dim, self.n_heads*self.head_dim, bias=False)
        self.wk = nn.Linear(self.dim, self.n_heads*self.head_dim, bias=False)
        self.wv = nn.Linear(self.dim, self.n_heads*self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads*self.head_dim, self.dim, bias=False)
    
    def forward(self, x):
        b, seq_len, dim = x.shape  # b: batch size, seq_len: sequence length
        
        assert dim == self.dim, "dim is not matching"
        
        q = self.wq(x)  # [b, seq_len, n_heads*head_dim]
        k = self.wk(x)  # [b, seq_len, n_heads*head_dim]
        v = self.wv(x)  # [b, seq_len, n_heads*head_dim]
        
        # Reshape the tensors for multi-head operations
        q = q.contiguous().view(b, seq_len, self.n_heads, self.head_dim)  # [b, seq_len, n_heads, head_dim]
        k = k.contiguous().view(b, seq_len, self.n_heads, self.head_dim)  # [b, seq_len, n_heads, head_dim]
        v = v.contiguous().view(b, seq_len, self.n_heads, self.head_dim)  # [b, seq_len, n_heads, head_dim]
        
        # Transpose to bring the head dimension to the front
        q = q.transpose(1, 2)  # [b, n_heads, seq_len, head_dim]
        k = k.transpose(1, 2)  # [b, n_heads, seq_len, head_dim]
        v = v.transpose(1, 2)  # [b, n_heads, seq_len, head_dim]
        
        # Compute attention scores and apply softmax
        attn = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)  # [b, n_heads, seq_len, seq_len]
        attn_scores = F.softmax(attn, dim=-1)  # [b, n_heads, seq_len, seq_len]
        
        # Compute the attended features
        out = torch.matmul(attn_scores, v)  # [b, n_heads, seq_len, head_dim]
        out = out.contiguous().view(b, seq_len, -1)  # [b, seq_len, n_heads*head_dim]
        
        return self.wo(out)  # [b, seq_len, dim]

The MultiHeadAttention module performing the following operations:

1. Linear transformations of the input tensor into “query” (Q), “key” (K), and “value” (V) representations.

q = self.wq(x)
k = self.wk(x)
v = self.wv(x)

2. Dividing these tensors into multiple “heads”.

q = q.contiguous().view(b, seq_len, self.n_heads, self.head_dim)
k = k.contiguous().view(b, seq_len, self.n_heads, self.head_dim)
v = v.contiguous().view(b, seq_len, self.n_heads, self.head_dim)

3. Computing attention scores via the dot product of Q and K.

    attn = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)

4. Applying softmax to these scores to procure attention weights.

    attn_scores = F.softmax(attn, dim=-1)

5. Multiplying the attention weights with the V tensor, yielding the attended features.

    out = torch.matmul(attn_scores, v)

6. Aggregating results across all heads and projecting to provide the concluding output.

    out = out.contiguous().view(b, seq_len, -1)
    return self.wo(out)

Attention Block

The AttentionBlock module encapsulates a typical block found within transformer architectures. It primarily consists of two significant components: a multi-head self-attention mechanism and a feed-forward neural network (FFN). Additionally, layer normalization and skip connections (residual connections) are employed to facilitate better learning and gradient flow.

let’s code the multi-head attention block.

class AttentionBlock(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.layer_norm_1 = nn.LayerNorm(args.dim)
        self.attn = MultiHeadAttention(args)
        
        self.layer_norm_2 = nn.LayerNorm(args.dim)
        
        self.ffn = nn.Sequential(
            nn.Linear(args.dim, args.hidden_dim),
            nn.GELU(),
            nn.Dropout(args.dropout),
            nn.Linear(args.hidden_dim, args.dim),
            nn.Dropout(args.dropout)
        )
    
    def forward(self, x):
        x = x + self.attn(self.layer_norm_1(x))
        x = x + self.ffn(self.layer_norm_2(x))
        return x

Let’s delve deeper into its structure:

1. Layer Normalization (Pre-Attention) Before feeding the input x into the multi-head attention mechanism, it's normalized using nn.LayerNorm.

self.layer_norm_1 = nn.LayerNorm(args.dim)
x = self.layer_norm_1(x)

2. Multi-Head Self-Attention This component allows the model to focus on different parts of the input sequence when generating its output.

self.attn = MultiHeadAttention(args)
x = x + self.attn(x)

3. Layer Normalization (Pre-Feed-Forward Network) Just like before the multi-head attention mechanism, the output is normalized again using LayerNorm before feeding it into the FFN.

self.layer_norm_2 = nn.LayerNorm(args.dim)
x = self.layer_norm_2(x)

4. Feed-Forward Neural Network (FFN) The FFN consists of two linear layers separated by a GELU activation function. There’s also dropout applied for regularization.

self.ffn = nn.Sequential(
    nn.Linear(args.dim, args.hidden_dim),
    nn.GELU(),
    nn.Dropout(args.dropout),
    nn.Linear(args.hidden_dim, args.dim),
    nn.Dropout(args.dropout)
)
x = x + self.ffn(x)

5. Residual Connections Residual or skip connections are vital for deep architectures like transformers. They help in preventing the vanishing gradient problem and helps in model convergence. In the code, these are represented by the addition operations where the input is added back to the output of both the attention mechanism and the FFN.

x = x + self.attn(...)
x = x + self.ffn(...)

By sequentially organizing the operations, this block ensures efficient and effective feature transformation, which is essential for the transformer’s performance.

Converting Image into Patches

Before creating our full vistion transformer model, we need to create a utility function that transforms images into non-overlapping patches.

def img_to_patch(x, patch_size, flatten_channels=True):
    # x: Input image tensor 
    # B: Batch size, C: Channels, H: Height, W: Width
    B, C, H, W = x.shape  # (B, C, H, W)
    
    # Reshape the image tensor to get non-overlapping patches
    x = x.reshape(B, C, H//patch_size, patch_size, W//patch_size, patch_size)  # (B, C, H/patch_size, patch_size, W/patch_size, patch_size)
    
    # Permute to group the patches and channels
    x = x.permute(0, 2, 4, 1, 3, 5)  # (B, H/patch_size, W/patch_size, C, patch_size, patch_size)
    
    # Flatten the height and width dimensions for patches
    x = x.flatten(1,2)  # (B, (H/patch_size * W/patch_size), C, patch_size, patch_size)
    
    # Option to flatten the channel and spatial dimensions
    if flatten_channels:
        x = x.flatten(2,4)  # (B, (H/patch_size * W/patch_size), (C * patch_size * patch_size))
    
    return x

The img_to_patch function takes an image tensor and converts it into non-overlapping patches of a specified size. This operation is typically used in vision transformers to represent an image as a sequence of flattened patches. The function provides an option to flatten the channels or keep them separate.

Vision Transformer

The VisionTransformer effectively integrates the previously discussed components to construct the final model. It operates as an encoder-only architecture similar to BERT, where all tokens attend to all other tokens. Moreover, we introduce an additional class token (cls_token) to every sequence in the batch, and this will be utilized later for classification, much like how BERT does with its special [CLS] token.

class VisionTransformer(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        # Define the patch size
        self.patch_size = args.patch_size
        
        # Embedding layer to transform flattened patches to desired dimension
        self.input_layer = nn.Linear(args.n_channels * (args.patch_size ** 2), args.dim)

        # Create the attention blocks for the transformer
        attn_blocks = []
        for _ in range(args.n_layers):
            attn_blocks.append(AttentionBlock(args))
        
        # Create the transformer by stacking the attention blocks
        self.transformer = nn.Sequential(*attn_blocks)
        
        # Define the classifier
        self.mlp = nn.Sequential(
            nn.LayerNorm(args.dim),
            nn.Linear(args.dim, args.n_classes)
        )
        
        # Dropout layer for regularization
        self.dropout = nn.Dropout(args.dropout)
        
        # Define the class token (similar to BERT's [CLS] token)
        self.cls_token = nn.Parameter(torch.randn(1, 1, args.dim))
        
        # Positional embeddings to give positional information to transformer
        self.pos_embedding = nn.Parameter(torch.randn(1, 1+args.n_patches, args.dim))
    
    def forward(self, x):
        # Convert image to patches and flatten
        x = img_to_patch(x, self.patch_size)
        b, seq_len, _ = x.shape
        
        # Transform patches using the embedding layer
        x = self.input_layer(x)
        
        # Add the class token to the beginning of each sequence
        cls_token = self.cls_token.repeat(b, 1, 1)
        x = torch.cat([cls_token, x], dim=1)
        
        # Add positional embeddings to the sequence
        x = x + self.pos_embedding[:,:seq_len+1]
        
        # Apply dropout
        x = self.dropout(x)
        
        # Process sequence through the transformer
        x = self.transformer(x)

        # Retrieve the class token's representation (for classification)
        x = x.transpose(0, 1)
        cls = x[0]

        # Classify using the representation of the class token
        out = self.mlp(cls)
        return out

Let’s dive into its key aspects:

1. Patch Embedding Instead of operating on raw pixels, the image is divided into fixed-size patches. Each patch is then linearly transformed (flattened and passed through a linear layer) to a specified dimension (args.dim).

self.patch_size = args.patch_size
self.input_layer = nn.Linear(args.n_channels * (args.patch_size ** 2), args.dim)
x = img_to_patch(x, self.patch_size) 
x = self.input_layer(x)

2. Transformer Blocks A sequence of attention blocks to process the embedded patches. The number of blocks is defined by args.n_layers.

attn_blocks = []
for _ in range(args.n_layers):
    attn_blocks.append(AttentionBlock(args))
self.transformer = nn.Sequential(*attn_blocks)
x = self.transformer(x)

3. CLS Token and Position Embeddings A class token is added to the sequence of embedded patches. This token is later used to obtain the final classification output. Positional embeddings are added to provide the transformer with information about the relative positions of patches.

self.cls_token = nn.Parameter(torch.randn(1, 1, args.dim))
self.pos_embedding = nn.Parameter(torch.randn(1, 1+args.n_patches, args.dim))
cls_token = self.cls_token.repeat(b, 1, 1)
x = torch.cat([cls_token, x], dim=1)
x = x + self.pos_embedding[:,:seq_len+1]

4. Dropout is applied for regularization purposes.

self.dropout = nn.Dropout(args.dropout)
x = self.dropout(x)

5. Classifier The classification head. It uses the class token’s[CLS] representation after it’s been processed by all transformer blocks.

self.mlp = nn.Sequential(
    nn.LayerNorm(args.dim),
    nn.Linear(args.dim, args.n_classes)
)
x = x.transpose(0, 1)
cls = x[0]
out = self.mlp(cls)

The Below code snippet provides a setup for preprocessing and loading the CIFAR10 dataset

# Path to the directory where CIFAR10 data will be stored/downloaded
DATA_DIR = "../data"

# Define the transformation for testing dataset:
# 1. Convert images to tensors.
# 2. Normalize the tensors using the mean and standard deviation of CIFAR10 dataset.
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.49139968, 0.48215841, 0.44653091], [0.24703223, 0.24348513, 0.26158784])
])

# Define the transformation for training dataset:
# 1. Apply random horizontal flip for data augmentation.
# 2. Perform random resizing and cropping of images for data augmentation.
# 3. Convert images to tensors.
# 4. Normalize the tensors using the mean and standard deviation of CIFAR10 dataset.
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomResizedCrop((32, 32), scale=(0.8, 1.0), ratio=(0.9, 1.1)),
    transforms.ToTensor(),
    transforms.Normalize([0.49139968, 0.48215841, 0.44653091], [0.24703223, 0.24348513, 0.26158784])
])

# Load the CIFAR10 training dataset with the defined training transformation.
# The dataset will be downloaded if not present in the DATA_DIR.
train_dataset = CIFAR10(root=DATA_DIR, train=True, transform=train_transform, download=True)

# Load the CIFAR10 testing dataset with the defined testing transformation.
# The dataset will be downloaded if not present in the DATA_DIR.
test_set = CIFAR10(root=DATA_DIR, train=False, transform=test_transform, download=True)

# Split the training dataset into training and validation sets.
# The training set will have 45000 images, and the validation set will have 5000 images.
train_set, val_set = torch.utils.data.random_split(train_dataset, [45000, 5000])

Let’s setup the data loaders for training, validation and test datasets

# Define the batch size for training, validation, and testing.
batch_size = 16

# Define the number of subprocesses to use for data loading.
num_workers = 4

# Create a DataLoader for the training and validation dataset:
# 1. Shuffle the training data for each epoch.
# 2. Drop the last batch if its size is not equal to `batch_size` to maintain consistency.
train_loader = torch.utils.data.DataLoader(dataset=train_set, 
                                           batch_size=batch_size, 
                                           shuffle=True,
                                           num_workers=num_workers, 
                                           drop_last=True)

# Do not drop any data; process all the validation data.
val_loader = torch.utils.data.DataLoader(dataset=val_set, 
                                         batch_size=batch_size, 
                                         shuffle=False,
                                         num_workers=num_workers, 
                                         drop_last=False)

# Create a DataLoader for the testing dataset:
# Do not drop any data; process all the test data.
test_loader = torch.utils.data.DataLoader(dataset=test_set, 
                                          batch_size=batch_size, 
                                          shuffle=False,
                                          num_workers=num_workers, 
                                          drop_last=False)

Let’s configure the model, optimization strategy, and training criterion using cross entropy loss

# Model, Loss and Optimizer
device = "cuda:0" if torch.cuda.is_available() else 0
args = ModelArgs()
model = VisionTransformer(args).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[80, 130], gamma=0.1)

Time to bring our model to life! Let’s train it.

num_epochs = 150  # example value, adjust as needed

for epoch in range(num_epochs):
    
    # Training Phase
    model.train()
    total_loss = 0.0
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device), labels.to(device)
        
        # Zero the parameter gradients
        optimizer.zero_grad()
        
        # Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # Backward pass and optimize
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    avg_train_loss = total_loss / len(train_loader)
    print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_train_loss:.4f}")

    # Validation Phase
    model.eval()
    total_val_loss = 0.0
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            # Forward pass
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            
            total_val_loss += loss.item()

            _, predicted = outputs.max(dim=-1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    avg_val_loss = total_val_loss / len(val_loader)
    val_accuracy = 100 * correct / total
    print(f"Epoch [{epoch+1}/{num_epochs}], Validation Loss: {avg_val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}%")

    # Update the learning rate
    lr_scheduler.step()

print("Training complete!")

Now, let’s evaluate our model’s accuracy on the test set.

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:  
        inputs, labels = inputs.to(device), labels.to(device)

        # Forward pass
        outputs = model(inputs)
        _, predicted = outputs.max(dim=-1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

test_accuracy = 100 * correct / total
print(f"Test Accuracy: {test_accuracy:.2f}%")

Test Accuracy: 75.07%

colab link:-

https://colab.research.google.com/github/pashanitw/pytorch_examples/blob/main/transformer/VisionTransformer.ipynb

Summary:-

Our model’s accuracy is lower than a basic CNN. This is because vision transformers need more data to excel. With the CIFAR-10 dataset, we expect around 75–80% accuracy. But on larger datasets, vision transformers often do better than CNNs.

Part-2 Sneak Peek:-

Coming up in Part 2: I’ll guide you through the hands-on implementation of LoRa (Low Rank Approximation) tailored for the vision transformer we’ve covered. Stay Tuned!

Conclusion:-

If you found value in this, connect or reach out to me on LinkedIn.

Why is Understanding MACs and FLOPs in Neural Networks Important?

In this session, we are going to delve deep into the concepts of MACs (Multiply-Accumulate Operations) and FLOPs (Floating Point Operations) within the context of neural networks. By learning how to calculate these manually using pen and paper, you’ll acquire a foundational understanding of the computational complexity and efficiency of various network structures.

Understanding MACs and FLOPs is not just an academic exercise; it is a critical component in optimizing neural networks for performance and efficiency. It helps in designing models that are both computationally efficient and effective, ultimately saving time and resources during the training and inference phases.

here is a fully functioning example in a colab notebook

Resource Efficiency:

Understanding FLOPs helps in estimating the computational cost of a neural network. By optimizing the number of FLOPs, one can potentially reduce the time taken to train or run a neural network.

Memory Efficiency:

MAC operations often dictate the memory usage of the network since they are directly related to the number of parameters and activations in the network. Reducing MACs can help in making the network memory efficient.

Energy Consumption:

Power Efficiency:

Both FLOPs and MAC operations contribute to the power consumption of the hardware on which the neural network is running. By optimizing these metrics, one can potentially reduce the energy requirements of running the network, which is particularly important in mobile and embedded devices.

Model Optimization:

Pruning and Quantization:

Understanding FLOPs and MACs can assist in optimizing a neural network through techniques like pruning (removing unnecessary connections) and quantization (reducing the precision of weights and activations), which aim to reduce computational and memory costs.

Performance Benchmarking:

Comparison between Models:

FLOPs and MACs provide a means to compare different models in terms of their computational complexity, which can be a criterion for selecting models for specific applications.

Hardware Benchmarking:

These metrics can also be used to benchmark the performance of different hardware platforms in running neural networks.

Deployment on Edge Devices:

Real-time Applications:

For real-time applications, especially on edge devices with limited computational resources, understanding and optimizing these metrics is critical in ensuring that the network can run within the time constraints of the application.

Battery Life:

In battery-powered devices, reducing the computational cost (and hence energy consumption) of neural networks can help in extending the battery life.

Research and Development:

Designing New Algorithm:

Researchers can use these metrics as guidelines when developing new algorithms or neural network architectures, aiming to improve computational efficiency without sacrificing accuracy.

Step 1: Understand the Definitions

FLOP

A FLOP (Floating Point OPeration) is considered to be either an addition, subtraction, multiplication, or division operation.

MAC

A MAC (Multiply-ACCumulate) operation is essentially a multiplication followed by an addition, i.e., MAC = a * b + c. It counts as two FLOPs (one for multiplication and one for addition).

Step 2: Analyze Each Layer

calculating the number of floating-point operations or multiply-accumulate operations to understand the computational complexity of each layer.

1. Fully Connected Layer (Dense Layer)

Now, we will create a simple neural network with 3 layers and begin counting the operations involved. Here is the formula for calculating the operations in the first linear layer, which is a fully connected (or dense) layer:

For a fully connected layer with I inputs and O outputs, the number of operations are as follows:

• MACs: I × O
• FLOPs: 2 × (I × O) (since each MAC counts as two FLOPs)

class SimpleLinearModel(nn.Module):
    def __init__(self):
        super(SimpleLinearModel,self).__init__()
        self.fc1 = nn.Linear(in_features=10, out_features=20, bias=False)
        self.fc2 = nn.Linear(in_features=20, out_features=15, bias=False)
        self.fc3 = nn.Linear(in_features=15, out_features=1, bias=False)
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        F.relu(x)
        x = self.fc3(x)
        return x

linear_model = SimpleLinearModel().cuda()
sample_data = torch.randn(1, 10).cuda()

step 1: Identifying Layer Parameters For the given model, we have three linear layers defined as:

• fc1: 10 input features, 20 output features
• fc2: 20 input features, 15 output features
• fc3: 15 input features, 1 output feature

Step 2: Calculating FLOPs and MACs Now, calculate MACs and FLOPs for each layer:

Layer fc1:

MACs = 10 × 20 = 200

FLOPs = 2 × MACs = 2 × 200 = 400

Layer fc2:

MACs = 20 × 15 = 300

FLOPs = 2 × MACs = 2 × 300 = 600

Layer fc3:

MACs = 15 × 1 = 15

FLOPs = 2 × MACs = 2 × 15 = 30

Step 3: Summing Up the Results Finally, to find the total number of MACs and FLOPs for a single input passing through the entire network, we sum the results from all layers:

• Total MACs = MACs(fc1) + MACs(fc2) + MACs(fc3) = 200 + 300 + 15 = 515
• Total FLOPs = FLOPs(fc1) + FLOPs(fc2) + FLOPs(fc3) = 400 + 600 + 30 = 1030

Verifying FLOPs and MACs with torchprofile Library

You can use the torchprofile library to verify the FLOPs and MACs calculations for the given neural network model. Here's how to do it:

macs = profile_macs(linear_model, sample_data)
print(macs)

#515

2. Convolutional Neural Networks(CNNs)

Now, let’s determine the MACs (Multiply-Accumulates) and FLOPs (Floating-Point Operations) for a straightforward convolutional model. This calculation can be a bit more involved than our previous example with dense layers, mainly due to factors like stride, padding, and kernel size. However, I’ll break it down to make it easier for our learning purpose.

class SimpleConv(nn.Module):
    def __init__(self):
        super(SimpleConv, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.fc =  nn.Linear(in_features=32*28*28, out_features=10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = x.view(x.shape[0], -1)
        x = self.fc(x)
        return x

x = torch.rand(1, 1, 28, 28).cuda()
conv_model = SimpleConv().cuda()

Important Consideration for Calculating Convolutional Operations

• When calculating operations for convolutional kernels, it’s crucial to remember that the number of channels in the kernel should match the number of channels in the input. For instance, if our input is an RGB image with three color channels, the kernel’s dimensions will be 3x3x3 to account for the input's three channels.
• For the purpose of our demonstration, we’ll maintain a consistent image size throughout the convolutional layers. To achieve this, we’ll set both the padding and stride values to 1.

step 1: Identifying Layer Parameters

For the given model, we have two conv layers and one linear layer defined as:

• conv1: 1 input channels, 16 output channels, kernel size is 3
• conv2: 16 input channels, 32 output channels
• fc: 32*28*28 input features, 1 output feature. because our image is not changed in the convolutional layers

Step 2: Calculating FLOPs and MACs Now, calculate MACs and FLOPs for each layer:

formula is output_image_size * kernel shape * output_channels

Layer conv1:

MACs = 28 * 28 * 3 * 3 * 1 * 16 = 1,12,896

FLOPs = 2 × MACs = 2 × 200 = 2,25,792

Layer conv2:

MACs = 28 × 28 * 3 * 3 * 16 * 32 = 3,612,672

FLOPs = 2 × MACs = 2 × 300 = 600 = 7,225,344

Layer fc:

MACs = 32 * 28 * 28 * 10 = 250,880

FLOPs = 2 × MACs = 2 × 15 = 501,760

Step 3: Summing Up the Results Finally, to find the total number of MACs and FLOPs for a single input passing through the entire network, we sum the results from all layers:

• Total MACs = MACs(conv1) + MACs(conv2) + MACs(fc) = 1,12,896 + 3,612,672 + 250,880 = 39,76,448
• Total FLOPs = FLOPs(fc1) + FLOPs(fc2) + FLOPs(fc3) = 2,25,792 + 7,225,344 + 501,760 = 7,952,896

Verifying Operations with torchprofile Library

macs = profile_macs(conv_model,(x,))
print(macs)

#3976448

3. Self-Attention Block

Having covered MACs for linear and convolutional layers, our next step is to determine the FLOPs (Floating-Point Operations) for a Self-Attention block, a crucial component in large language models. This calculation is essential for understanding the computational complexity of such models. Let’s delve into it.

class SimpleAttentionBlock(nn.Module):
    def __init__(self, embed_size, heads):
        super(SimpleAttentionBlock, self).__init__()
        self.embed_size = embed_size
        self.heads = heads
        self.head_dim = embed_size // heads

        assert (
            self.head_dim * heads == embed_size
        ), "Embedding size needs to be divisible by heads"

        self.values = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.keys = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.queries = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

    def forward(self, values, keys, queries, mask):
        N = queries.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], queries.shape[1]
        print(values.shape)
        values = self.values(values).reshape(N,  self.heads, value_len, self.head_dim)
        keys = self.keys(keys).reshape(N, self.heads, key_len, self.head_dim)
        queries = self.queries(queries).reshape(N,  self.heads, query_len, self.head_dim)


        energy = torch.matmul(queries, keys.transpose(-2, -1))        

        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.nn.functional.softmax(energy, dim=3)
        out = torch.matmul(attention, values).reshape(
            N, query_len, self.heads * self.head_dim
        )

        return self.fc_out(out)

step 1: Identifying Layer Parameters

Linear Transformations

let’s define hyper_params

• batch_size = 1
• seq_len = 10
• embed_size = 256

In the attention block, we have three linear transformations (for queries, keys, and values), and one at the end (fc_out).

• Input Size: [batch_size, seq_len, embed_size]
• Linear transformation matrix: [embed_size, embed_size]
• MACs: batch_size×seq_len×embed_size×embed_size

Query, Key, Value linear transformation:

MACs for Query Transformation = 1 * 10 * 256 * 256 = 6,55,360

MACs for Key Transformation = 1 * 10 * 256 * 256 = 6,55,360

MACS for Value Transformation = 1 * 10 * 256 * 256 = 6,55,360

Energy Calculation Calculation: queries (reshaped) dot keys (reshaped) — a dot product operation.

Macs: batch_size×seq_len×seq_len×heads×head_dim

query and key dot product

MACS = 1 * 10 * 10 * 32 [32 because 256/8 divide by heads] = 25,600

Output from Attention Weights and Values Calculation: attention weights dot values (reshaped) — another dot product operation.

Macs : batch_size×seq_len×seq_len×heads×head_dim

attention and value dot product

Macs = 1 * 10 * 10 * 32 = 25,600

Fully Connected Output (fc_out)

Macs: batch_size×seq_len×heads×head_dim×embed_size

Macs = 1 * 10 * 8 * 32 * 256 = 6,55,360

Step 2: Summing Up the Results

• Total MACs = MACs(conv1) + MACs(conv2) + MACs(fc) = 6,55,360 + 6,55,360 + 6,55,360 + 25,600 + 25,600 + 6,55,360 = 26,72,640
• Total FLOPs = 2 * Total MACs = 53,45,280

Verifying Operations with torchprofile Library

# Create an instance of the model
model = SimpleAttentionBlock(embed_size=256, heads=8).cuda()

# Generate some sample data (batch of 5 sequences, each of length 10, embedding size 256)
values = torch.randn(1, 10, 256).cuda()
keys = torch.randn(1, 10, 256).cuda()
queries = torch.randn(1, 10, 256).cuda()

# No mask for simplicity
mask = None
# Forward pass with the sample data
macs = profile_macs(model, (values, keys, queries, mask))
print(macs)

#2672640

Summary: Scaling MACs and FLOPs for Different Batch Sizes

Throughout our calculations, we’ve primarily considered a batch size of 1. However, it’s important to note that scaling MACs and FLOPs for larger batch sizes is straightforward.

To compute MACs or FLOPs for a batch size greater than one, you can simply multiply the total MACs or FLOPs obtained for batch size 1 by the desired batch size value. This scaling allows you to estimate computational requirements for various batch sizes in your neural network models.

Keep in mind that the results will directly scale linearly with the batch size. For instance, if you have a batch size of 32, you can obtain the MACs or FLOPs by multiplying the values for batch size 1 by 32.

Conclusion

I hope you found this article insightful and valuable. If you enjoyed the content and wish to stay updated with future posts, please consider following me on LinkedIn. Your support motivates me to continue sharing knowledge. Thank you for reading!