One Brain, Two Gears: Solving the Reasoning Model “Tax”

Reasoning models are built to think through everything. You could ask the simplest math question—like $2+2$—and the model will still go back and start a long chain of thought before giving an answer. The reality is that during training, the model must have seen the result; it knows $2+2$ is $4$. The model needs  to learn how to discriminate. The model needs to learn when to just give the answer and when to use the “Chain of Thought.”
The knowledge is already there. We only need to help the model get the ability to skip the reasoning path and get straight to answer.

The “Router” Trap

One approach is to use a router that decides which question needs reasoning and which does not. But this is a headache. It introduces another layer to maintain and adds latency to every call. Why add more infrastructure when you can fix the behavior in the model itself?

Thinking Mode Fusion: The Control Token

The fix is simpler than most people think: Control Tokens. By giving the model a token—for instance, <no_think>—you can fine-tune it on examples where that token signals a direct answer. No reasoning chain, just the result.
Initially, the model will resist. It’s been “over-trained” to be verbose, so it will still try to prompt those long answers. That’s where our loss function kicks in. By backpropagating that error through the network, we adjust the weights responsible for that verbose behavior.

Reshaping the Path, Not the Knowledge

You aren’t touching the “knows math” weights. You’re touching the “must explain everything” weights. After enough training, the model learns that <no_think> means cut to the chase.
The gradient update is just the standard formula:
$$w_{new} = w_{old} – (\eta \cdot \nabla L)$$
Nothing new here. You’re just creating a new path through the model’s decision space—one that goes straight to the answer instead of taking the long Chain of Thought route.

The Production Payoff

With this approach, you have one model that can do both. When you need deep reasoning, you get it. When you need a quick answer, you add the token and get it.
  • Same weights.
  • Same deployment.
  • No extra infrastructure.
The real work is in the dataset. You have to decide what counts as “simple enough” for a direct answer. If the model is not trained well it might miss reasoning when it has to use chain of thoughts
But for high-volume production use cases with mixed complexity, it’s worth the effort. You can cut your token costs significantly without spinning up a single extra server.

The Prediction Pipeline: How Your LLM Guesses the Next Word

This blog post explores the complete technical journey of a token, revealing how an LLM transforms raw text into a confident guess about what comes next.
I. The First Step: Tokenization and Numerical IDs
The LLM cannot directly process human language; it must first convert text into a numerical format.

  1. Breaking Down the Stream: The input text is broken into smaller units called tokens. While common words are single tokens, rare or complex words (like “unbelievable”) are split into sub-word units (e.g., [“un”, “believe”, “able”]). This efficiency allows the model to handle millions of unique words with a limited, core vocabulary.
  2. The Lookup Table: Every token is assigned a unique, simple numerical ID. The text input is immediately converted into a sequential list of these IDs.

II. From ID to Information: Token Embeddings

A raw ID holds no meaning. The model must look up its Token Embedding.

  1. The Embedding Matrix: The token’s ID points to a row in the massive, pre-trained Embedding Matrix. This row is the Token Embedding (E(Word)), a high-dimensional vector (e.g., a list of 768 numbers).
  2. Semantic Meaning as Features: These numbers are not pre-defined categories (like ‘Is it a city?’). Instead, they are abstract, learned features that capture the word’s statistical usage patterns during training. This vector space encodes:
    • Semantic Similarity: Words used in similar contexts (e.g., “king” and “queen”) have numerically similar vectors.
    • Syntactic Role: Information about whether the word is commonly used as a Noun, Verb, or Adjective is embedded within these feature values.

III. Injecting Order: The Positional Encoding

The Transformer architecture’s efficiency comes from its parallel processing, which makes it fundamentally order-blind. To fix this, we must explicitly inject sequence information.

  1. The Positional Vector (P(pos)): A second vector, the Positional Encoding (P(pos)), is generated for every possible slot (Position 1, Position 2, etc.). This vector acts as a unique “Position ID” for that slot.
  2. The Final Input Vector: The positional vector is mathematically added to the semantic embedding. This creates the final, enriched vector that is fed into the Transformer’s main layers:
    Input = E(word) + P(pos)
  3. The Key Insight: ID, not Grammar: The P(pos)vector only signals where a word is (Position 1, Position 2). The model still has to learn what that position means grammatically.

IV. The Core Mechanism: Self-Attention and Contextualization

The Multi-Head Self-Attention layer is the central computing unit. Here, the model combines the word’s meaning, position, and the context of the entire sentence.

For every token, the attention mechanism calculates an Attention Score between itself and every other token, determining how much influence each neighbor should have.

Example: Learning Subject vs. Object
Consider the ambiguity in word order, which is crucial for grammar:

  1. “Man” (P_1) bites (P_2) “dog” (P_4)
  2. “Dog” (P_1) bites (P_2) “man” (P_4)

When the model processes the verb “bites” (at P_2), a trained Attention Head will learn a rule:

“If I am a verb, I pay maximum attention to the token with the P_1 positional signature, because that token is the grammatical Subject.”

By attending highly to the P_1 token (whether it’s “Man” or “Dog”), the model correctly assigns the Subject role, which is essential for accurate prediction. The output is a Contextualized Vector for each word, reflecting its full meaning and role in the sentence.

V. The Final Act: Prediction

After passing through multiple layers of attention and feed-forward networks, the contextualized vector is ready to make a guess.

  1. The Output Vector: The vector for the last position is passed through a final linear layer, which projects the information into a vector the size of the entire vocabulary (e.g., 50,000 possibilities).
  2. The Softmax Probability: A Softmax function is applied to these scores, generating a probability distribution (all probabilities sum to 1.0). The model states: “There is an X% chance the next token is ‘the’, Y% chance it is ‘a’, etc.”
  3. Generating the Next Token: The model selects the next word ID based on this distribution, converts it back to text, and the entire pipeline restarts to generate the following word.

Happy learning 🙂 .

How Categorical Variables Are Transformed for Better Machine Learning

What is an Embedding?

An embedding is a way to represent categorical variables as dense, continuous vectors in a lower-dimensional space. This helps machine learning models, especially deep learning models, learn relationships between categories more effectively than using simple integer encodings.

Why Not Use Just Category Codes?

When we encode categories as integers using let’s say pandas category data type, the model might assume a false ordinal relationship (e.g., Monday = 1, Tuesday = 2, … which suggests Tuesday > Monday).
Embeddings solve this by mapping each category to a vector of continuous values that capture relationships more naturally.

Example:-

Imagine You’re in a Classroom
Let’s say you have a classroom of students with different favorite subjects:

Student Favorite Subject
Alice Math
Bob Science
Charlie History
David Math
Eva Science

These subjects are categorical because they are names, not numbers.

Step 1: Assigning Category Codes
To use these subjects in a machine learning model, we first convert them to numbers:

Subject Code
Math 0
Science 1
History 2

So now our data looks like this:

Student Favorite Subject (Code)
Alice 0
Bob 1
Charlie 2
David 0

The problem?

  • The model might think Math (0) is smaller than Science (1), which is smaller than History (2).
  • But Math is not “less than” Science! These are just different categories.

Step 2: Converting to Embeddings

Instead of just assigning a number, we represent each subject as a set of numbers (a vector).

🔹 Think of each subject as a point in a space rather than a single number.
🔹 We give each subject multiple numbers (features) that define its “meaning”.

For example, instead of:

  • Math = 0
  • Science = 1
  • History = 2

We assign each subject a list of numbers (a vector), like this:

Subject Embedding (Vector)
Math [0.8, 1.2, -0.5]
Science [1.5, -0.7, 0.9]
History [-0.2, 0.3, 1.1]

Now:

  • Each subject is represented in a continuous space, not just a single number.
  • The model can learn relationships between subjects (e.g., maybe Math and Science are closer than Math and History).

Understanding This with a Real-World Analogy

Think about GPS coordinates 🌍.

Imagine we want to find the locations of three cities:

  • New York
  • Los Angeles
  • Chicago

We don’t just say:

  • New York = 0, Los Angeles = 1, Chicago = 2 ❌ (this would be useless)

Instead, each city has GPS coordinates:

  • New York → [40.7, -74.0]
  • Los Angeles → [34.0, -118.2]
  • Chicago → [41.8, -87.6]

This tells us how far apart the cities are and their relationships in real-world space.
Similarly, embeddings place categories in a “meaningful” space where the model can learn relationships.

Why Are Embeddings Useful?

1️⃣ They prevent false relationships (e.g., Math isn’t “less than” Science).
2️⃣ They help the model find similarities (e.g., Math and Science might be closer than Math and History).
3️⃣ They allow the model to learn meaningful patterns instead of just memorizing numbers.

Final Analogy: Movie Recommendations 🍿

Imagine a movie recommendation system like Netflix:

  • Instead of saying “Comedy = 0, Action = 1, Drama = 2”,
  • It learns embedding vectors like:
    • Comedy = [0.8, 1.5, -0.6]
    • Action = [1.2, -0.4, 1.1]
    • Drama = [-0.3, 0.9, 0.7]

Now, if you like Comedy, the system can recommend other movies that are “closer” in the embedding space.

  • Embeddings turn categories into a set of numbers that capture their relationships.
  • Instead of using simple numbers (0,1,2), embeddings give each category multiple numbers (like GPS coordinates).
  • This helps models understand similarities and patterns better.

Linear regression using PyTorch

If you have a variable Y that is linearly dependent on X, the relationship can be expressed using the polynomial function:

Y=2X+1

In this blog, we’ll explore how to use a simple neural network in PyTorch to predict future values of Y based on X. Our goal is to train the network to learn the polynomial function, even when the data contains noise.

Objective:

The neural network will learn a function of the form:
Y=wX+b
where w (weight) and b (bias) are parameters the model needs to optimize during training. The goal is to make w close to 2 and b close to 1, aligning with the original function.

Additionally, we introduce an error term e in the data to simulate real-world imperfections and test how well the model handles noise. The function becomes:

Y=2X+1+e

Generating Sample Data:
We create sample data using PyTorch, introducing noise through a random error term e:

import torch
import matplotlib.pyplot as plt

X = torch.linspace(1, 50, 50).reshape(-1, 1)  # Input values
torch.manual_seed(75)  # For reproducibility
e = torch.randint(-6, 9, (50, 1), dtype=torch.float)  # Random error
Y = 2 * X + 1 + e

We can visualize this data using a scatter plot:

plt.scatter(X.numpy(), Y.numpy())
plt.show()


Defining the Model:

Using PyTorch, we define a simple neural network with a single dense (linear) layer:

import torch.nn as nn

class Model(nn.Module):
    def __init__(self, input_features, output_features):
        super().__init__()
        self.linear = nn.Linear(input_features, output_features)
        
    def predict(self, x):
        return self.linear(x)

Next, we initialize the model and examine its starting weight (w) and bias (b):

model = Model(1, 1)
print("Initial weight:", model.linear.weight.item())
print("Initial bias:", model.linear.bias.item())

At this point, the model is untrained, and the initial w and b values are random.

Evaluating the Untrained Model:
Using the initial values of w and b, we plot the model’s predictions against the data:

initial_weight = model.linear.weight.item()
initial_bias = model.linear.bias.item()

x1 = np.array([X.min().item(), X.max().item()])
y1_predicted = initial_weight * x1 + initial_bias

plt.scatter(X.numpy(), Y.numpy())
plt.plot(x1, y1_predicted, 'r')
plt.show()


The plot clearly shows that the untrained model fails to predict the polynomial function
𝑌=2𝑋+1

Training the Model:

To train the model, we:

  • Define a loss function to measure the difference between predictions and actual values. We use Mean Squared Error (MSE).
  • Use Stochastic Gradient Descent (SGD) as the optimizer, with a learning rate of 0.005.
  • Perform multiple epochs (iterations) of training to adjust the model’s parameters.
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.005)

epochs = 50
losses = []

for epoch in range(epochs):
    y_pred = model.forward(X)  # Predictions
    loss = criterion(y_pred, Y)  # Calculate loss
    losses.append(loss.item())

    optimizer.zero_grad()  # Reset gradients
    loss.backward()  # Backpropagation
    optimizer.step()  # Update weights and bias

Results:
After training, the model’s weight and bias are:

trained_weight = model.linear.weight.item()
trained_bias = model.linear.bias.item()

print("Trained weight:", trained_weight)
print("Trained bias:", trained_bias)

The trained values of w and b will be close to 2 and 1, respectively, allowing the model to accurately predict the polynomial function:

𝑌=2𝑋+1

Key Takeaways:

  • The neural network successfully learned the underlying polynomial function despite noisy data.
  • This example demonstrates the process of training a simple model, including defining the architecture, selecting a loss function, and optimizing the parameters.

Tracing with LightStep and Nginx

Tracing your webservice/microservice helps in knowing the bottleneck of the component that is taking the most time to respond to. It also helps in pinpointing the location/log/function where the error occurs and what causes the poor performance.

In this blog post, we are trying to cover how we can leverage opentracing with Nginx while using Lighstep vendor for tracing our webapp.

OpenTracing Vendor-neutral APIs and instrumentation for distributed tracing.

Lightstep provides unlimited cardinality, dynamic service maps, and immediate (shockingly accurate) root cause correlation across traces, metrics, and logs anywhere in your system
Lightstep is just another vendor that provides connector for configuring distributed tracing based on opentracing standards.

We need to configure Nginx to use nginx-opentracing module and provide a vendor tracer , we are going to use LightStep vendor.

This GitGist should help you get started very easily.

Install.txt describes the steps to be followed for downloading and configuring the required lib.

Note these two lines –

load_module modules/ngx_http_opentracing_module.so;
Instructs Nginx to load the opentracing module

opentracing_load_tracer /usr/local/lib/liblightstep_tracer_plugin.so /etc/nginx/lightstep-config.json;
The config describes Lightstep vendor lib location and the vendor lib config

Note- In Step3 when you run the command
strings /lib64/libstdc++.so.6 | grep GLIBCXX
make sure you have available
GLIBCXX_3.4.25
GLIBCXX_3.4.26

Once you have configured all the above things you need to restart the Nginx and you should be able to see traces for all the requests your web app is serving.

This would look like