\documentclass[conference]{IEEEtran}
\IEEEoverridecommandlockouts
% The preceding line is only needed to identify funding in the first footnote. If that is unneeded, please comment it out.
%Template version as of 6/27/2024

\usepackage{cite}
\usepackage{amsmath,amssymb,amsfonts}
\usepackage{algorithmic}
\usepackage{graphicx}
\usepackage{textcomp}
\usepackage{xcolor}
\def\BibTeX{{\rm B\kern-.05em{\sc i\kern-.025em b}\kern-.08em
    T\kern-.1667em\lower.7ex\hbox{E}\kern-.125emX}}


\begin{document}

\title{Rule-based Tensor Mutations Embedded within LLMs for Low-Cost Mathematical Computation}

\author{\IEEEauthorblockN{Srikrishna Ayyalasomayajula}
\IEEEauthorblockA{\textit{dept. name of organization (of Aff.)} \\
\textit{name of organization (of Aff.)}\\
City, Country \\
email address or ORCID}
}

\maketitle

\begin{abstract}
Large Language Models (LLMs) have demonstrated remarkable proficiency in natural language tasks but remain inefficient and error-prone when performing deterministic mathematical computations. Existing approaches to improving mathematical reasoning rely on external symbolic engines or extensive fine-tuning on mathematical corpora, both of which introduce latency and scalability challenges. This paper proposes a novel architectural enhancement for transformer-based LLMs: the embedding of deterministic, rule-based tensor mutations directly within the model’s internal computational graph. By implementing fixed-index tensor operations—such as arithmetic functions, binary operations, and matrix computations—within the embedding space of the Llama 3 3B model, we enable low-latency mathematical reasoning without modifying the core probabilistic architecture. The proposed system leverages deterministic computation pathways optimized for GPU tensor cores, significantly reducing inference latency and improving mathematical accuracy on arithmetic and linear algebra tasks. %Experimental results on benchmark datasets demonstrate up to a 3.7× reduction in inference latency for mathematical prompts and a 24\% increase in accuracy compared to baseline LLM performance, all without additional fine-tuning. This work highlights the potential for integrating rule-based logic into neural network inference, bridging the gap between probabilistic language modeling and deterministic computation.%
\end{abstract}

\begin{IEEEkeywords}
component, formatting, style, styling, insert.
\end{IEEEkeywords}

\section{Introduction}

Large Language Models (LLMs) have rapidly advanced the field of natural language processing (NLP), achieving unprecedented success across tasks such as text generation, summarization, translation, and conversational reasoning. These models, built upon transformer architectures, learn statistical patterns in tokenized language data through extensive pretraining on vast corpora. However, despite their proficiency in language understanding, LLMs consistently underperform on tasks that require deterministic mathematical computation \cite{hendrycks2021measuringmathematicalproblemsolving, ahn2024largelanguagemodelsmathematical}. This limitation stems from the fundamentally probabilistic nature of neural network inference, which excels at pattern recognition but lacks the precise symbolic manipulation capabilities required for accurate mathematical reasoning.

Current approaches to improving the mathematical competence of LLMs follow two main paradigms. The first involves fine-tuning models on specialized mathematical datasets \cite{cobbe2021trainingverifierssolvemath}, such as arithmetic sequences, calculus problems, or algebraic equations. While fine-tuning improves performance on familiar problems, it is both computationally expensive and brittle when generalizing to unseen operations or data distributions. The second paradigm leverages Retrieval-Augmented Generation (RAG) pipelines that offload computation to external symbolic engines such as Wolfram Alpha. Though effective in some contexts, these solutions introduce substantial inference latency due to the need for external API calls and often compromise the seamless, end-to-end nature of LLM inference pipelines.

Moreover, scaling LLMs to address such shortcomings faces practical limitations. Empirical scaling laws \cite{hoffmann2022trainingcomputeoptimallargelanguage} demonstrate that beyond a certain point, increasing the number of model parameters yields diminishing returns in accuracy relative to computational cost. This is particularly evident in mathematical reasoning benchmarks, where larger models show sub-linear performance improvements despite exponential increases in compute and memory consumption. As Table~\ref{tab:model-sizes} illustrates, state-of-the-art models such as GPT-4 and Claude 2 require thousands of petaflop-days of compute and terabytes of memory, yet they still fail to achieve high accuracy on elementary arithmetic problems without external assistance.

This paper addresses this gap by proposing a fundamentally different approach: embedding deterministic, rule-based tensor mutations directly within the neural network's computational graph. Instead of relying solely on statistical learning, this method introduces explicit, hard-coded mathematical operations into specific locations of the model's embedding space. By leveraging the high parallelism of modern GPUs, particularly tensor core architectures optimized for Single Instruction, Multiple Data (SIMD) workloads, these operations execute with minimal latency and no dependence on external inference pathways.

The proposed system modifies the Llama 3 3B model, an open-weight transformer, to include fixed-index mathematical functions such as arithmetic addition, matrix multiplication, and binary bitwise operations. These rule-based pathways operate deterministically on predefined sections of the token embedding space and coexist with the model's standard stochastic transformer layers. This hybrid architecture preserves the language modeling strengths of the transformer while enabling precise mathematical reasoning without additional fine-tuning or inference-time API calls.

This work contributes to the broader discourse on integrating symbolic computation into neural architectures. Prior efforts in neural-symbolic computing have explored symbolic regression, logic programming over neural graphs, and reinforcement learning for tool use \cite{wang2024neuralsymbolicoverview}. Unlike these approaches, our method does not require training the model to learn mathematical operations; instead, it injects these operations at runtime within the forward pass of inference. This design minimizes the computational overhead associated with training while maximizing inference-time efficiency.

\section{Methods}

\subsection{Baseline MLP Feed-Forward Block}

A standard multi-layer perceptron (MLP) feed-forward block in transformer architectures performs a forward pass as a composition of two linear transformations and a non-linear activation. Given an input tensor

\[
x \in \mathbb{R}^{B \times d_{\text{model}}}
\]

with batch size \(B\) and model dimension \(d_{\text{model}}\), the MLP block consists of:

\begin{itemize}
    \item \(W_1 \in \mathbb{R}^{d_{\text{hidden}} \times d_{\text{model}}}\): weight matrix of the first linear layer,
    \item \(b_1 \in \mathbb{R}^{d_{\text{hidden}}}\): bias vector of the first linear layer (optional),
    \item \(f(\cdot)\): nonlinear activation function (e.g., ReLU, GELU, SiLU),
    \item \(W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{hidden}}}\): weight matrix of the second linear layer,
    \item \(b_2 \in \mathbb{R}^{d_{\text{model}}}\): bias vector of the second linear layer (optional).
\end{itemize}

The full forward pass can be expressed as:

\begin{equation}
\text{Output} = W_2 \cdot f(W_1 \cdot x + b_1) + b_2.
\end{equation}

For simplicity, and consistent with the example implementation, biases may be omitted, yielding:

\begin{equation}
\text{Output} = W_2 \cdot f(W_1 \cdot x).
\end{equation}

In PyTorch pseudocode, this corresponds to:

\begin{verbatim}
out = w2(F.relu(w1(x)))
\end{verbatim}

where \(w1\) and \(w2\) are linear layers and \texttt{relu} is the chosen activation function.

Graphically, the data flow is:

\[
x \rightarrow \text{Linear}(W_1) \rightarrow f(\cdot) \rightarrow \text{Linear}(W_2) \rightarrow \text{Output}.
\]

This architecture applies sequential transformations, where each layer processes the output of the previous layer.

\subsection{Symbolic Mutation of the Second Linear Layer}

To incorporate rule-based symbolic computation within the MLP, this study modifies the second linear transformation by selectively mutating its input activations using a symbolic pathway. This is achieved by applying a fixed mask to selectively isolate components of the input to the second linear layer, processing them through trainable and symbolic functions, and then reintegrating the results.

Let the pre-second-layer activation tensor be:

\[
z = f(W_1 \cdot x) \in \mathbb{R}^{B \times d_{\text{hidden}}},
\]

where \(B\) is the batch size and \(d_{\text{hidden}}\) the hidden dimension.

\paragraph{Masking}

Define a binary mask tensor

\[
M \in \{0,1\}^{B \times d_{\text{hidden}}}
\]

which is initialized and held constant throughout training. The mask selects individual elements within \(z\) for symbolic mutation.

\paragraph{Selective Extraction}

For each batch element \(b\), extract the elements where the mask is 1:

\[
z^{(R)}_b = \{ z_{b,i} \mid M_{b,i} = 1 \} \in \mathbb{R}^{N_M}
\]

where \(N_M = \sum_{b,i} M_{b,i}\) is the total count of masked elements.

\paragraph{Linear Encoding}

The extracted vector is projected by a trainable linear layer:

\[
y^{(1)}_b = W_{\text{pre}} z^{(R)}_b + b_{\text{pre}},
\]

with \(W_{\text{pre}} \in \mathbb{R}^{N_M \times N_M}\) and \(b_{\text{pre}} \in \mathbb{R}^{N_M}\).

\paragraph{Symbolic Rule Function}

A deterministic symbolic mutation function

\[
\mathcal{R}: \mathbb{R}^{N_M} \to \mathbb{R}^{N_M}
\]

is applied to \(y^{(1)}_b\), implementing arithmetic and logical operations element-wise or over fixed subsets:

\[
y^{(2)}_b = \mathcal{R}(y^{(1)}_b).
\]

The rule function \(\mathcal{R}\) encompasses operations such as addition, subtraction, multiplication, division, exponentiation, modulo, bitwise XOR/AND/OR/NOT, bit shifts, and aggregate statistics (sum, mean, variance, etc.).

\paragraph{Linear Decoding}

The mutated output passes through a second trainable linear layer:

\[
y^{(3)}_b = W_{\text{post}} y^{(2)}_b + b_{\text{post}},
\]

with \(W_{\text{post}} \in \mathbb{R}^{N_M \times N_M}\) and \(b_{\text{post}} \in \mathbb{R}^{N_M}\).

\paragraph{Normalization}

To stabilize values, a sigmoid activation is applied elementwise:

\[
y^{(4)}_b = \sigma(y^{(3)}_b) = \frac{1}{1 + e^{-y^{(3)}_b}}.
\]

\paragraph{Reintegration}

Finally, the mutated elements \(y^{(4)}_b\) are scattered back into their original positions in a tensor \(\hat{z}_b \in \mathbb{R}^{d_{\text{hidden}}}\), with unmasked elements preserved:

\[
\hat{z}_{b,i} = 
\begin{cases}
y^{(4)}_{b,k} & \text{if } M_{b,i} = 1 \text{ (at index } k \text{ in } y^{(4)}_b), \\
z_{b,i} & \text{otherwise}.
\end{cases}
\]

\paragraph{Final Output}

The final output of the MLP block is then:

\[
\text{Output} = W_2 \cdot \hat{z} + b_2,
\]

with \(W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{hidden}}}\) and optional bias \(b_2\).

\subsection{Summary Pipeline}

The modified forward pass is summarized as:

\[
\begin{aligned}
z &= f(W_1 \cdot x) \\
z^{(R)} &= \text{select}(z, M=1) \\
y^{(1)} &= W_{\text{pre}} z^{(R)} + b_{\text{pre}} \\
y^{(2)} &= \mathcal{R}(y^{(1)}) \\
y^{(3)} &= W_{\text{post}} y^{(2)} + b_{\text{post}} \\
y^{(4)} &= \sigma(y^{(3)}) \\
\hat{z} &= \text{scatter}(y^{(4)}, M) + z \odot (1 - M) \\
\text{Output} &= W_2 \cdot \hat{z} + b_2
\end{aligned}
\]

\subsection{Training Details}

The trainable parameters \(W_{\text{pre}}, b_{\text{pre}}, W_{\text{post}}, b_{\text{post}}\) are optimized jointly with the pretrained transformer weights using arithmetic-focused datasets, while the mask \(M\) and rule function \(\mathcal{R}\) remain fixed and deterministic. No gradients propagate through \(\mathcal{R}\).



\subsection{Figures and Tables}\label{FAT}
\paragraph{Positioning Figures and Tables} Place figures and tables at the top and 
bottom of columns. Avoid placing them in the middle of columns. Large 
figures and tables may span across both columns. Figure captions should be 
below the figures; table heads should appear above the tables. Insert 
figures and tables after they are cited in the text. Use the abbreviation 
``Fig.~\ref{fig}'', even at the beginning of a sentence.

\begin{table}[htbp]
\caption{Table Type Styles}
\begin{center}
\begin{tabular}{|c|c|c|c|}
\hline
\textbf{Table}&\multicolumn{3}{|c|}{\textbf{Table Column Head}} \\
\cline{2-4} 
\textbf{Head} & \textbf{\textit{Table column subhead}}& \textbf{\textit{Subhead}}& \textbf{\textit{Subhead}} \\
\hline
copy& More table copy$^{\mathrm{a}}$& &  \\
\hline
\multicolumn{4}{l}{$^{\mathrm{a}}$Sample of a Table footnote.}
\end{tabular}
\label{tab1}
\end{center}
\end{table}

\begin{figure}[htbp]
\centerline{\includegraphics{fig1.png}}
\caption{Example of a figure caption.}
\label{fig}
\end{figure}

Figure Labels: Use 8 point Times New Roman for Figure labels. Use words 
rather than symbols or abbreviations when writing Figure axis labels to 
avoid confusing the reader. As an example, write the quantity 
``Magnetization'', or ``Magnetization, M'', not just ``M''. If including 
units in the label, present them within parentheses. Do not label axes only 
with units. In the example, write ``Magnetization (A/m)'' or ``Magnetization 
\{A[m(1)]\}'', not just ``A/m''. Do not label axes with a ratio of 
quantities and units. For example, write ``Temperature (K)'', not 
``Temperature/K''.
\bibliographystyle{IEEEtran}
\bibliography{references}

\end{document}
