| --- |
| license: mit |
| --- |
| |
| --- |
|
|
| language: |
|
|
| * en |
| library_name: pytorch |
| license: mit |
| tags: |
| * quadratic-equations |
| * mathematical-reasoning |
| * interpretability |
| * mechanistic-interpretability |
| * symbolic-reasoning |
| * synthetic-data |
| pipeline_tag: text-generation |
|
|
| --- |
|
|
| # SEAI: Learning Mathematical Representations for Quadratic Equation Solving |
|
|
| **SEAI (Square Equation AI)** is a research neural network designed to solve quadratic equations with real roots while enabling direct analysis of its learned internal representations. |
|
|
| The primary purpose of SEAI is not only numerical prediction. The project investigates whether a neural network trained to predict quadratic roots can develop internal representations correlated with mathematically meaningful quantities involved in the analytical solution. |
|
|
| A central observation of the current experiment is that hidden activations in deeper layers exhibit strong correlations with the normalized quadratic discriminant |
|
|
| $$ |
| p^2-q |
| $$ |
|
|
| and with its square-root transformation, despite neither quantity being explicitly provided as an input feature. |
|
|
| --- |
|
|
| ## Model Summary |
|
|
| SEAI receives a normalized representation of a quadratic equation |
|
|
| $$ |
| ax^2+bx+c=0 |
| $$ |
|
|
| and predicts the two ordered real roots. |
|
|
| The model contains exactly: |
|
|
| **8,040,002 trainable parameters.** |
|
|
| The architecture consists of four custom quadratic layers followed by a linear output projection. |
|
|
| ```text |
| Input (2) |
| ↓ |
| QuadraticLayer: 2 → 1000 |
| ↓ |
| LayerNorm |
| ↓ |
| GELU |
| ↓ |
| QuadraticLayer: 1000 → 2000 |
| ↓ |
| LayerNorm |
| ↓ |
| GELU |
| ↓ |
| QuadraticLayer: 2000 → 2000 |
| ↓ |
| LayerNorm |
| ↓ |
| GELU |
| ↓ |
| QuadraticLayer: 2000 → 1000 |
| ↓ |
| LayerNorm |
| ↓ |
| GELU |
| ↓ |
| Linear: 1000 → 2 |
| ↓ |
| Ordered normalized roots |
| ``` |
|
|
| --- |
|
|
| ## Mathematical Formulation |
|
|
| A quadratic equation is invariant under multiplication of all coefficients by the same non-zero scalar. |
|
|
| We therefore use the scale-invariant quantities |
|
|
| $$ |
| P=\frac{b}{a}, |
| \qquad |
| Q=\frac{c}{a}. |
| $$ |
|
|
| The training distribution generated from roots in |
|
|
| $$ |
| [-50,50] |
| $$ |
|
|
| produces |
|
|
| $$ |
| P\in[-100,100], |
| \qquad |
| Q\in[-2500,2500]. |
| $$ |
|
|
| The model receives the normalized variables |
|
|
| $$ |
| p=\frac{P}{100}, |
| \qquad |
| q=\frac{Q}{2500}. |
| $$ |
|
|
| The target roots are normalized as |
|
|
| $$ |
| x_n=\frac{x}{50}. |
| $$ |
| |
| Under this normalization, the exact solution for the normalized roots is |
| |
| $$ |
| \boxed{ |
| x_{n,1/2}=-p\pm\sqrt{p^2-q} |
| } |
| $$ |
|
|
| with the roots ordered so that |
|
|
| $$ |
| x_{n,1}\leq x_{n,2}. |
| $$ |
|
|
| The classical discriminant is |
|
|
| $$ |
| \Delta=b^2-4ac. |
| $$ |
|
|
| Under the SEAI normalization, |
|
|
| $$ |
| \boxed{ |
| \Delta=10000a^2(p^2-q) |
| } |
| $$ |
|
|
| and therefore |
|
|
| $$ |
| \boxed{ |
| p^2-q=\frac{\Delta}{10000a^2}. |
| } |
| $$ |
|
|
| Thus, \(p^2-q\) is exactly the discriminant normalized by the coefficient scale. |
|
|
| --- |
|
|
| ## Dataset |
|
|
| The dataset is entirely synthetic. |
|
|
| Two real roots are sampled from |
|
|
| $$ |
| x_1,x_2\sim U(-50,50) |
| $$ |
|
|
| and sorted so that |
|
|
| $$ |
| x_1\leq x_2. |
| $$ |
|
|
| A non-zero coefficient \(a\) is then selected and the remaining coefficients are constructed using |
|
|
| $$ |
| b=-a(x_1+x_2) |
| $$ |
|
|
| and |
|
|
| $$ |
| c=ax_1x_2. |
| $$ |
|
|
| This guarantees that the generated quadratic has the selected roots as its exact solutions. |
|
|
| The current training experiment uses: |
|
|
| ```text |
| Training examples: 150,000 |
| Root range: [-50, 50] |
| Task: real-root quadratic equations |
| ``` |
|
|
| The current experiment does not evaluate equations with complex roots. |
|
|
| --- |
|
|
| ## QuadraticLayer |
|
|
| The central custom component is: |
|
|
| ```python |
| class QuadraticLayer(nn.Module): |
| def __init__(self, in_features, out_features): |
| super().__init__() |
| |
| self.linear = nn.Linear(in_features, out_features) |
| |
| self.A = nn.Parameter(torch.ones(out_features) * 0.01) |
| self.B = nn.Parameter(torch.ones(out_features)) |
| self.C = nn.Parameter(torch.zeros(out_features)) |
| |
| def forward(self, x): |
| z = self.linear(x) |
| return self.A * z**2 + self.B * z + self.C |
| ``` |
|
|
| Each output channel therefore computes |
|
|
| $$ |
| y_i=A_i z_i^2+B_i z_i+C_i |
| $$ |
|
|
| where |
|
|
| $$ |
| z=Wx+b. |
| $$ |
|
|
| This gives the model an explicit mechanism for constructing quadratic nonlinear transformations of learned linear projections. |
|
|
| --- |
|
|
| ## Training Configuration |
|
|
| ```text |
| Dataset size: 150,000 |
| Batch size: 256 |
| Epochs: 125 |
| |
| Optimizer: AdamW |
| Learning rate: 3e-4 |
| Weight decay: 1e-4 |
| |
| Loss: SmoothL1Loss |
| Gradient clipping: max_norm = 1.0 |
| |
| Activation: GELU |
| Normalization: LayerNorm |
| ``` |
|
|
| --- |
|
|
| ## Evaluation |
|
|
| One recent independent evaluation run produced the following results: |
|
|
| | Metric | Result | |
| | ---------------------- | -------: | |
| | MAE | 0.084864 | |
| | RMSE | 0.137898 | |
| | Maximum observed error | 1.652145 | |
| | Error < 0.1 | 69.34% | |
| | Error < 0.01 | 0.27% | |
|
|
| These metrics are reported in the normalized root representation unless otherwise specified. |
|
|
| The current evaluation results should be considered preliminary. Earlier experiments used independently generated evaluation sets, so future releases will use fixed train/validation/test splits and fixed random seeds for reproducible comparisons. |
|
|
| The model tends to perform worse when the two roots become very close. |
|
|
| --- |
|
|
| ## Interpretability Analysis |
|
|
| One of the main purposes of SEAI is to investigate whether mathematically meaningful structures emerge inside the network. |
|
|
| Forward hooks are used to record hidden activations from each `QuadraticLayer`. |
|
|
| Neuron activations are compared against candidate mathematical features including |
|
|
| $$ |
| p,\quad q,\quad p^2,\quad q^2,\quad pq, |
| $$ |
|
|
| $$ |
| p^2-q,\qquad p^2+q, |
| $$ |
|
|
| and |
|
|
| $$ |
| \sqrt{|p^2-q|}. |
| $$ |
|
|
| Because the current dataset contains real-root equations, the discriminant-related quantity satisfies |
|
|
| $$ |
| p^2-q\geq0 |
| $$ |
|
|
| up to numerical precision. |
|
|
| The analysis investigates which candidate mathematical quantities are most strongly correlated with individual hidden activations. |
|
|
| --- |
|
|
| ## Observed Representation Structure |
|
|
| The current experiment shows a qualitative progression across layers. |
|
|
| ### First Quadratic Layer |
|
|
| Hidden activations primarily show strong correlations with the original variables |
|
|
| $$ |
| p |
| $$ |
|
|
| and |
|
|
| $$ |
| q. |
| $$ |
|
|
| ### Second Quadratic Layer |
|
|
| Nonlinear mathematical combinations become more prominent, including |
|
|
| $$ |
| p^2+q |
| $$ |
|
|
| and |
|
|
| $$ |
| p^2-q. |
| $$ |
|
|
| ### Third and Fourth Quadratic Layers |
|
|
| The strongest recurring candidate features include |
|
|
| $$ |
| p^2-q |
| $$ |
|
|
| and |
|
|
| $$ |
| \sqrt{p^2-q}. |
| $$ |
|
|
| This produces an experimentally observed progression consistent with the algebraic structure required by the analytical solution: |
|
|
| $$ |
| \boxed{ |
| p,q |
| \rightarrow |
| \text{quadratic combinations} |
| \rightarrow |
| p^2-q |
| \rightarrow |
| \sqrt{p^2-q} |
| \rightarrow |
| x_1,x_2 |
| } |
| $$ |
|
|
| The interpretation of this progression remains exploratory. |
|
|
| --- |
|
|
| ## Main Interpretability Finding |
|
|
| The central observation of the current study is: |
|
|
| > **An 8.04-million-parameter neural network developed hidden representations strongly correlated with the normalized quadratic discriminant and its square-root transformation, despite not receiving either quantity explicitly as an input.** |
|
|
| This does **not** establish that an individual neuron literally performs symbolic discriminant or square-root computation. |
|
|
| Instead, it provides evidence that the trained network contains internal representations that are statistically aligned with mathematically meaningful quantities involved in the classical solution. |
|
|
| --- |
|
|
| ## Error Behaviour Near Repeated Roots |
|
|
| The largest errors tend to occur when |
|
|
| $$ |
| x_1\approx x_2. |
| $$ |
|
|
| For the normalized representation this corresponds to |
|
|
| $$ |
| p^2-q\approx0. |
| $$ |
|
|
| This is the regime in which the square-root term in the analytical solution approaches zero. |
|
|
| A future systematic analysis will measure prediction error as a function of |
|
|
| $$ |
| |x_2-x_1| |
| $$ |
|
|
| and |
|
|
| $$ |
| p^2-q. |
| $$ |
|
|
| --- |
|
|
| ## Limitations |
|
|
| The following limitations are important. |
|
|
| ### Correlation is not proof of symbolic computation |
|
|
| A high correlation between a neuron and a mathematical expression does not prove that the neuron explicitly computes that expression. |
|
|
| ### Preliminary evaluation protocol |
|
|
| The current reported metrics include experiments performed with independently generated evaluation sets. Fixed evaluation data and fixed seeds are required for rigorous model comparison. |
|
|
| ### Restricted data distribution |
|
|
| The model is currently trained on real-root quadratic equations generated from roots in |
|
|
| $$ |
| [-50,50]. |
| $$ |
|
|
| Generalization outside this range has not yet been established. |
|
|
| ### No complex-root regime |
|
|
| The current training and evaluation setup does not investigate equations with |
|
|
| $$ |
| p^2-q<0. |
| $$ |
|
|
| ### Interpretability is exploratory |
|
|
| The current analysis is primarily based on neuron-level correlations. Stronger evidence would require symbolic regression, reproducibility across independent runs, and causal interventions on candidate neurons or representations. |
|
|
| --- |
|
|
| ## Future Research |
|
|
| Planned experiments include: |
|
|
| * fixed reproducible train/validation/test splits; |
| * multiple random seeds; |
| * standard MLP baselines; |
| * ablations of the quadratic transformation; |
| * systematic neuron-level symbolic regression; |
| * causal interventions on neurons correlated with \(p^2-q\); |
| * out-of-distribution evaluation; |
| * larger root and coefficient ranges; |
| * analysis of representation stability across independently trained models; |
| * comparison with conventional polynomial architectures; |
| * extension to higher-degree polynomial equations. |
|
|
| A key research question is whether the observed discriminant-related representations arise systematically from the quadratic architecture or can be reproduced equally well by standard neural networks of comparable size. |
|
|
| --- |
|
|
| ## Intended Use |
|
|
| SEAI is intended primarily as a research prototype for: |
|
|
| * mechanistic interpretability; |
| * mathematical representation learning; |
| * neural symbolic reasoning; |
| * analysis of learned algebraic structure; |
| * controlled studies of neural networks on mathematically defined tasks. |
|
|
| It is not intended to replace conventional numerical solvers in production environments. |
|
|
| For arbitrary quadratic equations, a conventional analytical or numerical solver remains preferable. |
|
|
| --- |
|
|
| ## Repository Structure |
|
|
| ```text |
| SEAI/ |
| ├── train.py |
| ├── model.py |
| ├── interpretability.py |
| ├── evaluation.py |
| ├── requirements.txt |
| ├── README.md |
| └── weights/ |
| └── model_weights_SEAI.pth |
| ``` |
|
|
| The final reproducible release should additionally contain: |
|
|
| ```text |
| ├── data/ |
| │ ├── train.pt |
| │ ├── validation.pt |
| │ └── test.pt |
| ├── configs/ |
| │ └── seai_config.json |
| └── results/ |
| ├── metrics.json |
| └── interpretability/ |
| ``` |
|
|
| --- |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{seai2026, |
| title={SEAI: Learning Mathematical Representations for Quadratic Equation Solving}, |
| author={ALEXFLR}, |
| year={2026}, |
| note={Research prototype} |
| } |
| ``` |
|
|
| The corresponding research paper is intended for publication on arXiv. |
|
|
| --- |
|
|
| ## Project Status |
|
|
| **Research prototype — experimental** |
|
|
| SEAI is primarily a research project investigating whether neural networks can develop internal representations aligned with known mathematical structures. |
|
|
|
|
| 🤖AI |
|
|
| class QuadraticLayer(nn.Module): |
| def __init__(self, in_features, out_features): |
| super().__init__() |
| self.linear = nn.Linear( |
| in_features, |
| out_features |
| ) |
| self.A = nn.Parameter( |
| t.ones(out_features) * 0.01 |
| ) |
| |
| self.B = nn.Parameter( |
| t.ones(out_features) |
| ) |
| |
| self.C = nn.Parameter( |
| t.zeros(out_features) |
| ) |
| |
| def forward(self, x): |
| |
| z = self.linear(x) |
| |
| return ( |
| self.A * z**2 |
| + self.B * z |
| + self.C |
| ) |
| |
| class SEAI(nn.Module): |
| def __init__(self): |
| super().__init__() |
| |
| self.net = nn.Sequential( |
| |
| QuadraticLayer(2, 1000), |
| nn.LayerNorm(1000), |
| nn.GELU(), |
| |
| QuadraticLayer(1000, 2000), |
| nn.LayerNorm(2000), |
| nn.GELU(), |
| |
| QuadraticLayer(2000, 2000), |
| nn.LayerNorm(2000), |
| nn.GELU(), |
| |
| QuadraticLayer(2000, 1000), |
| nn.LayerNorm(1000), |
| nn.GELU(), |
| |
| nn.Linear(1000, 2) |
| ) |
| |
| def forward(self, x): |
| return self.net(x) |
| |
|
|
| model = SEAI().to(device) |
|
|
| model.eval() |
|
|
| with t.no_grad(): |
| ... |
| weights in another file. |
