![\"High-level](\"https://www.tommasovaccari.com/static/high-level-flow-2a2290c9.webp\")
This report presents the work I carried out to introduce a new component into the muDock codebase.
\nThe complete implementation is available in muDock PR #105.
\nAs explained in section 2 of the muDock vectorization paper, the general context of this work is drug discovery: the process through which candidate molecules are investigated as potential effective treatments for a disease. One of the initial phases of this process is virtual screening, where computational techniques are used to analyse a large set of candidate molecules and select the most promising ones before experimental validation in a laboratory. Virtual screening provides an initial filtering step, reducing the amount of work that must later be performed in the laboratory and therefore reducing time and costs.
\nMolecular docking is one of the techniques used in virtual screening. Given a receptor, typically a protein, and a ligand, meaning a candidate molecule, the objective is to estimate how the ligand may arrange itself inside an interaction region of the protein called the binding site. The quality, and therefore the probability, of an arrangement is evaluated through an energy term used as a proxy for binding affinity. A single spatial arrangement of the ligand is called a pose and includes the position, orientation, and internal torsions of the molecule.
\nTo compare poses, docking uses a scoring function: a function that assigns a numerical value to each configuration, approximating the quality of the interaction between the ligand and the receptor. The task of docking is therefore to explore the pose space and identify the poses associated with the most favourable scores.
\nThe computational problem is divided into two levels. The first concerns the number of candidate molecules to evaluate, which can generally be very large. The second concerns the number of possible poses to explore for each receptor-ligand pair. The degrees of freedom of the complex come from the ligand being able to move in space and rotate internally. In general, an exhaustive search over the pose space is not practical.
\nFor this reason, docking engines usually use heuristics and optimisation algorithms that seek a compromise between computational cost and solution quality. In addition, many evaluations of different poses or molecules can be performed independently, making molecular docking a perfect use case for HPC.
\nThe objective of the project was to implement a scoring function in C++ by using the existing infrastructure and pipeline of muDock. In particular, the implemented scoring function is Vinardo, described in the original paper. This required adapting the data representation, atom typing, preprocessing of atom pairs, and score computation to the existing components.
\nThe main difficulty was therefore not only implementing the scoring function formula, but integrating it into an already structured codebase. The work required understanding the muDock architecture, comparing the behaviour with smina as the reference implementation, and introducing all the infrastructure required around the scoring kernel: atom typing, PDBQT topology information, pair preprocessing, and data management within the pipeline.
\nAt the end of the project, the produced output was:
\nThis section explains at a high level the muDock architecture and the essential components for docking. It then explains the theoretical operation of Vinardo.
\nmuDock was created as a microapp for molecular docking, built to replicate an AutoDock-like workflow. In addition to providing a docking engine, the codebase is designed as an environment for experimenting with optimisation and porting techniques, such as vectorisation and support for different execution backends.
\nIn the context of this work, the most relevant aspect is the deliberate separation between the phases of the pipeline: input parsing, molecule representation, scoring, search, and output production. This organisation makes it possible to modify a single component, such as the scoring function, without rewriting the entire execution flow.
\nAt a high level, the execution flow is the following:
\n
In general, a pipeline consists of one or more stages, where each stage represents a unit of work in the pipeline. In muDock, the stages relevant to this work are those related to scoring and genetic optimisation. Each stage follows a common cycle:
\nprepare: prepares the data required to execute the stage on the current batch;operator(): performs the actual work;teardown: collects the results and writes them into the output data structures.This structure can be seen as a stateful functor. The object maintains the state required for execution, such as references to buffers and scratchpads, and makes the computation invocable through operator().
For example, the scoring stage can be represented as follows:
\ntemplate<typename queue_type>\nstruct scoring_stage {\n std::shared_ptr<scratchpad<queue_type>> scratch;\n void prepare(batch<static_molecule>& batch);\n void operator()();\n void teardown(batch<static_molecule>& batch);\n};The available pipelines include a scoring-only pipeline, which calculates the score associated with the interaction between the protein and ligand in the given configuration, and a genetic pipeline, where the genetic stage uses the scoring function to evaluate the poses generated during the heuristic exploration of the pose space. In the second case, the parameters called knobs in the codebase control the number of individuals, number of generations, mutation probability, and other aspects of the search.
\nAt this point it is necessary to introduce the abstractions used to manage memory and data because this context requires enough flexibility to work independently from the execution backend and the framework used on that backend.
\nmuDock mainly uses two abstractions: buffer and scratchpad. A buffer represents a typed memory area that can be used to read and write data, but also to move it between the host, usually the machine launching the computation, and the device, the machine used to perform the intensive computation, for example one equipped with a GPU.
\nThe scratchpad is a shared registry of buffers. It allows stages to retrieve the data they need through common keys, avoiding unnecessary reconstruction or copying. The scratchpad can contain data that remains stable during the run, such as the representation of the protein, as well as data produced by one stage and consumed by the next. One example is the buffers associated with poses, containing ligand coordinates, which are produced during genetic exploration and then used by the scoring stage to evaluate the generated configurations.
\nThe fact that stages can pass data through shared buffers and key conventions is fundamental to guarantee stable contracts between different stages and compose them flexibly. This becomes clearer later, when the integration of the Vinardo-based scoring stage only requires implementing a stage that respects the input and output contracts, without modifying the genetic pipeline or the other existing scoring stages.
\nThe fundamental component of each stage is the kernel, which represents the actual numerical computation. The kernel is invoked through operator() and works on data already prepared by the previous phase. Since this is an HPC context and different execution backends must be supported, there can be different implementations of the same kernel, for example a CUDA version for GPU execution and a standard C++ version for CPU execution. For this work, the generic interface of the Vinardo kernel was implemented but specialised only for the CPU, leaving a GPU version as a possible future development.
Vinardo was developed as an evolution of the AutoDock Vina scoring function. For this reason, it preserves many of its characteristics and the same general structure. The quality of a pose is evaluated through a scoring function that sums energy contributions calculated over pairs of atoms.
\nBefore describing the scoring function, it is useful to distinguish between the conformation-dependent score, meaning the spatial arrangement, and the affinity value reported as an estimate of binding free energy. In the Vina model inherited by Vinardo, the conformation-dependent score is denoted by C and depends on the specific ligand pose.
In particular, C consists of an intermolecular contribution associated with protein-ligand interactions and an intramolecular contribution associated with interactions internal to the ligand:
C = C_{inter} + C_{intra}The value of C is calculated as the sum of the contributions over the atom pairs considered valid:
C = \\sum_{(i,j) \\in \\mathcal{P}} F_{T_i,T_j}(d_{ij})where \\mathcal{P} represents the set of atom pairs selected as valid during preprocessing. The considered pairs are:
In both cases, pairs containing at least one hydrogen atom are discarded. For ligand-ligand pairs, only atom pairs separated by more than three covalent bonds and able to move relative to each other are retained.
\nThe contribution associated with a single atom pair must then be defined. In the previous formula, T_i and T_j represent the atom types of the two atoms, while d_{ij} is the surface distance between the atoms in the pair. It is obtained by subtracting the radii associated with the two atom types from the Euclidean distance between the atomic centres:
d_{ij} = r_{ij} - R_i - R_jThe function F_{T_i,T_j} consists of multiple weighted energy terms:
F_{T_i,T_j}(d) =\nw_1 \\cdot \\text{Gauss}(d) +\nw_2 \\cdot \\text{Repulsion}(d) +\nw_3 \\cdot \\text{Hydrophobic}(d) +\nw_4 \\cdot \\text{HBond}(d)The Gauss and Repulsion terms are evaluated for all pairs, while Hydrophobic and HBond depend on the atom types involved. In particular, the hydrophobic term is active only for pairs of hydrophobic atoms, while the hydrogen bond term is active when the pair may represent a donor-acceptor interaction.
\nThe individual terms are defined as follows:
\n\\text{Gauss}(d) = e^{-\\left(\\frac{d - o_1}{s_1}\\right)^2}\\text{Repulsion}(d) =\n\\begin{cases}\nd^2 & \\text{if } d < 0 \\\\\n0 & \\text{if } d \\geq 0\n\\end{cases}\\text{Hydrophobic}(d) =\n\\begin{cases}\n1 & \\text{if } d \\leq p_1 \\\\\n\\dfrac{p_2 - d}{p_2 - p_1} & \\text{if } p_1 < d < p_2 \\\\\n0 & \\text{if } d \\geq p_2\n\\end{cases}\\text{HBond}(d) =\n\\begin{cases}\n1 & \\text{if } d \\leq h_1 \\\\\n\\dfrac{d}{h_1} & \\text{if } h_1 < d < 0 \\\\\n0 & \\text{if } d \\geq 0\n\\end{cases}From an implementation perspective, the Hydrophobic and HBond terms are enabled through two flags precomputed during preprocessing. This means that scoring does not need to reevaluate the atomic properties of the pair and only has to apply the function to the prepared data.
\nThe weights used by Vinardo are those reported in the original paper:
\nw_{\\text{gauss}} = -0.045,\\quad\nw_{\\text{repulsion}} = 0.800,\\quad\nw_{\\text{hydrophobic}} = -0.035,\\quad\nw_{\\text{hbond}} = -0.600The main parameters of the terms are:
\no_1 = 0.0,\\quad\ns_1 = 0.8,\\quad\np_1 = 0.0,\\quad\np_2 = 2.5,\\quad\nh_1 = -0.6Once the conformation score C has been calculated for the poses of a ligand, it is converted into affinity through a function g, producing an estimate of the binding free energy.
For the best pose:
\ns_1 = g(C_1 - C_{intra,1}) = g(C_{inter,1})For subsequent poses, Vina instead subtracts the intramolecular contribution of the best pose to preserve the relative ranking:
\ns_k = g(C_k - C_{intra,1})The Vina paper also proposes a concrete form for g, which was later inherited by Vinardo. The function normalises the value with respect to the number of rotatable bonds in the ligand:
g(x) = \\frac{x}{1 + w_{rot} \\cdot N_{rot}}Operationally, N_{rot} refers to the number of rotatable bonds in the ligand. This value is not retrieved directly from the field in the PDBQT file, but reconstructed from an in-memory topological representation of the ligand structure. In particular, the rotors described by the PDBQT torsion tree are considered. A rotatable bond is counted only when both endpoint atoms are not hydrogens and each has at least two bonds with non-hydrogen atoms.
From an implementation perspective, this structure therefore requires several components: atom typing compatible with Vinardo, a data representation suitable for scoring, a preprocessing phase to construct valid atom pairs, and a scoring function.
\nThis section describes the components implemented to integrate the new scoring function into the existing pipeline. The description follows the data flow: from the topological information extracted while parsing the PDBQT file, through the construction of the Vinardo view of the molecules, to atom-pair preprocessing and the scoring stage. During development, smina was used as the reference implementation to compare atom typing, pair preprocessing, and scoring values. The final integration of the components described in this section is collected in PR #105.
\nThe first required modification concerned the management of the topological information in the PDBQT file. The standard molecule representation contains atoms, coordinates, atom types, and bonds, but reproducing the behaviour of smina also requires information related to the ligand torsion tree.
\nIn particular, two pieces of data are needed:
\nN_{rot}.These data are derived from the ROOT/BRANCH/ENDBRANCH structure of the PDBQT file. When a ligand in PDBQT format is parsed, an intermediate representation of the torsion tree is constructed. The ROOT node contains the atoms of the main rigid block, while each BRANCH represents a rotatable bond between a from atom and a to atom, together with the subtree of atoms that depends on that bond.
This representation is used during parsing to construct the mobility matrix with semantics consistent with smina and to obtain the list of ligand rotors. Only the data required by subsequent phases are kept in static_molecule, through the pdbqt_ligand_data field.
This choice is forced by the existing pipeline, where ligands processed in batches are represented as static_molecule, while the protein is loaded separately as dynamic_molecule, stable data for the run. A possible refactor of the type used by the ligand queue remains open. If the pipeline supported a more specific type, the PDBQT data could be moved out of static_molecule and collected in a layer dedicated to Vinardo.
This pragmatic solution avoids reparsing the PDBQT file inside the scoring stage. The stage can instead directly use the mobility matrix and rotor list during batch preparation.
\nOnce the molecule representation is available, a view of the atoms consistent with the Vinardo scoring model must be constructed. The parsed molecule contains atom types following the AutoDock scheme, while Vinardo uses a set of atom types and properties derived from the model used by Vina.
\nFor this reason, a vinardo_layer was introduced on top of the existing molecule, following the same approach already used by the other layers in the codebase. The layer adds a Vinardo-specific view, exposing the Vinardo type, atomic radius, and flags required to determine whether each atom can contribute to hydrophobic or hydrogen-bond interactions.
Type conversion takes place in multiple steps. First, the AutoDock type used internally by muDock is normalised into an intermediate domain compatible with the Vinardo conversion table. An initial Vinardo type is assigned from this domain, but it may still be ambiguous. Ambiguous cases are then resolved using the topological context of the molecule.
\nIn particular, the bond graph is used to distinguish atoms that have the same initial type but different properties for scoring purposes, for example based on the presence of heteroatom neighbours or polar hydrogens.
\nThe result of the layer is then used by preprocessing to construct valid atom pairs and precompute the pair properties required by the scoring function.
\nThe scoring function is applied to a set of atom pairs selected in advance. For this reason, a preprocessing phase was introduced that constructs the pairs to evaluate during scoring from the Vinardo layers of the protein and ligand.
\nProtein-ligand pairs are generated by considering every combination of protein and ligand atoms, excluding pairs that contain hydrogens. The ligand-ligand case requires more work because not every internal pair in the ligand must contribute to the score. First, the ligand bond graph is constructed and a Breadth First Search limited to depth three is performed for each atom. The result is saved in a within_three_bonds matrix, which indicates whether two atoms are reachable within three covalent bonds.
A ligand-ligand pair is therefore retained only when both conditions hold: the two atoms are not reachable within three covalent bonds and they are marked as relatively mobile in the mobility matrix constructed from the PDBQT torsion tree.
\nFor each selected pair, the indices of the two atoms, the sum of the atomic radii, and two Boolean flags are saved. One flag indicates whether the pair can contribute to the hydrophobic term and the other whether it can contribute to the hydrogen-bond term. These data are independent from the specific pose and can therefore be calculated once before kernel execution.
\nIn the current implementation, these data are kept in the private buffers of the stage for the current batch. One possible future optimisation would expose them in the pipeline scratchpad if they needed to be reused by multiple stages or preserved beyond a single stage preparation. During genetic exploration, the generated poses change, but ligand topology, atom types, and the set of valid pairs do not. The kernel therefore only needs to update the transformed coordinates and geometric distances, reusing the already computed pair indices, radius sums, and interaction flags.
\nIn general, this preprocessing phase reduces the work performed during scoring. The kernel does not need to query the chemical structure of the molecule and can instead iterate over prepared arrays, calculating only the geometric distances and the energy terms that depend on them for each pose.
\nThe central point of the integration is the scoring stage, where the chemical representation built in the previous phases is transformed into the data actually used to calculate the score.
\nLike other existing stages, the vinardo_score stage implements the three-phase contract. In prepare(), it reads the data available in the scratchpads and constructs the buffers required by the kernel. In operator(), the computation is executed and the SCORES buffer is populated. In teardown(), the result is retrieved and written into the ligand properties.
The stage contains different categories of data with different roles and lifetimes in the computation. They are organised into three levels.
\nThe first level is the device scratchpad, maintained in the stage code as device_scratch and used for data that remains stable across batches. For Vinardo, it contains the protein coordinates in the PROT_X, PROT_Y, and PROT_Z buffers because the protein is shared by every ligand processed during the run.
The second level is the pipeline scratchpad, accessible in the stage through the scratch member and shared by stages operating on the same batch. This scratchpad contains flat buffers used as communication points between stages, produced from the C++ objects constructed previously. For example, in the genetic pipeline, the genetic stage produces new poses by modifying the ligand coordinates and writes them into the X_SCRATCH, Y_SCRATCH, and Z_SCRATCH buffers. The scoring stage then reads the already transformed coordinates from those same buffers and writes the results into SCORES.
The third level consists of the private buffers of the vinardo_score stage. These contain data required only by the Vinardo kernel, such as offsets, indices of protein-ligand and ligand-ligand pairs, sums of atomic radii, and interaction flags. These buffers represent a flattened version of the data described through more convenient C++ objects in the previous phases. The stage converts them into contiguous arrays that are more suitable for sequential kernel access.
The prepare() phase is where this data organisation is concretely constructed for the current batch. It receives a batch of static_molecule objects and first loads data shared with the other stages into scratch, such as atom counts and the working coordinate buffers of the ligands.
Next, a vinardo_ligand is constructed for every ligand in the batch. This object uses the vinardo_layer and the PDBQT data already associated with the molecule to calculate ligand-ligand pairs and the number of rotatable torsions N_{rot}. Protein-ligand pairs are instead generated inside the stage.
The prepare() phase ends by moving the private stage buffers to the device and constructing the vinardo_score_kernel, which receives pointers to those buffers. From this point onward, the computation uses only the linearised data prepared for the current batch.
The operator() phase simply invokes the kernel and delegates the actual score computation to it.
During teardown(), the stage copies the SCORES buffer from the device to the host and assigns the score to the SCORE property of the ligand. In this way, the result reenters the normal output mechanism of the pipeline.
The kernel is the component that consumes the buffers produced by prepare() and calculates the score of the poses in the batch when the stage invokes operator(). The general interface is defined by vinardo_score_kernel<queue_type>. The object maintains pointers to ligand coordinates, protein coordinates, pair arrays, interaction flags, the number of rotatable torsions, and the SCORES buffer.
In this project, the kernel was implemented only for the CPU backend through queue_cpp. In this version, operator() calls calc_vinardo_energy, passing the pointers to the buffers prepared by the stage.
The calc_vinardo_energy() function iterates over the ligands in the batch and the poses associated with each ligand. For each pose, it separately calculates the protein-ligand contribution and the ligand-ligand contribution. For each pair, it reads the atom indices, retrieves the coordinates, and calculates the Euclidean distance between the two atoms. If the distance is greater than or equal to 8 Angstrom, the pair is discarded. Otherwise, the surface distance is calculated by subtracting the precomputed sum of atomic radii from the Euclidean distance. The surface distance and the precomputed hydrophobic and hydrogen-bond interaction flags are then passed to compute_vinardo_pair_energy().
The function that calculates the energy contribution of a single pair was kept separate from the kernel. This choice made it possible to validate the Vinardo formula independently of the execution infrastructure. The separation also makes it possible to expose debugging APIs, useful for observing the individual energy terms and verifying them against the reference.
\nThe kernel keeps the intermolecular and intramolecular contributions separate. This is required to calculate affinity correctly when multiple poses are present for the same ligand. First, the best pose is identified according to the raw score. Each pose is then corrected using the intramolecular contribution of the reference pose and normalised through N_{rot}. The final value is written into the SCORES buffer.
The pipeline integration is completed in pipeline.hpp, where the vinardo_score_pipeline and genetic_vinardo_pipeline aliases were added. The first instantiates the vinardo_score stage in the scoring-only pipeline, while the second uses it as the scoring stage inside the genetic pipeline. In this way, Vinardo follows the same pattern already used by AutoDock without requiring a separate flow.
The search algorithm and scoring function can be selected at runtime through the CLI options --search, with the values none and genetic, and --score, with the values adt and vinardo. This makes it possible to freely combine the four available configurations.
This section describes the process used to validate the Vinardo implementation and the results obtained from comparison with the reference.
\nValidation was organised incrementally, using smina as the reference implementation for the expected Vinardo behaviour. The final affinity value was therefore not the only value validated. The work was instead divided into the main components of the pipeline: atom typing, pair preprocessing, reconstruction of ligand topology information, and affinity calculation.
\nTo avoid introducing debugging code into the final PR, part of the work was carried out on parallel validation branches. These branches were derived from the corresponding development branches and preserved the same implementation logic, but added dump and instrumentation tools to export data that could be compared with smina. This made it possible to associate atoms, atom pairs, and energy contributions through consistent indices across the two codebases, identifying bugs and divergences at the level of individual components. Development therefore moved to the next component only after the newly introduced one had been completely aligned with the reference.
\nDuring development, a set of 15 complexes from PDBbind v2024 was used to provide broader validation. The considered complexes were identified by the following IDs: 1BNM, 1FKF, 1NLO, 1NLP, 1ABT, 1B2I, 1AO8, 1CWB, 1CWC, 4SGA, 1STD, 220L, 1ILQ, 1B56, and 1EVH.
\nThe following table summarises the main validation points used during development. For each component, the output produced by muDock was compared with the output obtained from the reference.
\n| Component | Comparison | Result | Reference |
|---|---|---|---|
| Atom typing | Vinardo types assigned to protein and ligand atoms | Aligned with smina after correction | PR #5 / PR #8 |
| Protein-ligand preprocessing | Protein-ligand atom pairs selected for scoring | Aligned after excluding hydrogens | PR #6 / PR #7 |
| Ligand-ligand preprocessing | Internal ligand pairs selected through topological constraints and relative mobility | Aligned using the PDBQT torsion tree | PR #6 / PR #10 |
| Number of rotatable torsions | N_{rot} used to normalise affinity | Reconstructed as in the reference | PR #10 |
| Scoring | Final affinity value on benchmark cases | Close to the reference, with an explainable residual delta | PR #12 / PR #13 |
The incremental comparison showed that alignment did not depend only on the formula in the paper, but also on several implementation choices. The most relevant cases are reported below because they led to corrections or implementation decisions in the final code.
\nThe per-atom comparison with smina highlighted errors in the table of data associated with the Vinardo types, in particular in the flags and values later used by preprocessing and the scoring function. The table was corrected and the comparison repeated until the same typing as the reference was obtained.
\nIn the first preprocessing comparison, muDock generated more atom pairs than smina. The difference came from the reference excluding pairs that contain at least one hydrogen atom, while the initial implementation retained them. Preprocessing was therefore modified to apply the same filter to both protein-ligand and ligand-ligand pairs.
\nAfter excluding hydrogens, a divergence remained in the internal ligand pairs. The problem was not related to the topological distance between atoms, but to the definition of relative mobility.
\nAs explained in this comment, the initial implementation used a component already present in muDock: fragments. For each rotatable bond, this component temporarily removes the corresponding connection from the graph, assuming that the rotatable bond represents a valid cut of the graph. A BFS is then performed over the components reachable from the two endpoints of the bond, identifying the atoms on both sides of the cut and constructing the mobile-fragment mask.
\nThis concept was not semantically identical to the one used by the reference, which derives relative mobility from the PDBQT torsion tree. The concrete divergence mainly concerned pairs involving the atoms at the endpoints of the rotor. In the semantics reconstructed from the reference, the from/to pair and the pairs between each of these two atoms and the atoms in the sub-branch are marked as not relatively mobile. The fragment-based logic instead included these pairs among the mobile ones. As a result, preprocessing selected more ligand-ligand pairs than the reference.
To align the behaviour, a mobility matrix with semantics consistent with smina was introduced.
\nAnother divergence concerned the value of N_{rot} used to normalise affinity. Here too, directly using either the number of rotors already available in the molecule or the value present in the PDBQT file was not sufficient because the reference applies its own semantics to the rotors that are actually counted.
To align the behaviour, the torsion-tree representation introduced for the mobility matrix was reused. Starting from the rotors described in the PDBQT file, those where at least one endpoint is a hydrogen or does not have at least two bonds with non-hydrogen atoms are excluded. The resulting value is then used in the affinity calculation, making the normalisation consistent with smina.
\nAfter aligning the previous components, a residual delta from smina remained. This delta was traced back to curl, a transformation applied by the reference to positive energy contributions before the final normalisation of affinity.
However, curl is applied to intermediate computation results at a level that is not present in the implementation developed in this project. It is also an implementation choice of the reference rather than a term described in the theoretical formulation of Vinardo. For this reason, it was not included in the final implementation.
A test was still performed on an instrumented branch, restructuring the code to expose the same intermediate results. When curl was also applied, the final affinity value aligned with the reference up to machine precision.
Final validation on the dataset compared the affinity calculated by muDock using Vinardo with the affinity produced by smina on the same inputs. The results are shown below.
\n| PDB ID | muDock affinity | smina affinity | Absolute delta |
|---|---|---|---|
| 1BNM | -6.669172 | -6.669182 | 1.001600 \\cdot 10^{-5} |
| 1FKF | -11.700834 | -11.700848 | 1.381000 \\cdot 10^{-5} |
| 1NLO | 1.007401 | 0.992802 | 1.459884 \\cdot 10^{-2} |
| 1NLP | -3.374581 | -3.378935 | 4.354485 \\cdot 10^{-3} |
| 1ABT | -8.456276 | -8.456291 | 1.554700 \\cdot 10^{-5} |
| 1B2I | -2.721796 | -2.723378 | 1.582090 \\cdot 10^{-3} |
| 1AO8 | -6.664783 | -6.664826 | 4.205000 \\cdot 10^{-5} |
| 1CWB | -7.065887 | -7.065935 | 4.775400 \\cdot 10^{-5} |
| 1CWC | -7.466530 | -7.466600 | 6.955900 \\cdot 10^{-5} |
| 4SGA | -8.104091 | -8.104555 | 4.640370 \\cdot 10^{-4} |
| 1STD | -9.236320 | -9.236327 | 6.967000 \\cdot 10^{-6} |
| 220L | -4.433152 | -4.433151 | 8.110000 \\cdot 10^{-7} |
| 1ILQ | -5.019547 | -5.020099 | 5.520600 \\cdot 10^{-4} |
| 1B56 | -5.262350 | -5.262351 | 1.259000 \\cdot 10^{-6} |
| 1EVH | -7.126778 | -7.126916 | 1.374770 \\cdot 10^{-4} |
Across the 15 cases, the mean absolute delta was 1.459784 \\cdot 10^{-3}, with a maximum delta of 1.459884 \\cdot 10^{-2}. These results show that the implementation is aligned with the reference, with residual differences explained by the absence of curl in this implementation.
In addition to numerical validation against smina, a functional forward-screening experiment was performed on the CASF-2016 dataset. The Vinardo test was conducted on the 57 targets of the CASF-2016 benchmark, using the scores produced by the CPU scoring-only pipeline and evaluating them with the official forward-screening script provided with CASF.
\nIn the forward-screening test, 285 ligands and 100 poses per ligand are considered for each target. Each ligand is assigned the best score among its poses and the ligands are ranked according to that score. The top 1\\%, 5\\%, and 10\\% are then evaluated. With 285 ligands and CASF rounding, these correspond to 3, 14, and 29 ligands respectively.
Let T be the set of targets, R_t^\\alpha the ligands selected in the top \\alpha for target t, B_t the set of known binders for that target, and L1_t the binder with the best experimental affinity. Two metrics were reported:
SR_\\alpha = \\frac{1}{|T|}\\sum_{t \\in T}\\mathbf{1}[L1_t \\in R_t^\\alpha]\\cdot 100Success Rate measures the percentage of targets for which the best experimental binder appears in the top \\alpha. Enrichment Factor measures how much the ranking enriches true binders compared with random selection:
EF_\\alpha = \\frac{1}{|T|}\\sum_{t \\in T}\n\\frac{|R_t^\\alpha \\cap B_t|}{|B_t| \\cdot \\alpha}An EF value close to 1 indicates approximately random behaviour. The denominator represents the expected number of known binders obtained through random selection.
| Method | Ranking | Top 1% (3) | Top 5% (14) | Top 10% (29) |
|---|---|---|---|---|
| Vinardo | lower is better | 33.3% / 9.34 | 49.1% / 4.18 | 63.2% / 3.19 |
Each cell reports Success Rate / Enrichment Factor.
\nThe results show very good enrichment for Vinardo across the complete benchmark of 57 targets, with the best experimental binder recovered in the top 1\\%, 5\\%, and 10\\% for 19, 28, and 36 targets respectively.
The extended validation over the 15 PDBbind complexes was used during development, but it was not included directly in the CTests of the final PR. The automatic tests instead use the complexes already present in muDock, avoiding the introduction of new datasets into the codebase.
\nFor each complex used in the CTests, a reference affinity value produced by smina was saved. The automatic test runs the muDock Vinardo pipeline on the same PDBQT files and verifies that the calculated affinity remains within an absolute tolerance of 10^{-3} from the reference.
In this way, the CTests automatically verify that subsequent changes do not alter the behaviour of the Vinardo pipeline integrated into muDock.
\nThe work introduced the Vinardo scoring function into muDock and integrated it into the existing pipeline. The integration required constructing a Vinardo view of the molecules, converting atom types, preprocessing atom pairs, reconstructing ligand topology information from PDBQT, and implementing the scoring stage with its CPU kernel.
\nValidation was performed by progressively comparing the components with smina, used as the reference implementation. This made it possible to identify and correct divergences such as hydrogen management, the semantics of relative mobility inside the ligand, and the calculation of the number of rotatable torsions. The final comparison over 15 PDBbind complexes shows good alignment with the reference. The residual differences can be traced back to the curl applied by smina to intermediate contributions, which was not included in the final implementation.
Some future developments remain. During PR review, a problem emerged in estimating the memory required by protein-ligand buffers. An initial version restricted Vinardo to a batch size of one because the static estimation API receives only the number of ligand atoms. This constraint was removed by introducing a conservative estimate based on an upper bound for the number of considered protein atoms. The estimate could still be refined, especially for accelerated backends and more precise preallocation.
\nA second point concerns the auxiliary data extracted from PDBQT. It is currently constructed during static molecule parsing even when it is not required, although it is used only by Vinardo. With runtime selection of the scoring function, this data could be constructed only in the path that uses it.
\nAnother future development is extending the Vinardo kernel to backends other than the CPU version implemented in this project. Finally, parts of preprocessing could be optimised by storing pose-independent data in the scratchpad, such as valid atom pairs and their associated data, avoiding reconstruction when ligand topology does not change.
\nOverall, the project added a functional Vinardo scoring path to muDock, validated it against an external reference, and integrated it into the stage-based model of the codebase. It provides complete runtime support through the CLI for selecting the search algorithm and scoring function, while leaving a solid foundation for a future port to accelerated backends.
\nNow that I have covered the fundamentals of neural networks and have a grasp of how a transformer works, I decided to build a small GPT from scratch, train it on different datasets, and see what happens.
\nThe first step, as usual, was to implement it without relying on libraries that already provide the complete architecture. Rebuilding it from scratch helps me remember that there is nothing magical about it: it is math and scale working together to produce something that can appear surprisingly complex.
\nI wanted to have a solid base that I could fork whenever I wanted to play around with a model, so I created a repository that I can use as a starting point for different experiments. The details need to change depending on the experiment, but the core architecture remains the same.
\nYou can play with this model even on a laptop, but if you want to experiment with a bit more scale, you will need a GPU. For these experiments, I am using an old PC that I built in 2020. It is not particularly powerful, but it is a good starting point. I turned it into a server running Ubuntu LTS and connected it through Tailscale, so I can access it over SSH. It has 32 GB of RAM, a Ryzen 5 3600, and, most importantly, an RX 5600 XT. To use the GPU as a backend, I had to compile a custom PyTorch build and configure it so that the code can access the GPU through the CUDA device interface exposed by PyTorch.
\nBefore looking at the architecture, I want to derive the computational cost of training and inference for a model like this one.
\nComputational cost is usually measured in FLOPs (floating-point operations). In this approximation, one multiplication and one addition each count as one FLOP. The goal is to derive a rough mental model of the computational cost of training and inference for an LLM.
\nWe can start from the basics. The fundamental building block is matrix multiplication. Suppose we want to multiply the following matrices:
\nC = A @ Bwhere A has shape M \\times K, B has shape K \\times N, and C has shape M \\times N. Each element of C is obtained from the dot product between one row of A and one column of B, so:
FLOPs = 2 * \\ K * \\ N * MFor each dot product, we perform K multiplications and K-1 additions, which we approximate as 2K operations. We repeat this for every row of A and every column of B.
Now we can extend that reasoning to a classic linear layer:
\ny = x@wNothing that we have not seen yet, it's still a matrix multiplication.\nSuppose that x has shape (B,T,C), where B is the batch size, T is the sequence length, and C is the embedding dimension. The weight matrix W has shape (C,D).
We can treat the first two dimensions of x as B \\cdot T rows, each multiplied by W. The cost is therefore:
flops_{linear} = 2 \\cdot B \\cdot T \\cdot C \\cdot DWhen C=D, this becomes:
flops_{linear} = 2 \\cdot B \\cdot T \\cdot C^2Now we can try to estimate the cost of an attention block.\nIn an attention block we have the following matrices:
\nx: (B, T, C) inputW_q: (C, head\\_size)W_k: (C, head\\_size)W_v: (C, head\\_size)Then we have to perform the following operation, for a single head:
\nq = x @ W_q → (B, T, head_size)k = x @ W_k → (B, T, head_size)v = x @ W_v → (B, T, head_size)scores = q @ k.transpose(-2, -1) → (B, T, T)attention = softmax(scores) → (B, T, T)y = attention @ v → (B, T, head_size)Now we can compute the FLOPs required by each step for a single head:
\nflops = 2 \\cdot B \\cdot C \\cdot T \\cdot head\\_size\nflops = 2 \\cdot B \\cdot C \\cdot T \\cdot head\\_size\nflops = 2 \\cdot B \\cdot C \\cdot T \\cdot head\\_size\nflops = 2 \\cdot B \\cdot head\\_size \\cdot T^2\nflops \\approx 4 \\cdot B \\cdot T^2 (derived below)flops = 2 \\cdot B \\cdot T^2 \\cdot head\\_size\nSince every attention head performs the same computation independently, and there are h heads, the total cost is h times the cost derived above. Using h \\cdot head\\_size = C:
h \\cdot \\big( \\underbrace{6 \\cdot B \\cdot T \\cdot C \\cdot head\\_size}_{\\text{qkv}} + \\underbrace{4 \\cdot B \\cdot T^2 \\cdot head\\_size}_{\\text{scores + weighted sum}} \\big) + \\underbrace{4 \\cdot B \\cdot T^2 \\cdot h}_{\\text{softmax}}\n= 6BTC^2 + 4BT^2C + 4BT^2hAfter the heads are concatenated back into a (B,T,C) tensor, there's one more matmul — the output projection — which is not per-head, so it isn't scaled by h:
flops_{out\\_proj} = 2 \\cdot B \\cdot T \\cdot C^2Adding it in:
\n\\boxed{flops_{attention} = 8 \\cdot B \\cdot T \\cdot C^2 \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot C \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot h}On the softmax cost: Softmax on a row of length T requires: a max-reduction for numerical stability (T-1 comparisons), a subtract-and-exponentiate pass (2T flops — one subtraction and one exp per element, counting exp as 1 flop by convention), a sum-reduction (T-1 additions), and a final divide (T divisions). That's \\approx 5T flops per row, though 4T is the more common rounding in the literature and the two are interchangeable at this level of approximation. There are B \\cdot T rows per head (one per query position, per batch element), so one head's softmax costs \\approx 4 \\cdot B \\cdot T^2 — independent of head\\_size. This is why it has to be added on top of the per-head total above rather than folded into it via the h \\cdot head\\_size = C substitution: it's the one term whose cost doesn't factor through that product, since it scales with the number of heads (independent softmax calls), not with how wide each head is.
Now we can derive the cost of the other fundamental block: the feedforward network. It usually contains two linear layers with a non-linearity between them. The first linear layer expands the input dimension by a factor of four, and the second projects it back to the original dimension:
\nWe already have the formula for a linear layer, so we can reuse it directly for step 1 and step 3.
\nflops = 2 \\cdot B \\cdot T \\cdot C \\cdot 4C = 8 \\cdot B \\cdot T \\cdot C^2\nB \\cdot T \\cdot 4C), negligible compared to the matmuls, so we drop itflops = 2 \\cdot B \\cdot T \\cdot 4C \\cdot C = 8 \\cdot B \\cdot T \\cdot C^2\nSumming:
\n\\boxed{flops_{FFN} = 16 \\cdot B \\cdot T \\cdot C^2}Notice this is exactly twice the 8 \\cdot B \\cdot T \\cdot C^2 term from attention (the qkv + output projection part) — the FFN is the more expensive of the two sub-blocks per layer, by a factor of 2, for the standard 4x expansion ratio.
We can now combine the two components to obtain the computational cost of one transformer layer, which contains an attention block and a feedforward network. A complete transformer stacks multiple layers of this type.
\n\\boxed{flops_{layer} = flops_{attention} + flops_{ffn} = 8 \\cdot B \\cdot T \\cdot C^2 \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot C \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot h \\;+\\;16 \\cdot B \\cdot T \\cdot C^2 }So at the end we obtain
\n\\boxed{flops_{layer} = 24 \\cdot B \\cdot T \\cdot C^2 \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot C \\;+\\; 4 \\cdot B \\cdot T^2 \\cdot h }If the model contains l layers, their contribution to one forward pass is l times flops_{layer}. The complete model also contains the embedding layer and the final projection over the vocabulary.
The embedding layer is basically a lookup, so we can consider its arithmetic cost equal to zero in this approximation. At the end of the model we also have the language model head, which projects the representation of every token from C to vocab_size. This is another linear layer, so its cost is:
flops_{lm\\_head} = 2 \\cdot B \\cdot T \\cdot C \\cdot vocab\\_sizeEven when the language model head shares its weights with the token embedding, we still have to perform this matrix multiplication. Weight tying reduces the number of parameters, not the computational cost of the projection.
\nPutting together the l transformer layers and the final projection, we obtain the cost of one complete forward pass:
\\boxed{flops_{forward} = l \\cdot \\left(24 \\cdot B \\cdot T \\cdot C^2 + 4 \\cdot B \\cdot T^2 \\cdot C + 4 \\cdot B \\cdot T^2 \\cdot h\\right) + 2 \\cdot B \\cdot T \\cdot C \\cdot vocab\\_size}Often it is more useful to reason about the cost per token. In one forward pass we process B \\cdot T tokens, so we can divide the previous formula by this value:
\\boxed{flops_{forward\\_per\\_token} = 24 \\cdot l \\cdot C^2 + 4 \\cdot l \\cdot T \\cdot C + 4 \\cdot l \\cdot T \\cdot h + 2 \\cdot C \\cdot vocab\\_size}This is the cost of a complete forward pass over a sequence of length T. In this implementation, autoregressive generation runs another forward pass for every generated token. Implementations with a KV cache can reuse previous keys and values, changing the inference cost.
These formulas make the two main scaling behaviours more visible. The projections inside the transformer grow quadratically with the embedding dimension C. In the complete forward pass, attention grows quadratically with the context length T; after dividing by the number of tokens, the corresponding per-token cost grows linearly with T.
For training we also have to perform the backward pass. A useful rough approximation is that the backward pass costs around two times the forward pass:
\nflops_{backward} \\approx 2 \\cdot flops_{forward}The complete training cost is then approximately three times the forward cost. Per token we obtain:
\n\\boxed{flops_{training\\_per\\_token} \\approx 72 \\cdot l \\cdot C^2 + 12 \\cdot l \\cdot T \\cdot C + 12 \\cdot l \\cdot T \\cdot h + 6 \\cdot C \\cdot vocab\\_size}Clearly I have not developed from scratch this model architecture, is derived from the original attention paper, plus some other sources, like karpathy's videos and also deep learning books. It's pretty easy, nothing too complicated, I just want to graps the fundamentals, so it's deliberately written in clanky python and pytorch, nothing that must be taken for serious.
\nYou may already be familiar with the architecture, so you can skip this section. To better retain what I am learning, I like to sketch the systems I am building. In this case, I used Excalidraw, which I found very helpful:
\n![\"architecture\"](\"https://www.tommasovaccari.com/static/architecture-64a5bdb9.webp\")
As we have already seen, the classic transformer block contains a self-attention layer followed by a feedforward network. The main difference is the placement of layer normalization: the original paper uses post-norm, while this implementation uses pre-norm. Pre-norm generally makes optimization more stable, especially as the number of layers increases.
\nThe implementation is available on GitHub, together with the code required for training and inference. Instead of loading the entire training set into RAM, I used NumPy memory mapping to create a disk-backed array and load only the required portions. This prevented crashes and out-of-memory errors during training.
\nNow we can look at some examples of this custom GPT at work.
\nThe first experiment is SumGPT, a small model trained to sum two three-digit numbers. The input has a format like 123+456= and the model has to generate the result followed by a newline. I decided to generate the digits of the result in reverse order, because this aligns the carry operation with the autoregressive direction of the model: it can start from the units and then move to tens, hundreds and thousands.
For this task I used a deliberately small configuration, with an embedding size of 128, 4 attention heads, 4 transformer blocks and a context length of 12. The complete operation fits inside this context, so there is no reason to use a larger model. On my GPU the training took around three minutes, and on 1000 randomly generated sums it reached around 99% accuracy.
\nAn interesting thing that I noticed during training was that the loss was not measuring only the ability to perform the sum. The operands before the = sign are generated randomly, so asking the model to predict them introduces a part of the loss that cannot be reduced: there is no pattern that allows it to know which random digit comes next. I fixed this by masking all targets before =, so the cross-entropy only measures the generated result.
This is a small and constrained experiment, but it shows that the architecture can learn an algorithmic pattern and not only the statistical structure of natural language. It also shows how important it is to understand what the loss is actually measuring: a high loss does not always mean that the model is failing at the task we care about.
\nThe second experiment is BookGPT, a character-level language model trained on Italian books. Here the task is more open: given a sequence of characters, the model has to predict the next one and generate prose one character at a time.
\nThe first dataset was composed of 19 books from an Hugging Face dataset, plus I Promessi Sposi. During the experiment I discovered that most of those books were English Gutenberg texts translated automatically into Italian. The corpus contained duplicated fragments, strange sentences and even a book with the wrong title. The model learned these artifacts: the generated text had a recognizable Italian structure, but it also produced many invented words and English names like Sighbury and Mr. Wergomestane.
For this reason I rebuilt the dataset using 15 native Italian works from authors like Manzoni, Verga, Pirandello, Nievo, De Amicis, Deledda and Collodi. The final clean corpus contains around 9.48 million characters and a vocabulary of 140 characters. Every book was split independently into 90% training and 10% validation before concatenating them, to avoid having a validation set biased toward only the last books in the corpus.
\nI trained three different configurations and compared their validation loss and generated text at matching training steps.
\nThe first one was the baseline trained on the old corpus. It had around 852K parameters, an embedding size of 128, 4 heads, 4 blocks and a context length of 256. It reached a validation loss of 1.4355 at step 5000 and 1.3416 at step 8250. Training and validation loss remained very close, so there was no clear sign of overfitting. The generated text started to reproduce Italian syntax, but the bad quality of the corpus was clearly visible in its vocabulary.
\nThe second run used the clean corpus and a larger model, with around 4.09M parameters, an embedding size of 256, 8 heads, 6 blocks and the same context length of 256. This was the best result. At step 2000 its validation loss was 1.5673, compared with 1.7487 for the baseline, and at step 4750 it reached 1.3969, compared with 1.4497. The validation loss continued to decrease, while the larger gap from the training loss stabilized instead of continuing to grow.
\nThe generated text improved too. The model started to generate names present in the corpus, such as zio Trao from I Malavoglia, and it often maintained a plausible agreement between subjects and verbs. However, the text was still not coherent for long. Subjects could disappear or change inside the same generation, and many words were still invented even if they had a plausible Italian morphology.
Since the main limitation seemed to be coherence, I tried a third run with the context length increased from 256 to 512. The hypothesis was that seeing more previous characters would help the model maintain subjects and references for longer. The result was negative: at step 4750 the validation loss was 1.5241, worse than both the baseline at 1.4497 and the larger model at 1.3969. The generated text did not show a clear improvement either. With prompts like Roma or lungo i fiumi, the model produced locally plausible Italian fragments, but the sentence still lost its direction and introduced invented words.
On this hardware the larger context was not free. Without an efficient attention kernel, its memory cost forced me to reduce the batch size from 128 to 64. The model was therefore learning a more difficult task with noisier gradient updates, without receiving more representational capacity. In this setup, investing the available compute in width and depth produced better results than investing it in a longer context.
\nThe main observation from BookGPT is that context length alone does not create coherence. A wider window only gives the model the possibility to look further back; it does not guarantee that it will learn how to use that information. In this experiment, a clean corpus and more model capacity had a much clearer effect than simply doubling the context length.
\nThese experiments made the trade-offs of a GPT architecture much more concrete. By implementing it from scratch I also got a better grasp of its fundamental building block: a transformer block composed of causal self-attention and a feedforward network, together with residual connections and layer normalization. Deriving its computational cost also made it clear which parts scale with the model dimension and which ones become expensive as the context grows.
\nSumGPT showed that even a very small model can learn a precise algorithmic task, but only if the training objective measures the right thing. BookGPT showed that for language modeling the architecture is only one part of the problem: cleaning the dataset and increasing useful model capacity improved the results more than simply giving the model a longer context.
\nThe longer-context run was probably the most useful negative result. A larger context window does not automatically produce better coherence, especially when it forces a smaller batch size and the model does not have enough capacity to use the additional information. On this hardware and at this scale, a context of 256 with a wider and deeper model was a better allocation of compute than a context of 512.
\nThe generated text is still far from coherent prose, but that is consistent with the size of the model and the character-level setup. The useful result was not producing a good language model; it was understanding which changes actually improved it, which ones did not, and why.
","date_published":"2026-07-06T00:00:00.000Z","tags":["Neural Networks","Deep Learning","Transformers","Attention","Python"]},{"id":"https://www.tommasovaccari.com/blog/introduction-to-blas","url":"https://www.tommasovaccari.com/blog/introduction-to-blas","title":"Introduction to BLAS","summary":"An introduction to BLAS routines, the roofline model, and SGEMM optimization","authors":[{"name":"Tommaso Vaccari"}],"content_html":"While reasoning on some ML concepts I just stumbled on an interesting question, who and what is providing fast low level APIs, for example for computing matmul etc?
\nThe answer happened to lie in the BLAS interface definition, the Basic Linear Algebra Subprograms. BLAS defines a standard set of routines for common linear-algebra operations, while libraries such as Apple Accelerate, Intel oneMKL, OpenBLAS, cuBLAS,and BLIS provide optimized implementations. Netlib also provides reference implementations, which prioritize portability and correctness rather than peak performance. As you can see hardware vendors usually ship also their BLAS implementation taylored for the hardware.
\nBLAS was originally specified through Fortran routines. C programs can call the same operations through CBLAS, a C interface with conventions that are more natural in C, such as explicit row-major and column-major layouts.
\nHigh-level machine-learning frameworks rely heavily on operations such as matrix multiplication. Their performance depends on optimized numerical kernels supplied by the framework, hardware vendor, or an external library. These notes grew out of studying BLAS and the roofline model, with the practical goal of writing a custom SGEMM kernel and comparing it with Apple Accelerate's cblas_sgemm.
The complete implementation used in this article is available in the custom-blas repository.
\nBLAS routines are divided into three levels. For dense square inputs of dimension n, their characteristic computational costs are:
| Level | Main operation | Example | Computational cost |
|---|---|---|---|
| 1 | Vector-vector | AXPY | O(n) |
| 2 | Matrix-vector | GEMV | O(n^2) |
| 3 | Matrix-matrix | GEMM | O(n^3) |
I am mainly interested in the third category: matrix-matrix multiplication. The conventional algorithm has a time complexity of O(n^3). Asymptotically faster algorithms exist: Strassen's algorithm, for example, runs in O(n^{\\log_2 7}) \\approx O(n^{2.807}). In practice, conventional blocked GEMM is usually preferred for ordinary matrix sizes because it has smaller constants, better numerical properties, and maps efficiently to modern memory hierarchies. Some libraries use Strassen-like methods selectively for sufficiently large matrices.
BLAS routine names encode the data type and operation. Common precision prefixes are:
\n| Prefix | Data type |
|---|---|
S | Single-precision real (float) |
D | Double-precision real (double) |
C | Single-precision complex |
Z | Double-precision complex |
Many matrix routines then include a code describing the matrix structure and operation:
\n| Code | Meaning |
|---|---|
GE | General matrix |
SY | Symmetric matrix |
TR | Triangular matrix |
GB | General band matrix |
MV | Matrix-vector operation |
MM | Matrix-matrix operation |
For example, SGEMM means single-precision general matrix-matrix multiplication.
The Netlib documentation describes both the cblas_sgemm and Fortran SGEMM interfaces.
The roofline model helps determine whether a kernel is limited by compute throughput or memory bandwidth. It uses the following quantities, that we can define:
\nW: the number of floating-point operations performed, measured in FLOPs.Q: the number of bytes transferred between the relevant levels of the memory hierarchy.I: the ratio W/Q, measured in FLOPs per byte.P: the achieved rate of computation, measured in FLOP/s.For a given architecture, let P_{\\text{peak}} be its peak floating-point throughput and B_{\\text{peak}} its peak memory bandwidth. The roofline bound is
P \\leq \\min\\left(P_{\\text{peak}}, B_{\\text{peak}} I\\right).![\"Roofline](\"https://www.tommasovaccari.com/static/image-1a64c4ae.webp\")
The intersection between the horizontal compute roof and the diagonal bandwidth roof is the ridge point:
\nI_{\\text{ridge}} = \\frac{P_{\\text{peak}}}{B_{\\text{peak}}}.Kernels to the left of this point are bandwidth-bound under the model; kernels to the right are compute-bound.
\nFor an AXPY-like vector operation, both the work and memory traffic scale as O(n), so arithmetic intensity remains O(1). Dense matrix-vector multiplication performs O(n^2) work but must also read O(n^2) matrix elements, so its intensity is also O(1).
Dense matrix-matrix multiplication performs O(n^3) work on O(n^2) input and output data. With a blocked implementation that reuses data effectively, its arithmetic intensity can grow as O(n). This reuse is why GEMM can move into the compute-bound region. A naive loop nest does not automatically achieve the minimum possible memory traffic, so its realized intensity may be much lower.
The goal is to implement several simplified SGEMM variants and compare them with Apple Accelerate:
\ni-j-k loop orderi-k-j loop order with vectorization disabledi-k-j loop order with vectorization enabledThe standard CBLAS signature is:
\nvoid cblas_sgemm(\n CBLAS_ORDER Order,\n CBLAS_TRANSPOSE TransA,\n CBLAS_TRANSPOSE TransB,\n int M,\n int N,\n int K,\n float alpha,\n const float *A,\n int lda,\n const float *B,\n int ldb,\n float beta,\n float *C,\n int ldc\n);It computes
\nC \\leftarrow \\alpha AB + \\beta C.The parameters have the following roles:
\nOrder selects row-major or column-major storage. C supports both; this experiment uses CblasRowMajor.TransA and TransB specify whether each input is used as-is or transposed.M, N, and K define the matrix dimensions: A is M \\times K, B is K \\times N, and C is M \\times N when neither input is transposed.alpha scales AB`, while `beta` scales the previous value of C$.lda, ldb, and ldc are the leading dimensions. For contiguous row-major matrices without padding, they are respectively K, N, and N in the non-transposed case.A × B = C\n(M × K) (K × N) (M × N)To keep the experiment focused, every custom implementation supports only square, row-major, non-transposed matrices and is called with \\alpha = 1 and \\beta = 0:
cblas_sgemm(\n CblasRowMajor, CblasNoTrans, CblasNoTrans,\n n, n, n,\n 1.0f,\n A, n,\n B, n,\n 0.0f,\n C, n\n);The custom functions mimic the CBLAS parameter list but do not implement the complete cblas_sgemm contract.
The naive implementation computes the dot product between each row of A and each column of B:
for (size_t i = 0; i < N; i++) {\n for (size_t j = 0; j < N; j++) {\n float acc = 0.0f;\n for (size_t k = 0; k < N; k++) {\n acc += *(A + lda * i + k) * *(B + ldb * k + j);\n }\n *(C + ldc * i + j) = acc;\n }\n}The accesses to A are contiguous, but the accesses to B use a stride of ldb. This is inefficient for row-major storage because consecutive iterations touch different rows of B.leading to cache misses, that are very costly.
Changing the loop order from i-j-k to i-k-j makes the innermost loop traverse rows of B and C contiguously:
for (size_t i = 0; i < N; i++) {\n for (size_t k = 0; k < N; k++) {\n float a_i_k = *(A + lda * i + k);\n for (size_t j = 0; j < N; j++) {\n *(C + ldc * i + j) += a_i_k * *(B + ldb * k + j);\n }\n }\n}Because this version accumulates into C, the benchmark initializes C to zero before each run. The measured scalar variant also explicitly disables Clang's loop vectorization and unrolling, isolating the effect of loop reordering as far as this compiler configuration allows.
The final custom variant keeps the same memory-access pattern and asks Clang to vectorize and interleave the inner loop:
\nfor (size_t i = 0; i < N; i++) {\n for (size_t k = 0; k < N; k++) {\n float a_i_k = *(A + lda * i + k);\n\n #pragma clang loop vectorize(enable) interleave(enable)\n for (size_t j = 0; j < N; j++) {\n *(C + ldc * i + j) += a_i_k * *(B + ldb * k + j);\n }\n }\n}On the Apple M2, NEON vector registers are 128 bits wide, so one vector can hold four 32-bit floats. This makes the loop suitable for SIMD, although register width alone does not predict the complete speedup: instruction scheduling, fused multiply-add throughput, memory traffic, and compiler decisions also matter.
\nThe benchmark produced the following results on an Apple M2:
\n| Implementation | N = 1024 | N = 2048 |
|---|---|---|
Naive (i-j-k) | 1.72 GFLOP/s | 0.92 GFLOP/s |
Reorganized, scalar (i-k-j, no auto-vectorization) | 6.37 GFLOP/s | 6.38 GFLOP/s |
Reorganized and vectorized (i-k-j, SIMD) | 27.33 GFLOP/s | 26.54 GFLOP/s |
Accelerate cblas_sgemm | 789.81 GFLOP/s | 1035.06 GFLOP/s |
At N = 1024, loop reordering improves performance from 1.72 to 6.37 GFLOP/s, a gain of approximately 3.7\\times. The main difference is the access pattern: the naive version reads B column by column with stride ldb, whereas the reorganized version accesses B and C sequentially in the inner loop.
Enabling vectorization raises performance from 6.37 to 27.33 GFLOP/s, approximately another 4.3\\times. This is consistent with SIMD being effective on the contiguous inner loop, but it should not be interpreted as proof that the speedup comes exclusively from processing four floats per vector instruction.
The total improvement from the naive to the vectorized custom version is approximately 15.9\\times. Accelerate remains much faster: about 29\\times faster at N = 1024 and 39\\times faster at N = 2048.
This remaining gap cannot be closed by a vectorization pragma alone. High-performance GEMM implementations typically combine several techniques:
\nBLAS routines are fundamental building blocks for scientific computing and machine learning. This experiment shows that two understandable changes—improving memory access and enabling SIMD—can produce a substantial speedup, while also demonstrating how much additional engineering separates a simple loop nest from a production GEMM implementation.
\nThe next step is to add cache blocking and a register-level microkernel, then evaluate them with a more rigorous benchmark harness.
","date_published":"2026-06-27T00:00:00.000Z","tags":["HPC","Linear Algebra","C"]},{"id":"https://www.tommasovaccari.com/blog/attention-mechanism-to-small-gpt","url":"https://www.tommasovaccari.com/blog/attention-mechanism-to-small-gpt","title":"From Self-Attention to a Small GPT","summary":"A personal from-scratch derivation of causal self-attention, starting from a character language model and ending with the core pieces of a small GPT.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"This post is a continuation of my notes on language modeling, adapted from the notes that I took while studying the topic.\nThe goal is to start from the simplest character-level model, as explained in other posts, and then build the core ideas that lead to a small GPT-style transformer.
\nThe main focus is the attention mechanism: why it is needed, how causal self-attention works, and how it fits inside a transformer block.
\nThe derivation is grounded in the scaled dot-product and multi-head attention mechanisms introduced by Vaswani et al. in Attention Is All You Need.
\nThe path is:
\nThe objective is to understand what each piece is doing.
\nA language model predicts the next token given the previous tokens:
\nP(x_t \\mid x_1, x_2, \\ldots, x_{t-1})In this project the model is character-based, so every character is mapped to an integer. This keeps the vocabulary small and makes the implementation easier to inspect.
\nThere are more powerful tokenizers, for example:
\nbut for a first implementation a character-level tokenizer is enough.
\nIf the vocabulary is:
\nchars = sorted(list(set(text)))\nstoi = {ch: i for i, ch in enumerate(chars)}\nitos = {i: ch for ch, i in stoi.items()}then encoding and decoding are just table lookups:
\nencode = lambda s: [stoi[c] for c in s]\ndecode = lambda xs: \"\".join(itos[i] for i in xs)The model never sees strings directly. It sees tensors of integer token ids.
\nThis is the first important shift: for the model, characters are not characters. They are numbers. But a raw integer id is not a good\nrepresentation by itself, because the number 10 is not “more character” than the number 3. The id is only an index.
For this reason we use an embedding table. Each token id selects a learned vector, and that vector becomes the internal representation\nof the token.
\nAs humans, we have a strong geometric intuition only up to two or three dimensions. In deep learning, instead, we work with high-\ndimensional spaces. The idea is that a vector with many dimensions can encode many useful features at the same time.
\nIn a character-level model these features are not necessarily human-readable meanings. The model may learn that some characters behave\nlike vowels, that some often appear after others, or that some are common near word endings. The important point is that the\nrepresentation is learned from data.
\nA good way to build intuition for this is to look at pretrained word embeddings. In that case, words with similar meaning often end up\nclose in the embedding space. Classic examples show that relationships such as king - man + woman can point near queen.
Our model is much smaller and works at character level, but the principle is the same: instead of treating tokens as isolated symbols,\nwe map them into vectors that the neural network can modify, combine, and use for prediction.
\nBefore seeing an example, we need another important concept from linear algebra: the standard dot product, also known as the standard inner product.
\nWhen we work with a vector space, we can associate it with an inner product, which is a function that takes two vectors and returns a real number:
\n\\langle \\cdot, \\cdot \\rangle : V \\times V \\to \\mathbb{R}An inner product must satisfy some properties:
\n\\langle v, v \\rangle \\ge 0\\langle v, v \\rangle = 0 only when v = 0When we work in the standard space \\mathbb{R}^n, the standard dot product between two vectors
a = (a_1, a_2, \\ldots, a_n)and
\nb = (b_1, b_2, \\ldots, b_n)is defined as:
\na \\cdot b = \\sum_{i=1}^{n} a_i b_iSo we multiply the entries in the same position and then we sum everything.
\nThere is also a geometric interpretation. If \\|a\\| and \\|b\\| are the lengths of the two vectors, then:
a \\cdot b = \\|a\\| \\|b\\| \\cos(\\theta)where \\theta is the angle between the two vectors.
This is important for embeddings because now we have a way to compare two vectors. If two embedding vectors point in a similar direction, then intuitively they are carrying similar information. If they point in very different directions, then they are representing something different.
\nThe problem is that the raw dot product also depends on the length of the vectors. So if we want to focus only on the direction, we use cosine similarity:
\n\\mathrm{cosine\\_similarity}(a,b) =\n\\frac{a \\cdot b}{\\|a\\|\\|b\\|}The value is between -1 and 1:
1, the two vectors point in a similar direction0, the two vectors are almost orthogonal-1, the two vectors point in opposite directionsFor embeddings, this gives us a practical measure of similarity. If two words have similar meaning, we expect their vectors to have a high cosine similarity.
\n![\"Cosine](\"https://www.tommasovaccari.com/static/cosine-similarity-vectors-cc32c236.webp\")
In the next example, we can use pretrained GloVe embeddings and test this directly with PyTorch.
\nReference: GloVe: Global Vectors for Word Representation.
\nglove = GloVe(name=\"6B\", dim=100);\n\ndef cosine(a, b):\n return F.cosine_similarity(glove[a], glove[b], dim=0).item()\n\nprint(\"king / queen:\", cosine(\"king\", \"queen\"))\nprint(\"king / man:\", cosine(\"king\", \"man\"))\nprint(\"queen / woman:\", cosine(\"queen\", \"woman\"))\nprint(\"man / woman:\", cosine(\"man\", \"woman\"))\n\ntarget = glove[\"king\"] - glove[\"man\"] + glove[\"woman\"]\n# Compute the cosine similarity between the target vector and all vectors\n# in the GloVe vocabulary\nscores = F.cosine_similarity(target.unsqueeze(0), glove.vectors)\n\nexclude = {\"king\", \"man\", \"woman\"}\nfor word in exclude:\n scores[glove.stoi[word]] = -1\n# Now we can see if it's effectively close to \"queen\"\nbest_score, best_idx = torch.topk(scores, k=5)\nprint(\"Top 5 closest words to 'king - man + woman':\")\nfor i in range(5):\n print(f\" {glove.itos[best_idx[i]]}: {best_score[i].item():.4f}\")This produces:
\nTop 5 closest words to 'king - man + woman':\n queen: 0.7834\n monarch: 0.6934\n throne: 0.6833\n daughter: 0.6809\n prince: 0.6713When training a transformer, we do not push the entire dataset into the model at once. We sample fixed-length chunks.
\nIf block_size = 8, then each row of the batch contains 8 input tokens and 8 target tokens. The target sequence is just the input sequence shifted by one position:
x = data[i:i + block_size]\ny = data[i + 1:i + block_size + 1]So for one row:
\nx: [18, 47, 56, 57, 1, 58, 46, 43]\ny: [47, 56, 57, 1, 58, 46, 43, 1]This contains multiple training examples at the same time. At position t, the model should predict y[t] using only x[:t+1].
The batch dimension allows parallelization:
\ndef get_batch(split):\n data = train_data if split == \"train\" else val_data\n ix = torch.randint(len(data) - block_size, (batch_size,))\n x = torch.stack([data[i:i + block_size] for i in ix])\n y = torch.stack([data[i + 1:i + block_size + 1] for i in ix])\n return x, yThe shape is:
\nx: (B, T)\ny: (B, T)where B is the batch size and T is the context length.
The simplest neural language model is a bigram model.
\nIt predicts the next token using only the current token:
\nP(x_{t+1} \\mid x_t)In PyTorch this can be implemented with an embedding table of shape:
\nvocab_size x vocab_sizeWhen we pass an integer token id, the embedding table returns the corresponding row. That row contains the logits for the next character.
\nclass BigramLanguageModel(nn.Module):\n def __init__(self, vocab_size):\n super().__init__()\n self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)\n\n def forward(self, idx, targets=None):\n logits = self.token_embedding_table(idx)\n\n if targets is None:\n return logits, None\n\n B, T, C = logits.shape\n logits = logits.view(B * T, C)\n targets = targets.view(B * T)\n loss = F.cross_entropy(logits, targets)\n return logits, lossThe loss is cross entropy. Internally, cross entropy applies softmax to convert logits into probabilities and then computes the negative log likelihood of the correct target.
\nIf the target probability is high, the loss is low. If the target probability is low, the loss is high.
\nGeneration repeatedly asks the model for the next-token distribution and samples from it.
\ndef generate(self, idx, max_new_tokens):\n for _ in range(max_new_tokens):\n logits, _ = self(idx)\n logits = logits[:, -1, :]\n probs = F.softmax(logits, dim=-1)\n idx_next = torch.multinomial(probs, num_samples=1)\n idx = torch.cat((idx, idx_next), dim=1)\n return idxAt this stage, looking only at the last time step seems wasteful because the model computes logits for all positions. But this structure becomes useful once the model starts using the whole context.
\nThe bigram model has a hard limitation: every token is predicted from one previous token only.
\nThis is not enough for language. The meaning of a token depends on the previous context. A token should be able to receive information from previous tokens.
\nThe constraint is causal:
\npast tokens -> current token\nfuture tokens are hiddenSo information can flow from left to right, but not from right to left.
\nThe simplest way to let tokens communicate is to average the previous token vectors.
\nIf x has shape:
(B, T, C)then for each position t, we want:
\\text{out}_{b,t} = \\frac{1}{t+1}\\sum_{i=0}^{t} x_{b,i}This gives each token a \"bag of previous words\" representation.
\nThe naive implementation is a double loop:
\nout = torch.zeros_like(x)\nfor b in range(B):\n for t in range(T):\n xprev = x[b, :t + 1]\n out[b, t] = torch.mean(xprev, 0)The same operation can be written as matrix multiplication.
\nWe create a lower triangular matrix:
\ntril = torch.tril(torch.ones(T, T))\nwei = tril / tril.sum(1, keepdim=True)\nout = wei @ xThe matrix wei contains the averaging weights. Since it is lower triangular, each token can only aggregate information from itself and from previous tokens.
This is the first intuition for attention: each token aggregates information from previous tokens using a weight matrix.
\nThe limitation is that these weights are fixed. They do not depend on the actual content of the tokens.
\nSelf-attention makes the aggregation weights data-dependent.
\nEach token starts with a learnable private representation x that contains:
Then each token produces three vectors:
\nIn code:
\nkey = nn.Linear(n_embd, head_size, bias=False)\nquery = nn.Linear(n_embd, head_size, bias=False)\nvalue = nn.Linear(n_embd, head_size, bias=False)Given:
\nx # (B, T, C)\nq = query(x) # (B, T, head_size)\nk = key(x) # (B, T, head_size)\nv = value(x) # (B, T, head_size)we compute the affinity between tokens using the dot product between queries and keys:
\nwei = q @ k.transpose(-2, -1)The result has shape:
\n(B, T, T)For every token at position i, wei[i, j] tells how much token i wants to receive from token j.
In language modeling, token i cannot receive information from future tokens.
So we mask the upper triangular part of the attention matrix:
\nwei = wei.masked_fill(tril == 0, float(\"-inf\"))\nwei = F.softmax(wei, dim=-1)After softmax, every row is a probability distribution over the allowed previous positions.
\nThe output is:
\nout = wei @ vSo every token receives a weighted sum of the value vectors of previous tokens.
\nFor a fixed pair of tokens i, j, the raw attention score is:
\\mathrm{score}_{i,j} = q_i \\cdot k_j = \\sum_{h=1}^{d} q_{i,h}k_{j,h}where:
\nd = \\mathrm{head\\_size}Assume for intuition that:
\nq_{i,h} \\sim \\mathcal{N}(0,1), \\quad k_{j,h} \\sim \\mathcal{N}(0,1)and that the components are independent.
\nFor each term:
\nX_h = q_{i,h}k_{j,h}we have:
\n\\mathbb{E}[X_h] = 0and:
\n\\mathrm{Var}(X_h) = 1The score is a sum of d independent terms:
\\mathrm{score}_{i,j} = \\sum_{h=1}^{d} X_hso:
\n\\mathrm{Var}(\\mathrm{score}_{i,j}) = dand the standard deviation grows like:
\n\\sqrt{d}If d is large, the values entering softmax can become large in magnitude. Then softmax becomes very peaky, often close to a one-hot vector.
To keep the variance stable, we divide by:
\n\\sqrt{d}The scaled attention score is:
\n\\frac{q_i \\cdot k_j}{\\sqrt{d}}and now:
\n\\mathrm{Var}\\left(\\frac{\\mathrm{score}_{i,j}}{\\sqrt{d}}\\right)\n=\n\\frac{1}{d}\\mathrm{Var}(\\mathrm{score}_{i,j})\n=\n1This is scaled dot-product attention.
\nPutting the pieces together:
\nclass Head(nn.Module):\n def __init__(self, head_size):\n super().__init__()\n self.key = nn.Linear(n_embd, head_size, bias=False)\n self.query = nn.Linear(n_embd, head_size, bias=False)\n self.value = nn.Linear(n_embd, head_size, bias=False)\n self.register_buffer(\"tril\", torch.tril(torch.ones(block_size, block_size)))\n self.dropout = nn.Dropout(dropout)\n\n def forward(self, x):\n B, T, C = x.shape\n k = self.key(x)\n q = self.query(x)\n wei = q @ k.transpose(-2, -1) * C**-0.5\n wei = wei.masked_fill(self.tril[:T, :T] == 0, float(\"-inf\"))\n wei = F.softmax(wei, dim=-1)\n wei = self.dropout(wei)\n v = self.value(x)\n out = wei @ v\n return outThere is one important correction in this code: the scaling factor should use the head dimension, not the original embedding dimension. A clearer version is:
\nwei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5The output shape is:
\n(B, T, head_size)Each token now carries information aggregated from the previous tokens according to learned, data-dependent weights.
\nA single head gives one attention pattern.
\nMulti-head attention runs several heads in parallel:
\nclass MultiHeadAttention(nn.Module):\n def __init__(self, num_heads, head_size):\n super().__init__()\n self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n self.proj = nn.Linear(num_heads * head_size, n_embd)\n self.dropout = nn.Dropout(dropout)\n\n def forward(self, x):\n out = torch.cat([h(x) for h in self.heads], dim=-1)\n out = self.dropout(self.proj(out))\n return outIf n_embd = 384 and num_heads = 6, then each head can have:
head_size = 384 / 6 = 64The idea is that different heads can specialize in different relationships between tokens.
\nAs you can see, we are processing the heads separately. What we want to do now is process all heads in parallel, exploiting matrix multiplication and treating the head axis as an additional batch-like dimension.
\nIn the previous multi-head implementation, each head was composed of three matrices created inside the single head. The idea is to still create the three projections, but originate them from one bigger matrix and then act as if they were separated by head:
\nclass MultiHeadParallel(nn.Module):\n def __init__(self, config):\n super().__init__()\n assert config.n_embd % config.n_head == 0\n # key, query, value projections for all heads, but in a batch\n self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)\n # output projection\n self.c_proj = nn.Linear(config.n_embd, config.n_embd)\n # regularization\n self.attn_dropout = nn.Dropout(config.attn_pdrop)\n self.resid_dropout = nn.Dropout(config.resid_pdrop)\n # causal mask to ensure that attention is only applied to the left in the input sequence\n self.register_buffer(\"bias\", torch.tril(torch.ones(config.block_size, config.block_size))\n .view(1, 1, config.block_size, config.block_size))\n self.n_head = config.n_head\n self.n_embd = config.n_embd\n\n def forward(self, x):\n B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)\n\n # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n q, k ,v = self.c_attn(x).split(self.n_embd, dim=2)\n k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n\n # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))\n att = F.softmax(att, dim=-1)\n att = self.attn_dropout(att)\n y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n\n # output projection\n y = self.resid_dropout(self.c_proj(y))\n return yThis implementation is adapted from Karpathy's minGPT repository.\nThe interesting part, in my opinion, is that we are adding a head dimension and using it to compute multiple heads in parallel. At the end, we put everything back together to obtain the (B, T, C) matrix.
After attention, each token has received information from the previous context.
\nThen we let each token process its own representation independently with a feed-forward network:
\nclass FeedForward(nn.Module):\n def __init__(self, n_embd):\n super().__init__()\n self.net = nn.Sequential(\n nn.Linear(n_embd, 4 * n_embd),\n nn.ReLU(),\n nn.Linear(4 * n_embd, n_embd),\n nn.Dropout(dropout),\n )\n\n def forward(self, x):\n return self.net(x)Attention is the communication step. The feed-forward network is the computation step in which we let the channels of each token's internal representation mix together.
\nA transformer block combines:
\nclass Block(nn.Module):\n def __init__(self, n_embd, n_head):\n super().__init__()\n head_size = n_embd // n_head\n self.sa = MultiHeadAttention(n_head, head_size)\n self.ffwd = FeedForward(n_embd)\n self.ln1 = nn.LayerNorm(n_embd)\n self.ln2 = nn.LayerNorm(n_embd)\n\n def forward(self, x):\n x = x + self.sa(self.ln1(x))\n x = x + self.ffwd(self.ln2(x))\n return xThe residual connections give gradients a direct path through the network. This is important because the model becomes deep once we stack multiple transformer blocks.
\nThe layer normalization standardizes each token along the feature dimension. In this implementation it is applied before attention and before the feed-forward network.
\nThis is called a pre-norm transformer block.
\nThe full model uses:
\nclass GPTLanguageModel(nn.Module):\n def __init__(self):\n super().__init__()\n self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n self.position_embedding_table = nn.Embedding(block_size, n_embd)\n self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n self.ln_f = nn.LayerNorm(n_embd)\n self.lm_head = nn.Linear(n_embd, vocab_size)\n\n def forward(self, idx, targets=None):\n B, T = idx.shape\n tok_emb = self.token_embedding_table(idx)\n pos_emb = self.position_embedding_table(torch.arange(T, device=idx.device))\n x = tok_emb + pos_emb\n x = self.blocks(x)\n x = self.ln_f(x)\n logits = self.lm_head(x)\n\n if targets is None:\n loss = None\n else:\n B, T, C = logits.shape\n logits = logits.view(B * T, C)\n targets = targets.view(B * T)\n loss = F.cross_entropy(logits, targets)\n\n return logits, lossThe token embedding tells the model what token it is looking at. The positional embedding tells the model where the token is in the context window.
\nWithout positional embeddings, the transformer has no direct notion of order.
\nA compact way to see the full path is:
\nThis mechanism is called self-attention because queries, keys, and values all come from the same input x.
In cross-attention, the queries usually come from one sequence, while keys and values come from another sequence.
\nFor example, in an encoder-decoder transformer:
\nSo cross-attention lets one sequence attend to information produced by another sequence.
\nIn a GPT-style decoder-only model, we only use causal self-attention.
\nThe attention mechanism is a learned, data-dependent way to move information between tokens.
\nThe causal mask makes it usable for language modeling because each token can only look at the past. Scaling keeps the softmax numerically well-behaved. Multi-head attention lets the model learn several communication patterns at the same time.
\nOnce attention is combined with feed-forward layers, residual connections, dropout, and layer normalization, we obtain the basic transformer block used in a small GPT.
","date_published":"2026-06-22T00:00:00.000Z","tags":["Neural Networks","Deep Learning","Transformers","Attention","Python"]},{"id":"https://www.tommasovaccari.com/blog/neural-network-training-statistics-initialization-saturation-diagnostics","url":"https://www.tommasovaccari.com/blog/neural-network-training-statistics-initialization-saturation-diagnostics","title":"Neural Network Training: Statistics, Initialization and Diagnostics","summary":"A structured study of neural network training dynamics, covering statistical preliminaries, loss interpretation, saturation, gradient flow, initialization schemes, diagnostic tools, and stabilization mechanisms.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"This post comes from a cleaned-up version of the notes I wrote while studying neural network training. The goal is to turn those notes into something more readable, without losing the practical and exploratory style of the original work.
\nIn these notes I dive into a set of concepts that are fundamental for understanding how neural networks work and why they behave in a certain manner during training. These ideas are not only useful for building intuition, but also essential for debugging models, diagnosing training instabilities, and making optimization behave properly.
\nThe topics covered are:
\n\n\nThe code and plots in this post are developed from a notebook, available in this GitHub repository: Neural-Network-Statistic-and-diagnostic-tools.
\n
Firstly we have to introduce some basic concepts, starting from the definition of a random variable.
\nWe define X as a random variable when X represents the possible outcomes of a random process. This allows us to model uncertainty mathematically and to reason about events such as the probability that X falls inside a certain range.
Random variables fall into two main categories:
\nThe main difference lies in the values that X can assume. A discrete random variable can assume values from a finite or countable set, while a continuous random variable can assume infinitely many values inside an interval.
Every distribution is characterized by functions that describe how probability is assigned. For continuous random variables, the most important one is the probability density function, or pdf.
\nThe pdf can be misleading at first, because it does not directly give the probability that X is equal to a specific value. Instead, it describes density. Probabilities are obtained by integrating the density over an interval:
P(a \\le X \\le b) = \\int_a^b f_X(x)\\,dxThis also means that the probability of a single exact value is zero:
\nP(X = a) = \\int_a^a f_X(x)\\,dx = 0So the intuitive interpretation is that probability is the area under the pdf curve.
\nThe cumulative distribution function, or cdf, accumulates this area from -\\infty up to a value x:
F_X(x) = P(X \\le x) = \\int_{-\\infty}^{x} f_X(t)\\,dtUsually in machine learning we are interested in continuous random variables, and one of the most important examples is the Gaussian distribution.
\nFor a Gaussian random variable X \\sim \\mathcal{N}(\\mu, \\sigma^2), the probability density function is:
f_X(x) = \\frac{1}{\\sigma \\sqrt{2\\pi}}\ne^{-\\frac{(x - \\mu)^2}{2\\sigma^2}}In this case, the probability that X falls between a and b is:
P(a \\le X \\le b) =\n\\int_a^b\n\\frac{1}{\\sigma \\sqrt{2\\pi}}\ne^{-\\frac{(x - \\mu)^2}{2\\sigma^2}}\\,dxThe parameters \\mu and \\sigma^2 are respectively the mean and the variance of the distribution. The mean describes where the distribution is centered, while the variance describes how spread out the values are around the mean.
Their theoretical definitions are:
\n\\mu = \\mathbb{E}[X] = \\int_{-\\infty}^{+\\infty} x f_X(x)\\,dx\\sigma^2 = \\mathrm{Var}(X) = \\mathbb{E}[(X - \\mu)^2]\n= \\int_{-\\infty}^{+\\infty} (x - \\mu)^2 f_X(x)\\,dxWhen we work with data, we usually do not know the true distribution parameters. We only have samples, so we estimate the mean and variance from them.
\nGiven samples x_1, x_2, \\dots, x_n, the sample mean is:
\\bar{x} = \\frac{1}{n}\\sum_{i=1}^{n} x_iThe sample variance is:
\ns^2 = \\frac{1}{n - 1}\\sum_{i=1}^{n}(x_i - \\bar{x})^2So basically the mean is derived from averaging the values, and the variance is derived by seeing how each value differs from the mean. We square these differences, average them, and obtain a measure of how spread out the samples are around the mean. If you have noticed that in the second case we do not average by n but by n - 1, I suggest you read this explanation.
Now that we have defined these basic concepts, we can use Python and PyTorch to build some intuition through examples. We can use these tools to generate samples from a standard Gaussian distribution:
\nplt.figure(figsize=(12, 8))\nbins = torch.linspace(-4, 4, 51).numpy()\n\nfor idx, i in enumerate((100, 1000, 10000, 100000), start=1):\n x = torch.randn(i)\n mean = x.mean()\n variance = x.var(unbiased=True)\n\n plt.subplot(2, 2, idx)\n plt.hist(x.numpy(), bins=bins, density=True, alpha=0.7, color=\"g\")\n plt.title(f\"n = {i}\\nmean = {mean:.3f}, var = {variance:.3f}\")\n\nplt.tight_layout()\nplt.show()![\"Sample](\"https://www.tommasovaccari.com/static/sample-size-histograms-e2fbea40.webp\")
As we can see, with a small number of samples, the mean and variance can deviate significantly from their true values (0 and 1 for a standard normal distribution). As we increase the number of samples, the estimates become more accurate, and the histogram approaches the shape of the true Gaussian distribution. If you are interested to delve deeper, this is related to the law of large numbers and the central limit theorem.
\nNow we can see intuitively how mean and variance modify the shape of the bell curve, that is, the plot of the Gaussian distribution over a histogram. The mean shifts the center of the bell curve, while the variance controls how wide or narrow it is. A higher variance results in a wider bell curve, while a lower variance results in a narrower bell curve.
\nplt.figure(figsize=(12, 10))\nbins = torch.linspace(-8, 8, 80).numpy()\n\nplt.subplot(2, 1, 1)\nfor mu in (-3, 0, 3):\n x = torch.randn(10000) + mu\n plt.hist(x.numpy(), bins=bins, density=True, alpha=0.5, label=f\"mu = {mu}\")\n\nplt.title(\"Fixed variance, changing mean\")\nplt.legend()\n\nplt.subplot(2, 1, 2)\nfor sigma in (0.5, 1, 2):\n x = torch.randn(10000) * sigma\n plt.hist(x.numpy(), bins=bins, density=True, alpha=0.5, label=f\"sigma = {sigma}\")\n\nplt.title(\"Fixed mean, changing variance\")\nplt.legend()\n\nplt.tight_layout()\nplt.show()![\"Mean](\"https://www.tommasovaccari.com/static/mean-variance-histograms-1434e153.webp\")
Now that we have built some intuition about these concepts, we can look at what happens when we apply a nonlinear transformation to samples drawn from a random variable. In our context, we want to see how an initially Gaussian distribution changes when we apply the tanh activation function.
\nThe tanh activation function is defined as follows:
\n\\tanh(x) = \\frac{e^x - e^{-x}}{e^x + e^{-x}}Its output is bounded between -1 and 1:
\\tanh(x) \\in (-1, 1)![\"Tanh](\"https://www.tommasovaccari.com/static/tanh-function-6350b865.webp\")
If we apply tanh to every sample and then plot the result, we can see that the output distribution is no longer Gaussian. Its mean and variance are changed by the nonlinearity:
\n![\"Tanh](\"https://www.tommasovaccari.com/static/tanh-distribution-transformations-6d53dcfe.webp\")
The most important thing to grasp is that the output distribution depends strongly on the input variance. When the variance is large, many input values fall far from the mean. Since tanh is almost flat for large positive or negative values, those samples are pushed close to -1 or 1. This is called tanh saturation, and we will come back to it later.
On the other hand, when the variance is small, most values stay close to 0. Around zero, tanh is almost linear, so the transformed values also remain close to zero. You can see it respectively in row 2 and row 3 of the image.
To follow along during the explanation we are going to use a simple neural network from the previos post. The context is the same, we have a datasets composed of names and we are trying to train a bigram model to learn how to generate italian names.
\n# MLP \nn_embd = 10 # the dimensionality of the character embedding vectors\nn_hidden = 200 # the number of neurons in the hidden layer of the MLP\n\ng = torch.Generator().manual_seed(2147483647) # for reproducibility\nC = torch.randn((vocab_size, n_embd), generator=g)\nW1 = torch.randn((n_embd * block_size, n_hidden), generator=g)#* ((5/3)/(n_embd*block_size)**0.5)\nb1 = torch.randn(n_hidden, generator=g)#*0.01\nW2 = torch.randn((n_hidden, vocab_size), generator=g)*0.01\nb2 = torch.randn(vocab_size, generator=g)*0\n\n\nparameters = [C, W1, b1, W2, b2]\nprint(sum(list((p.nelement() for p in parameters)))) # number of parameters in total\nfor p in parameters:\n p.requires_grad = True\n# same optimization as last time\nmax_steps = 200000\nbatch_size = 32\nlossi = []\n\nfor i in range(max_steps):\n \n # minibatch construct\n ix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g)\n Xb, Yb = Xtr[ix], Ytr[ix] # batch X,Y\n \n # forward pass\n emb = C[Xb] # embed the characters into vectors\n embcat = emb.view(emb.shape[0], -1) # concatenate the vectors\n hpreact = embcat @ W1 #+ b1 # hidden layer pre-activation\n\n h = torch.tanh(hpreact) # hidden layer\n logits = h @ W2 + b2 # output layer\n loss = F.cross_entropy(logits, Yb) # loss function\n \n # backward pass\n for p in parameters:\n p.grad = None\n loss.backward()\n \n # update\n lr = 0.1 if i < 100000 else 0.01 # step learning rate decay\n for p in parameters:\n if p.grad is not None:\n p.data += -lr * p.grad\n \n # track stats\n if i % 10000 == 0: # print every once in a while\n print(f'{i:7d}/{max_steps:7d}: {loss.item():.4f}')\n lossi.append(loss.log10().item())\n![\"Forward](\"https://www.tommasovaccari.com/static/forward-pass-network-753d96ab.webp\")
As you can see the model and the forward pass are pretty simple, we have the following:
\nix: random indices used to sample a minibatch from the training set. Shape: (batch_size,).Xb: input batch, obtained with Xtr[ix]. Each row contains the context characters used to predict the next one. Shape: (batch_size, block_size).Yb: target batch, obtained with Ytr[ix]. Each element is the correct next character index. Shape: (batch_size,).emb = C[Xb]: embedding lookup. Each character index in Xb is replaced by its learned vector from C. Shape: (batch_size, block_size, n_embd).embcat = emb.view(emb.shape[0], -1): concatenation of the context embeddings into a single vector per example. Shape: (batch_size, block_size * n_embd).hpreact = embcat @ W1: hidden layer pre-activation, before applying the nonlinearity. Shape: (batch_size, n_hidden).h = torch.tanh(hpreact): hidden layer activation. Shape: (batch_size, n_hidden).logits = h @ W2 + b2: raw output scores, one score for each possible next character. Shape: (batch_size, vocab_size).loss = F.cross_entropy(logits, Yb): compares the logits with the correct targets and returns a scalar loss. Shape: ().At the end of the network we obtain logits, one raw score for each possible next character. These are not probabilities yet. They can be any real number, positive or negative.
In PyTorch, F.cross_entropy(logits, Yb) expects these raw logits directly.
Under the hood, PyTorch's cross-entropy combines two operations:
\nlog_probs = F.log_softmax(logits, dim=1)\nloss = F.nll_loss(log_probs, Yb)So when we write:
\nloss = F.cross_entropy(logits, Yb)we are already doing the stable version of softmax + negative log-likelihood.
\nFor a single example, let:
\nz \\in \\mathbb{R}^{V}be the logits, where V is the vocabulary size. Softmax turns these raw scores into probabilities:
p_i = \\frac{e^{z_i}}{\\sum_{j=1}^{V} e^{z_j}}This does two things: it makes every value positive, and it normalizes all values so that they sum to one:
\n\\sum_{i=1}^{V} p_i = 1So after softmax we can interpret p_i as the probability assigned to token i.
Log-softmax does the same transformation, but directly in log space:
\n\\log p_i = z_i - \\log\\left(\\sum_{j=1}^{V} e^{z_j}\\right)This is the numerically stable quantity used by the loss. If the correct target class is y, the negative log-likelihood is:
\\mathcal{L} = -\\log p_ySo the loss is small when the model assigns high probability to the correct token, and large when it assigns low probability to it.
\nSoftmax is still useful as intuition. If the logits are close to each other, the output distribution is close to uniform. If one logit is much larger than the others, softmax puts most of the mass on that class.
\nSo the scale of the final logits matters. Small logits usually mean a softer, less confident output distribution. Very large logits usually mean a very sharp, overconfident distribution.
\ng_softmax = torch.Generator().manual_seed(2147483647)\n\nfor scale in (0.1, 1.0, 10.0):\n logits = torch.randn(4, generator=g_softmax) * scale\n probs = torch.softmax(logits, dim=0)\n\n print(f\"scale = {scale}\")\n print(\"logits:\", logits.round(decimals=3))\n print(\"probs :\", probs.round(decimals=3))\n print()scale = 0.1\nlogits: tensor([-0.0980, -0.1660, -0.0060, -0.0340])\nprobs : tensor([0.2440, 0.2280, 0.2680, 0.2600])\n\nscale = 1.0\nlogits: tensor([ 2.0990, 0.8960, 0.3380, -0.2090])\nprobs : tensor([0.6360, 0.1910, 0.1090, 0.0630])\n\nscale = 10.0\nlogits: tensor([ -9.3910, -11.8100, -6.2840, -9.8260])\nprobs : tensor([0.0420, 0.0040, 0.9280, 0.0270])This is why the magnitude of the final layer matters at initialization. At the beginning of training, the model has not learned anything yet, so it should not be very confident about any particular token. Its predictions should be close to uniform, and the loss should start near the baseline log(vocab_size):
expected_initial_loss = torch.tensor(vocab_size).float().log()\nprint(\"vocab_size:\", vocab_size)\nprint(\"expected initial loss:\", expected_initial_loss.item())vocab_size: 27\nexpected initial loss: 3.295836925506592If the actual initial loss is much higher than this value, something is wrong with the initialization. The model is not just random: it is confidently random. It assigns too much probability to wrong classes, and cross-entropy punishes that heavily.
\nThis is exactly what creates the initial hockey-stick shape in the training loss curve:
\n![\"Previous](\"https://www.tommasovaccari.com/static/lossi-25b58328.webp\")
The loss starts too high, then quickly drops once the model learns to reduce its overconfident random predictions. But this early drop is not meaningful learning: it is mostly the model correcting a bad output scale.
\nA simple fix is to initialize the last layer with smaller weights:
\nW2 = torch.randn((n_hidden, vocab_size), generator=g) * 0.01\nb2 = torch.randn(vocab_size, generator=g) * 0We can see the effect directly by comparing a large output-layer scale with a smaller one:
\ng_compare = torch.Generator().manual_seed(2147483647)\nbatch_size = 32\nix = torch.randint(0, Xtr.shape[0], (batch_size,), generator=g_compare)\nXb, Yb = Xtr[ix], Ytr[ix]\n\nC_test = torch.randn((vocab_size, n_embd), generator=g_compare)\nW1_test = torch.randn((n_embd * block_size, n_hidden), generator=g_compare) * ((5/3) / (n_embd * block_size)**0.5)\n\nemb = C_test[Xb]\nembcat = emb.view(emb.shape[0], -1)\nhpreact = embcat @ W1_test\nh = torch.tanh(hpreact)\n\nfor scale in (1.0, 0.01):\n W2_test = torch.randn((n_hidden, vocab_size), generator=g_compare) * scale\n b2_test = torch.zeros(vocab_size)\n logits = h @ W2_test + b2_test\n loss = F.cross_entropy(logits, Yb)\n\n print(f\"W2 scale = {scale}\")\n print(\"logits std:\", logits.std().item())\n print(\"loss:\", loss.item())\n print()W2 scale = 1.0\nlogits std: 10.497655868530273\nloss: 25.23309326171875\n\nW2 scale = 0.01\nlogits std: 0.11374638229608536\nloss: 3.3066704273223877This keeps the initial logits close to zero. When logits are close to zero, softmax would produce probabilities close to uniform, and the initial loss starts near the theoretical baseline.
\nNow we can train again with this initialization and plot the loss curve. To make the trend easier to read, we average the loss every 1000 iterations:
\n# Now we can plot the loss curve to see how it evolves during training\nplt.figure(figsize=(10, 5))\nlossi_avg = torch.tensor(lossi).view(-1, 1000).mean(1)\nplt.plot(lossi_avg, label=\"Training Loss (log10)\")\nplt.axvline(100, color=\"r\", linestyle=\"--\", label=\"learning-rate decay\")\nplt.xlabel(\"1000-iteration blocks\")\nplt.ylabel(\"log10(loss)\")\nplt.legend()\nplt.show()![\"Training](\"https://www.tommasovaccari.com/static/training-loss-moving-average-a9ec4333.webp\")
Each point is the average of 1000 consecutive training iterations. The vertical line is at block 100, which corresponds to iteration 100000, where the learning rate changes from 0.1 to 0.01. With the smaller last-layer initialization, the curve no longer starts with the same artificial hockey-stick drop.
Now we want to understand what happens under the hood when we call:
\nloss.backward()The loss is just a function of all the parameters of the network:
\n\\mathcal{L} = f(C, W_1, b_1, W_2, b_2)Training means changing these parameters in the direction that reduces the loss. For each parameter \\theta, PyTorch computes:
\\frac{\\partial \\mathcal{L}}{\\partial \\theta}This tells us how much the loss changes when that parameter changes a little.
\nThen the update step is:
\n\\theta \\leftarrow \\theta - \\eta \\frac{\\partial \\mathcal{L}}{\\partial \\theta}where \\eta is the learning rate.
In code, this is the part:
\n# backward pass\nfor p in parameters:\n p.grad = None\nloss.backward()\n\n# update\nlr = 0.1 if i < 100000 else 0.01 # step learning rate decay\nfor p in parameters:\n if p.grad is not None:\n p.data += -lr * p.gradFirst we reset the old gradients, then .backward() computes the new gradients, and finally we move every parameter in the opposite direction of its gradient.
The key mathematical tool is the chain rule. If a function is built by composing smaller functions:
\ny = f(g(x))then:
\n\\frac{dy}{dx}\n=\n\\frac{dy}{dg}\n\\frac{dg}{dx}Backpropagation is just repeated chain rule applied from the loss backward through the network.
\nTo make the mechanism clear, we do not start from the full neural network. We start from a single simplified block:
\nz = wx + by = \\tanh(z)\\mathcal{L} = \\ell(y)So the full computation is a composition of functions:
\nx\n\\xrightarrow{z = wx + b}\nz\n\\xrightarrow{y = \\tanh(z)}\ny\n\\xrightarrow{\\mathcal{L} = \\ell(y)}\n\\mathcal{L}During the forward pass we go from left to right. During backpropagation we go from right to left.
\nThe first gradient comes from the loss itself:
\n\\frac{\\partial \\mathcal{L}}{\\partial y}Since:
\n\\mathcal{L} = \\ell(y)we have:
\n\\frac{\\partial \\mathcal{L}}{\\partial y}\n=\n\\ell'(y)This depends on the specific loss function. In the previous section we used cross-entropy. For softmax + cross-entropy, the gradient with respect to the logits has the clean form:
\n\\frac{\\partial \\mathcal{L}}{\\partial z_i}\n=\np_i - \\mathbb{1}_{i=y}Here, in the simplified example, we just call the first gradient:
\ng_y = \\frac{\\partial \\mathcal{L}}{\\partial y}Now we move one step backward. Since:
\ny = \\tanh(z)we want to know how the loss changes when z changes. This is a two-function composition:
z \\rightarrow y \\rightarrow \\mathcal{L}So the chain rule gives:
\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial y}\n\\frac{\\partial y}{\\partial z}The local derivative of tanh is:
\n\\frac{\\partial y}{\\partial z}\n=\n1 - \\tanh^2(z)Since y = \\tanh(z), this is also:
\\frac{\\partial y}{\\partial z}\n=\n1 - y^2Therefore:
\ng_z\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial y}\n(1 - y^2)\n=\n\\ell'(y)(1 - y^2)This quantity is the gradient that arrives at the linear part.
\nNow we move one more step backward. The linear part is:
\nz = wx + bThere are three things we may care about: w, b, and x.
First, for the weight w:
w \\rightarrow z \\rightarrow \\mathcal{L}so:
\n\\frac{\\partial \\mathcal{L}}{\\partial w}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n\\frac{\\partial z}{\\partial w}Since:
\n\\frac{\\partial z}{\\partial w} = xwe get:
\n\\frac{\\partial \\mathcal{L}}{\\partial w}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}xFor the bias b:
b \\rightarrow z \\rightarrow \\mathcal{L}so:
\n\\frac{\\partial \\mathcal{L}}{\\partial b}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n\\frac{\\partial z}{\\partial b}Since:
\n\\frac{\\partial z}{\\partial b} = 1we get:
\n\\frac{\\partial \\mathcal{L}}{\\partial b}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}For the input x:
x \\rightarrow z \\rightarrow \\mathcal{L}so:
\n\\frac{\\partial \\mathcal{L}}{\\partial x}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n\\frac{\\partial z}{\\partial x}Since:
\n\\frac{\\partial z}{\\partial x} = wwe get:
\n\\frac{\\partial \\mathcal{L}}{\\partial x}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial z}wSo the full backward pass for this small block is:
\ng_y = \\frac{\\partial \\mathcal{L}}{\\partial y}g_z = \\frac{\\partial \\mathcal{L}}{\\partial z}\\frac{\\partial \\mathcal{L}}{\\partial w} = g_z x\\frac{\\partial \\mathcal{L}}{\\partial b} = g_z\\frac{\\partial \\mathcal{L}}{\\partial x} = g_z wThis is the core idea. Each operation receives a gradient from the operation after it, multiplies it by its own local derivative, and passes the result backward.
\nThe full neural network is just the same mechanism repeated many times, with vectors and matrices instead of single numbers.
\nNow that we have seen what happens to a Gaussian distribution when it is passed through a tanh, and now that we know the local derivative of tanh, we can connect the two ideas.
The question is: what kind of distribution does a neuron see before the nonlinearity?
\nConsider one neuron:
\nz = w_1x_1 + w_2x_2 + \\dots + w_nx_n + bor, more compactly:
\nz = \\sum_{i=1}^{n} w_i x_i + bThis value z is the pre-activation. It is the input of the tanh:
h = \\tanh(z)To understand the scale of z, we need two basic facts about variance.
First, if random variables are independent, the variance of their sum is the sum of their variances:
\n\\mathrm{Var}(X_1 + X_2 + \\dots + X_n)\n=\n\\sum_{i=1}^{n} \\mathrm{Var}(X_i)Second, multiplying a random variable by a constant scales the variance by the square of that constant:
\n\\mathrm{Var}(aX) = a^2 \\mathrm{Var}(X)In our neuron, each term is w_i x_i. If we assume that the inputs and weights are independent and centered around zero, then the variance of the pre-activation is approximately:
\\mathrm{Var}(z)\n=\n\\sum_{i=1}^{n} \\mathrm{Var}(w_i x_i)and, under the usual independence assumptions:
\n\\mathrm{Var}(z)\n\\approx\nn \\, \\mathrm{Var}(w) \\, \\mathrm{Var}(x)From now on, assume that the input activations are normalized so that:
\n\\mathrm{Var}(x) \\approx 1Then the expression becomes:
\n\\mathrm{Var}(z)\n\\approx\nn \\, \\mathrm{Var}(w)This is the key point: under this assumption, the variance of the pre-activation grows with the number of inputs, also called fan_in.
If the weights are initialized too large, then \\mathrm{Var}(w) is large, so \\mathrm{Var}(z) becomes large. This means that many pre-activations fall far from zero.
But we already saw what tanh does in that case: it pushes large positive values close to 1 and large negative values close to -1.
So if z has high variance, then:
h = \\tanh(z)will be concentrated near -1 and 1.
This is a problem for backpropagation because the local derivative of tanh is:
\n\\frac{\\partial h}{\\partial z}\n=\n1 - h^2If h \\approx 1 or h \\approx -1, then:
1 - h^2 \\approx 0So the gradient that flows backward through the tanh gets multiplied by a number close to zero:
\n\\frac{\\partial \\mathcal{L}}{\\partial z}\n=\n\\frac{\\partial \\mathcal{L}}{\\partial h}\n(1 - h^2)This is local gradient attenuation. If many neurons are saturated, many gradients are killed locally. This is one way vanishing gradients appear in practice.
\nWe can see this directly in the network by plotting the hidden activations. Here we are still using the naive hidden-layer initialization: W1 is not fan-in normalized yet, while W2 is already scaled down to fix the initial output confidence.
plt.hist(h.view(-1).tolist(), 50);If a large amount of mass is close to -1 and 1, the hidden layer is saturated.
An even clearer diagnostic is:
\nplt.figure(figsize=(10, 20))\nplt.imshow(h.abs() > 0.99, cmap=\"grey\", interpolation=\"nearest\");![\"Tanh](\"https://www.tommasovaccari.com/static/tanh-saturation-map-1aa907ca.webp\")
Each white point is an activation with absolute value greater than 0.99. These are neurons that are almost fully saturated. For those values, the local tanh derivative is almost zero, so the gradient does not flow well through them.
The opposite problem is when the weights are initialized too small.
\nIf \\mathrm{Var}(w) is too small, then:
\\mathrm{Var}(z)\n\\approx\nn \\, \\mathrm{Var}(w) \\, \\mathrm{Var}(x)also becomes too small. In this case most pre-activations are very close to zero.
\nFor tanh, this does not kill the local gradient, because around zero:
\n\\tanh(z) \\approx zand:
\n\\tanh'(z) \\approx 1So the tanh itself is not saturated. The problem is more subtle: the activations are small because the pre-activations are small.
\nIf the next layer is:
\nz_{next} = hW_{next} + b_{next}then the gradient of the loss with respect to the next weight matrix has the usual form:
\n\\frac{\\partial \\mathcal{L}}{\\partial W_{next}}\n=\nh^T \\delta_{next}where \\delta_{next} is the gradient arriving from the next pre-activation. So if h is very small, the weight gradient is also small. The update becomes small not because tanh blocked the gradient locally, but because the activation that multiplies the gradient is tiny.
There is a second effect in the backward pass. The gradient sent to the previous layer is:
\n\\delta_{prev}\n=\n\\delta_{next} W_{next}^TIf the weights are very small, this backward signal is also scaled down. So with weights that are too small, the network can end up with small activations in the forward pass and small gradients in the backward pass.
\nSo we want a middle ground:
\nIt is important to be precise about what we are analyzing here. This is an initialization problem. We are looking at what happens at the first forward pass, before training has had the chance to move the parameters.
\nLater, during training, similar problems can appear again: activations can drift, distributions can shift, and neurons can still move into saturated regions. That is a different problem, and later we will look at mechanisms designed to keep activations well behaved during training.
\nFor now, we only want a good starting point.
\nFrom:
\n\\mathrm{Var}(z)\n\\approx\nn \\, \\mathrm{Var}(w) \\, \\mathrm{Var}(x)and assuming:
\n\\mathrm{Var}(x) \\approx 1we get:
\n\\mathrm{Var}(z)\n\\approx\nn \\, \\mathrm{Var}(w)So if we want:
\n\\mathrm{Var}(z) \\approx 1we need:
\nn \\, \\mathrm{Var}(w) \\approx 1which means:
\n\\mathrm{Var}(w) \\approx \\frac{1}{n}and therefore:
\n\\mathrm{Std}(w) \\approx \\frac{1}{\\sqrt{n}}This gives us a very simple first correction: scale the weights by the inverse square root of the fan-in.
\nW1 = torch.randn((n_embd * block_size, n_hidden), generator=g) * (1 / (n_embd * block_size)**0.5)This does not solve every training problem, but it fixes the first obvious one: the pre-activations are no longer exploding just because each neuron is summing many independent inputs.
\nAfter this change, we can rerun the same saturation diagnostic:
\nplt.figure(figsize=(10, 20))\nplt.imshow(h.abs() > 0.99, cmap=\"grey\", interpolation=\"nearest\");![\"Tanh](\"https://www.tommasovaccari.com/static/tanh-saturation-map-fan-in-75f334ec.webp\")
We can also look again at the activation histogram:
\nplt.hist(h.view(-1).tolist(), 50);![\"Tanh](\"https://www.tommasovaccari.com/static/tanh-activation-hist-fan-in-417cc91e.webp\")
Now we see far fewer saturated activations, because the initial pre-activation variance has been brought back to a reasonable scale. Clearly, there are more advanced techniques to adjust weights at initialization, such as Xavier/Glorot initialization, He/Kaiming initialization, and Fixup initialization.
\nBut later we are going to see that we can introduce some architectural changes that helps us take in control both the initialization and in training variance problem.
\nIn this section we are going to see some of the most common techniques used in modern neural networks to keep activations and gradients well behaved during training. These techniques are not just for initialization: they help maintain stable distributions throughout training, which is crucial for deep networks.
\nBatch Normalization was introduced by Sergey Ioffe and Christian Szegedy in the Google paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.
\nUntil now we mostly reasoned about initialization. We asked: if the input variance is approximately one, how should we initialize the weights so that the pre-activation variance does not explode or collapse at the first forward pass?
\nThat is important, but it is only the first forward pass.
\nDuring training, the weights of the previous layers keep changing. So the distribution seen by a layer also keeps changing. A layer may start with reasonable pre-activations, but after some updates those pre-activations can shift, widen, shrink, or move into the saturated region of the nonlinearity. This is the idea that the BatchNorm paper calls internal covariate shift: the input distribution of internal layers is not stable while the model is training.
\nThe idea of BatchNorm is simple:
\n\n\nbefore sending the pre-activations into the nonlinearity, normalize them over the mini-batch.
\n
In our case, the block becomes:
\nz = xW + b\\tilde{z} = \\mathrm{BatchNorm}(z)h = \\tanh(\\tilde{z})So the tanh does not directly see the raw linear output anymore. It sees a normalized version of it.
To understand the mechanism, start from the usual standardization formula. If:
\nX \\sim \\mathcal{N}(\\mu, \\sigma^2)then:
\nZ = \\frac{X - \\mu}{\\sigma}has:
\n\\mathbb{E}[Z] = 0and:
\n\\mathrm{Var}(Z) = 1If the original variable is Gaussian, then this produces a standard Gaussian:
\nZ \\sim \\mathcal{N}(0, 1)BatchNorm applies exactly this idea inside the network, but using mini-batch statistics instead of the true population mean and variance.
\nSuppose the linear layer output has shape:
\nz \\in \\mathbb{R}^{m \\times d}where m is the batch size and d is the number of neurons. Each column of z is one neuron evaluated over all examples in the batch. So BatchNorm normalizes each neuron independently across the batch.
For one neuron, we have the mini-batch values:
\n\\mathcal{B} = \\{z_1, z_2, \\dots, z_m\\}The batch mean is:
\n\\mu_{\\mathcal{B}} = \\frac{1}{m}\\sum_{i=1}^{m} z_iThe batch variance is:
\n\\sigma_{\\mathcal{B}}^2 = \\frac{1}{m}\\sum_{i=1}^{m}(z_i - \\mu_{\\mathcal{B}})^2Then we normalize:
\n\\hat{z}_i = \\frac{z_i - \\mu_{\\mathcal{B}}}{\\sqrt{\\sigma_{\\mathcal{B}}^2 + \\epsilon}}The small \\epsilon is only for numerical stability, so we never divide by zero.
At this point, for each neuron, the batch of normalized values has approximately zero mean and unit variance. This does not magically make every distribution Gaussian, but if the pre-activation is already roughly Gaussian, it brings it close to a standard Gaussian.
\nThere is one more important detail. If we forced every layer to always use zero-mean and unit-variance activations, we would reduce what the network can represent. Sometimes the network may actually want a shifted or scaled version of the normalized activation.
\nSo BatchNorm adds two learnable parameters per neuron:
\ny_i = \\gamma \\hat{z}_i + \\betawhere:
\n\\gamma is a learnable scale\\beta is a learnable shiftThis is crucial. BatchNorm normalizes the signal, but then gives the model the freedom to learn the right scale and offset again if that is useful.
\nThis is the implementation used in the notebook:
\nclass BatchNormalizationLayer:\n def __init__(self, fan_out, eps=1e-5, momentum=0.1):\n self.eps = eps\n self.momentum = momentum\n self.training = True\n\n self.gamma = torch.ones(fan_out, requires_grad=True)\n self.beta = torch.zeros(fan_out, requires_grad=True)\n\n self.running_mean = torch.zeros(fan_out)\n self.running_var = torch.ones(fan_out)\n\n def __call__(self, x):\n if self.training:\n batch_mean = x.mean(0)\n batch_var = x.var(0, unbiased=False)\n\n x_normalized = (x - batch_mean) / torch.sqrt(batch_var + self.eps)\n\n with torch.no_grad():\n self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean\n self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var\n else:\n x_normalized = (x - self.running_mean) / torch.sqrt(self.running_var + self.eps)\n\n out = self.gamma * x_normalized + self.beta\n self.out = out\n return out\n\n def parameters(self):\n return [self.gamma, self.beta]The important line is:
\nbatch_mean = x.mean(0)\nbatch_var = x.var(0, unbiased=False)The dimension 0 is the batch dimension. So we are not computing one global mean over the whole matrix. We are computing one mean and one variance for each neuron.
Then:
\nx_normalized = (x - batch_mean) / torch.sqrt(batch_var + self.eps)standardizes every column independently.
\nThe parameters:
\nself.gamma = torch.ones(fan_out, requires_grad=True)\nself.beta = torch.zeros(fan_out, requires_grad=True)are the learnable scale and shift. They are returned by parameters(), so the optimizer updates them like normal weights.
The running statistics are needed because training and inference are different. During training, using the current mini-batch statistics is fine. During inference, instead, we do not want the prediction for one example to depend on the other examples that happened to be in the same batch. So we use the running estimates accumulated during training:
\nself.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean\nself.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_varIn the model, we insert BatchNorm after the linear layer and before the tanh:
layers = [\n Linear(n_embd * block_size, n_hidden),\n BatchNormalizationLayer(n_hidden),\n Tanh(),\n Linear(n_hidden, n_hidden),\n BatchNormalizationLayer(n_hidden),\n Tanh(),\n Linear(n_hidden, vocab_size),\n]This placement matters. The paper also discusses applying BatchNorm before the nonlinearity, because the goal is to keep the nonlinearity input in a controlled range. For us this means: normalize the pre-activation, then apply tanh.
There is also a small practical consequence: if we normalize the output of xW + b, the bias before BatchNorm becomes less important, because subtracting the batch mean removes constant shifts. The learnable \\beta after normalization becomes the meaningful shift.
We will inspect the concrete effect in the diagnostics section below. The important point here is conceptual: initialization tries to make the first step healthy; BatchNorm keeps re-normalizing the signal while the network is changing.
\nLayer Normalization was introduced by Ba, Kiros and Hinton in Layer Normalization.
\nThe idea is close to BatchNorm, but the axis is different.
\nBatchNorm normalizes each neuron using the statistics of the mini-batch. So it looks across examples:
\nx \\in \\mathbb{R}^{batch \\times channels}LayerNorm, instead, normalizes the channel vector of a single example.
\nThis is the version that will be very useful when we move to transformers. In a transformer, each token is represented by a vector:
\nx_t \\in \\mathbb{R}^{d}where t is the token position and d is the number of channels, or embedding dimensions.
LayerNorm takes this vector and normalizes it across its channels:
\n\\mu_t = \\frac{1}{d}\\sum_{j=1}^{d} x_{t,j}\\sigma_t^2 = \\frac{1}{d}\\sum_{j=1}^{d}(x_{t,j} - \\mu_t)^2\\hat{x}_{t,j} = \\frac{x_{t,j} - \\mu_t}{\\sqrt{\\sigma_t^2 + \\epsilon}}and then, as usual, it gives the model a learnable scale and shift:
\ny_{t,j} = \\gamma_j \\hat{x}_{t,j} + \\beta_jThe important intuition is this: for each token, LayerNorm asks \"is this token vector well scaled across its channels?\" It does not need to look at the other examples in the batch.
\nThat is why it fits transformers so naturally. When we predict the next token, every token representation is a channel vector that is repeatedly processed by attention and MLP blocks. LayerNorm keeps those token vectors in a reasonable scale before or after the block, depending on the architecture.
\nSo for now I only want to remember the operational difference:
\nAnother modern mechanism that directly attacks the training problem is the residual connection, introduced by He, Zhang, Ren and Sun in Deep Residual Learning for Image Recognition.
\nThe problem is very pragmatic. We would like to make networks deeper, because deeper networks should be able to build more abstract representations. But after some point, simply stacking more layers does not automatically help. The optimization becomes harder, gradients have to pass through many transformations, and the deeper model can even train worse than a shallower one.
\nThe residual idea is to stop forcing a block to learn the whole transformation from scratch.
\nWithout a residual connection, a block learns:
\ny = F(x)With a residual connection, the block learns:
\ny = x + F(x)This changes the meaning of what the block has to learn. The block is not asked to produce the full output anymore. It only has to learn a correction to the input.
\nThat is why the name is residual: F(x) is the residual part, the difference between what we already have and what we want.
If the best thing for a layer is to do almost nothing, the residual block can learn:
\nF(x) \\approx 0and then:
\ny \\approx xSo the block can behave like an identity mapping. This is important because adding more layers should not make the optimization problem worse just because the model has to rediscover how to copy information forward.
\nThere is also a direct gradient intuition. If:
\ny = x + F(x)then:
\n\\frac{\\partial y}{\\partial x}\n=\nI + \\frac{\\partial F(x)}{\\partial x}So during backpropagation the gradient has a direct path through the identity term. It does not have to pass only through the nonlinear block F. This does not mean gradients can never vanish or explode, but it gives the network a much cleaner route for information and gradients to flow through depth.
In code, the idea is basically:
\nx = x + block(x)or, if the dimensions need to be adapted:
\nx = projection(x) + block(x)This is why residual connections fit naturally in this discussion. BatchNorm tries to keep the internal distributions under control. Good initialization tries to make the first forward/backward pass reasonable. Residual connections make depth easier by giving the network a stable path that can carry both activations and gradients across many layers.
\nDropout was introduced by Srivastava, Hinton, Krizhevsky, Sutskever and Salakhutdinov in Dropout: A Simple Way to Prevent Neural Networks from Overfitting.
\nThe motivation is different from BatchNorm and residual connections.
\nBatchNorm and residual connections mainly help with optimization and signal flow. Dropout is mostly a regularization technique: it tries to prevent the network from fitting the training data too specifically.
\nThe idea is simple. During training, for each neuron we sample a Bernoulli random variable that decides if that neuron is active or not.
\nm_i \\sim \\mathrm{Bernoulli}(p)where p is the probability that the neuron is kept active.
If the sampled value is 1, the neuron fires and its activation is used. If the sampled value is 0, the neuron is switched off for that forward pass.
So at every training step, the network is slightly different. Some neurons are present, some neurons do not fire. The model cannot rely too much on one specific activation always being there.
\nThe intuition is that dropout forces redundancy. If a feature is useful, the network should not encode it in one fragile path only. It should learn representations that still work even when some neurons are temporarily off.
\nDropout is not something I would add blindly to fix a broken training run. If activations are saturated or gradients are dead, dropout does not solve that. It can even make optimization noisier. But once the model trains and starts to overfit, dropout is a clean way to make the network less dependent on exact neuron co-adaptations.
\nSo in the mental map:
\nAt this point, only looking at the loss is not enough.
\nThe loss tells us if the model is improving, but it does not tell us why the training is healthy or unhealthy. If the loss is bad, the problem could be saturated activations, gradients that are too small, gradients that are too large, or updates that are completely out of scale with the parameters.
\nSo we need a small diagnostic toolbox.
\nThe first plot is still the loss curve. But we do not plot every single iteration, because that is too noisy. We average every 1000 iterations:
\n# Now we can plot the loss curve to see how it evolves during training\nplt.figure(figsize=(10, 5))\nlossi_avg = torch.tensor(lossi).view(-1, 1000).mean(1)\nplt.plot(lossi_avg, label=\"Training Loss (log10)\")\nplt.axvline(100, color=\"r\", linestyle=\"--\", label=\"learning-rate decay\")\nplt.xlabel(\"1000-iteration blocks\")\nplt.ylabel(\"log10(loss)\")\nplt.legend()\nplt.show()![\"Diagnostic](\"https://www.tommasovaccari.com/static/diagnostic-loss-curve-batchnorm-2f2aee05.webp\")
This gives us the global view. The model improves quickly at the beginning, then the curve becomes flatter. The red dashed line is the learning-rate decay: after 100 blocks, so after 100000 iterations, the learning rate goes from 0.1 to 0.01.
This plot is useful, but it is not enough. It tells us that training is moving, but it does not tell us what is happening inside the network.
\nSince we are using tanh, the first thing I want to inspect is saturation. A tanh neuron is saturated when its output is very close to -1 or 1. In that region, the local derivative is close to zero, so the gradient has a hard time flowing backward.
The following plot shows, for each Tanh layer, which activations are saturated:
# Visualize saturated activations for the two Tanh layers\ntanh_layers = [(i, layer) for i, layer in enumerate(layers[:-1]) if isinstance(layer, Tanh)]\n\nplt.figure(figsize=(10, 4 * len(tanh_layers)))\n\nfor k, (i, layer) in enumerate(tanh_layers, 1):\n t = layer.out.detach()\n\n plt.subplot(len(tanh_layers), 1, k)\n plt.imshow((t.abs() > 0.99).float(), cmap=\"gray\", interpolation=\"nearest\", aspect=\"auto\")\n plt.title(f'layer {i} ({layer.__class__.__name__}) | saturated activations')\n plt.xlabel('neuron index')\n plt.ylabel('batch example')\n\nplt.tight_layout()![\"Tanh](\"https://www.tommasovaccari.com/static/diagnostic-tanh-saturation-by-layer-fb86375b.webp\")
Here the x-axis is the neuron index and the y-axis is the batch example. White pixels are activations with:
\n|h| > 0.99This plot is better than a single percentage because it preserves the structure. If we see a full vertical white stripe, that neuron is saturated for almost every example in the batch. That would be a bad sign. If instead we see sparse white pixels, then some examples are saturating but the whole layer is not dead.
\nIn this run, the first Tanh layer is still more saturated than the second one. So BatchNorm helped us control the signal, but it did not magically remove all saturation.
The saturation map is useful, but I also want the distribution view. The histogram tells us where the mass of the activations is.
\n# Inspect activation distributions layer by layer\ntanh_layers = [(i, layer) for i, layer in enumerate(layers[:-1]) if isinstance(layer, Tanh)]\n\nn = len(tanh_layers)\nplt.figure(figsize=(4 * n, 3))\n\nfor k, (i, layer) in enumerate(tanh_layers, 1):\n t = layer.out.detach().view(-1)\n\n mean = t.mean().item()\n std = t.std().item()\n sat = (t.abs() > 0.97).float().mean().item() * 100\n\n print(f'layer {i:2d} ({layer.__class__.__name__:>10s}) | mean {mean:+.2f} | std {std:.2f} | saturated {sat:.2f}%')\n\n plt.subplot(1, n, k)\n plt.hist(t.tolist(), bins=40, density=True)\n plt.title(f'layer {i}')\n plt.xlabel(f'μ={mean:+.2f}, σ={std:.2f}\\nsat={sat:.1f}%')\n plt.ylim(bottom=0)\n\nplt.suptitle('Activation distributions by layer')\nplt.tight_layout()layer 2 ( Tanh) | mean -0.01 | std 0.74 | saturated 19.19%\nlayer 5 ( Tanh) | mean -0.01 | std 0.64 | saturated 7.50%![\"Activation](\"https://www.tommasovaccari.com/static/diagnostic-activation-distributions-by-layer-d7d14b80.webp\")
This confirms the same story. The means are close to zero, which is good. But the first layer has a larger standard deviation and more mass close to -1 and 1, so it has more saturated activations.
This is the kind of plot that makes the problem visible: we are not just saying \"maybe tanh saturates\". We can actually inspect where and how much it is happening.
\nNow we look at the backward pass.
\nAfter loss.backward(), each layer output has a gradient. For the Tanh layers, this tells us the gradient signal that is flowing backward through those activations:
# Inspect gradient distributions layer by layer\n# We want gradients that are not collapsed to zero and not exploding.\ntanh_layers = [(i, layer) for i, layer in enumerate(layers[:-1]) if isinstance(layer, Tanh)]\n\nn = len(tanh_layers)\nplt.figure(figsize=(4 * n, 3))\n\nfor k, (i, layer) in enumerate(tanh_layers, 1):\n t = layer.out.grad.detach().view(-1)\n\n mean = t.mean().item()\n std = t.std().item()\n\n print(f'layer {i:2d} ({layer.__class__.__name__:>10s}) | mean {mean:+.3e} | std {std:.3e}')\n\n plt.subplot(1, n, k)\n plt.hist(t.tolist(), bins=40, density=True)\n plt.title(f'layer {i}')\n plt.xlabel(f'μ={mean:+.1e}, σ={std:.1e}')\n plt.ylim(bottom=0)\n\nplt.suptitle('Gradient distributions by layer')\nplt.tight_layout()layer 2 ( Tanh) | mean -3.638e-12 | std 3.133e-03\nlayer 5 ( Tanh) | mean -2.237e-05 | std 5.294e-03![\"Gradient](\"https://www.tommasovaccari.com/static/diagnostic-gradient-distributions-by-layer-0124301e.webp\")
The mean is close to zero, and that is not a problem. Gradients have signs, so positive and negative values can cancel.
\nThe more interesting quantity here is the standard deviation. If the gradient distribution is collapsed almost exactly at zero, the layer is not receiving a useful learning signal. If it is extremely wide, then the backward signal may be unstable.
\nSo here I do not want \"large variance\" in an absolute sense. I want visible, non-collapsed variance: gradients should carry information, but not explode.
\nFinally, we inspect the gradients of the parameters themselves.
\nFor each weight matrix, we plot the distribution of its gradients and compute:
\n\\frac{\\mathrm{std}(\\nabla W)}{\\mathrm{std}(W)}This is a scale-aware diagnostic. A gradient with standard deviation 0.01 may be big or small depending on the scale of the parameter it is updating. So we compare gradient scale against weight scale.
# Inspect weight gradient distributions parameter by parameter\nweight_params = [(i, p) for i, p in enumerate(parameters) if p.ndim == 2]\n\nn = len(weight_params)\nplt.figure(figsize=(8, 3 * n))\n\nfor k, (i, p) in enumerate(weight_params, 1):\n t = p.grad.detach().view(-1)\n\n mean = t.mean().item()\n std = t.std().item()\n ratio = (p.grad.std() / p.std()).item()\n\n print(f'weight {tuple(p.shape)!s:>12s} | mean {mean:+.3e} | std {std:.3e} | grad:data ratio {ratio:.3e}')\n\n plt.subplot(n, 1, k) # one below the other\n plt.hist(t.tolist(), bins=40, density=True)\n plt.title(f'param {i} - shape {tuple(p.shape)}')\n plt.xlabel(f'μ={mean:+.1e}, σ={std:.1e}, ratio={ratio:.1e}')\n plt.ylim(bottom=0)\n\nplt.suptitle('Weight gradient distributions')\nplt.tight_layout()weight (27, 10) | mean +3.311e-10 | std 1.384e-02 | grad:data ratio 1.252e-02\nweight (30, 200) | mean +1.307e-05 | std 7.136e-03 | grad:data ratio 1.838e-02\nweight (200, 200) | mean +1.558e-05 | std 4.308e-03 | grad:data ratio 2.301e-02\nweight (200, 27) | mean +2.208e-11 | std 1.813e-02 | grad:data ratio 8.428e-02![\"Weight](\"https://www.tommasovaccari.com/static/diagnostic-weight-gradient-distributions-06847694.webp\")
One important detail: this ratio tells us how large the gradient is compared to the weight. But during SGD the weight is not updated with the raw gradient. It is updated with the learning-rate-scaled gradient:
\n\\Delta W = -\\eta \\nabla WSo if we want to know how much the weight really moves, we have to include the learning rate:
\n\\frac{\\mathrm{std}(\\Delta W)}{\\mathrm{std}(W)}\n=\n\\eta \\frac{\\mathrm{std}(\\nabla W)}{\\mathrm{std}(W)}For example, if:
\nweight std = 1\ngradient std = 0.01\nlearning rate = 0.01then:
\ngrad:data ratio = 0.01 / 1 = 1e-2\nupdate:data ratio = (0.01 * 0.01) / 1 = 1e-4So the plot tells us how strong the raw gradient signal is. To understand how much the parameters really move after one SGD step, we also have to multiply by the learning rate.
\nThis is the intuition I want from this plot:
\nSo these plots are not just decoration. They are the basic instruments I want before trusting a training run: loss curve, activation saturation, activation distribution, gradient distribution, and parameter update scale.
\nThe main point of this post is that a neural network is not just a stack of matrix multiplications followed by loss.backward().
It is also a system that moves distributions forward and gradients backward. If the distributions are badly scaled, the nonlinearities saturate. If the gradients collapse, the model does not learn. If the updates are too large, training becomes unstable. The loss curve shows the final symptom, but the statistics of activations, gradients, and updates show the mechanism.
\nSo the intuition I want to keep is simple: before changing architectures randomly, inspect the signal. Look at means, variances, saturation, gradient distributions, and update ratios. These plots are not advanced tooling; they are the minimum instrumentation needed to understand if the network is actually trainable.
\nMain practical reference:
\nmakemore lectures on activations, gradients, BatchNorm, and manual backpropagation.Background math and signal/statistics intuition:
\nPapers cited:
\nThis post is a technical walkthrough of MovHex, my project for Algorithms and Principles of Computer Science at Politecnico di Milano.
\nThe assignment: build a routing engine over a rectangular map of hexagons. The program must support four commands:
\ninit — create or reinitialize the mapchange_cost — update hexagon costs within a hexagonal radiustoggle_air_route — add or remove a directed air routetravel_cost — answer shortest-path queriesConcretely, the program must accept the following command formats:
\ninit <n_columns> <n_rows>\nchange_cost <x> <y> <v> <radius>\ntoggle_air_route <x1> <y1> <x2> <y2>\ntravel_cost <xp> <yp> <xd> <yd>where:
\nn_columns, n_rows define the size of the rectangular hex mapx, y identify a hexagon using the specification's coordinate systemv is the signed update applied by change_costradius is the positive hex-distance radius of the updatex1, y1 are the source coordinates of a directed air routex2, y2 are the destination coordinates of a directed air routexp, yp are the source coordinates of a shortest-path queryxd, yd are the destination coordinates of a shortest-path queryAt first glance this looks like \"implement Dijkstra on a graph.\" In practice, the real difficulty is in the interaction between hex geometry, mutable state, and performance constraints.
\nThe verifier for the course included explicit thresholds for 30 cum laude:
\nMy final version meets both thresholds.
\nCorrectness alone was not enough. The solution had to be modeled carefully, implemented safely in C, and then profiled and specialized until it was fast enough.
\nWhenever I show code in this post, I show lightly reformatted excerpts from the final main.c in the repository. If a block contains ..., that means omitted lines, not rewritten logic.
The official statement describes a rectangular tiling of equal-sized hexagons. Hexagons are identified by column first, row second. Indices start from zero, increase left to right and bottom to top. The hexagon (0,1) sits on the upper-right side of (0,0).
That last detail is important: the coordinate system encodes an offset hex layout, and that affects both neighbor generation and distance computation.
\n![\"Hex](\"https://www.tommasovaccari.com/static/coordinate-system-07d434ed.webp\")
The coordinate system defined by the specification: columns grow left to right, rows grow bottom to top, and (0,1) sits on the upper-right side of (0,0).
Each hexagon stores a natural number in [0,100]:
0, the hexagon can still be visited, but it cannot be abandoned — no ground or air departure is possible from itinitinit n_columns n_rows creates or reinitializes a map of n_rows × n_columns hexagons. All costs are set to 1, no air routes exist. The program responds OK.
change_costchange_cost x y v radius updates every hexagon whose hex distance from (x,y) is strictly less than the given radius. The specification calls this distance DistEsagoni and defines it as the minimum number of hexagons traversed to reach the destination from the source, including the destination in the count, ignoring costs, blocked departures, and air routes.
That \"including the destination\" clause pins the semantics: adjacent hexagons have distance 1, and the hexagon at the center has distance 0 from itself.
\nIn the rest of this post I refer to the same quantity as hexDistance, because that is the name I used while reasoning about the implementation. The formula and the semantics are the same as DistEsagoni in the specification.
The update rule from the specification is:
\n\\operatorname{cost}(x_e,y_e)=\\operatorname{cost}(x_e,y_e)+\\left\\lfloor v \\cdot \\max\\left(0,\\frac{r-\\operatorname{DistEsagoni}\\big((x_e,y_e),(x,y)\\big)}{r}\\right)\\right\\rfloorIf I write d = \\operatorname{hexDistance}\\big((x_e,y_e),(x,y)\\big), the additive delta is:
\\Delta(d)=\\left\\lfloor v \\cdot \\max\\left(0,\\frac{r-d}{r}\\right)\\right\\rfloorThe result is clamped to [0,100]. The same update also applies to all outgoing air routes of every affected hexagon. The program responds KO if (x,y) is not a valid hexagon or if radius = 0, otherwise OK.
toggle_air_routetoggle_air_route x1 y1 x2 y2 adds or removes a directed air route. If the route does not exist, it is created. If it already exists, it is removed.
When a new route is created, its cost is the floor of the average of all already existing outgoing air-route costs from (x1,y1) together with the ground exit cost of (x1,y1). Each hexagon may have at most 5 outgoing air routes. The program responds OK if coordinates are valid and the limit is not exceeded, KO otherwise.
When a route is removed, no cost recalculation happens — the route is simply deleted.
\ntravel_costtravel_cost xp yp xd yd asks for the minimum cost of reaching (xd,yd) from (xp,yp). The rules:
-1 if coordinates are invalid or if no path existsThat third rule matters. It means a weight-0 hexagon can still produce a valid travel_cost answer when it is both source and destination.
The specification gives one more important clue: in realistic inputs, change_cost and toggle_air_route are rare, while travel_cost is used heavily, and most source/destination pairs concentrate in the same zones.
That hint shaped the whole approach:
\nThe central design question was not \"which graph algorithm should I use?\" It was: what state does this problem actually require me to store permanently?
\nThe final implementation is built on four decisions:
\nThe permanent map consists of three parallel arrays:
\nweight[] (uint8_t) — the ground exit cost of each hexagon, in [0,100]counter_air_route[] (uint8_t) — the number of outgoing air routes from each node, in [0,5]air_route[] (uint32_t* per node) — a dynamically allocated array of destination indices for each node, NULL when the node has no outgoing air routesThat is all. No adjacency lists for ground links, no per-route cost fields, no embedded graph metadata.
\nDijkstra's algorithm needs its own data structures, allocated once after init and reused across every travel_cost query:
distance[] (uint32_t) — tentative distances from the source, initialized to UINT32_MAX at the start of each run via memset(distance, 0xFF, ...)travel_cost resultsThis separation matters. Ground neighbors are computed on the fly from coordinates and row parity. Air routes are stored because they mutate. Dijkstra metadata is transient and does not belong inside the map.
\nThe mapping between 2D coordinates and the flat index is:
\n\\text{index} = i \\cdot \\text{COLS} + jstatic inline void convert2to1(uint32_t i, uint32_t j, uint32_t COLS, uint32_t *INDEX){\n *INDEX = i * COLS + j;\n}\n\nstatic inline void convert1to2(uint32_t *i, uint32_t *j, uint32_t COLS, uint32_t INDEX){\n *i = INDEX / COLS;\n *j = INDEX % COLS;\n}The original implementation used a struct hex_t containing weight, distance, predecessor, ground-link lists, and air-route lists — all packed into one type. Profiling showed that the struct was too fat for cache lines, and most fields were only needed during Dijkstra. The final architecture is the result of dismantling that struct into flat parallel arrays and removing everything that could be derived instead of stored.
The first hard problem was not shortest paths. It was getting the geometry right.
\nA square grid lets you get away with matrix intuition for a long time. A hex grid does not. Neighbor relations depend on row parity, and the specification numbers rows from bottom to top, while most standard hex-grid references assume top-based coordinates. Trying to reason in a matrix-like way immediately produces questions that do not exist on square grids: if I move to the upper-right neighbor, do I stay in the same column or shift? Does that answer change when the row is even?
\nThe external reference that made this click for me was the Red Blob Games guide on hexagonal grids. It gave me the right conceptual model: separate the problem into three coordinate layers instead of trying to make matrix intuition work directly.
\n(q, r, s) subject to the constraint q + r + s = 0, which makes distance computation clean and parity-independentThe progression was concrete: matrix intuition broke because it does not account for the row-dependent offset, and only after recognizing the offset layout explicitly did the cube-coordinate machinery become applicable.
\nIn the offset layout, moving to the upper-right neighbor from a hexagon depends on whether the current row is even or odd. In an even row, the upper-right neighbor is in the same column. In an odd row, it shifts one column to the right.
\nThe final code encodes this as two lookup tables:
\nconst int8_t neighbors_set[2][6][2] = {\n // Even rows (parity 0): (col_offset, row_offset)\n { {-1,+1}, {0,+1}, {+1,0}, {0,-1}, {-1,-1}, {-1,0} },\n // Odd rows (parity 1):\n { {0,+1}, {+1,+1}, {+1,0}, {+1,-1}, {0,-1}, {-1,0} }\n};During Dijkstra, the row parity of the current node selects the correct table, and each of the six neighbors is generated by adding the corresponding offset to the current (col, row) and checking bounds with isinrange. No adjacency list is needed.
Hex distance is the core quantity behind change_cost and the geometric model. Computing it directly in offset coordinates is error-prone because the parity-dependent shifts make the arithmetic inconsistent across rows. Cube coordinates solve this.
In cube coordinates, every hexagon is represented as (q, r, s) with the constraint q + r + s = 0. The key property: each single hex step changes exactly two of the three coordinates by \\pm 1 — one increases, one decreases, the third stays the same — so the constraint is always maintained. This means each step adds exactly 2 to the Manhattan distance |{\\Delta q}| + |{\\Delta r}| + |{\\Delta s}|. Therefore:
d(a, b) = \\frac{|q_1 - q_2| + |r_1 - r_2| + |s_1 - s_2|}{2}This formula is parity-independent and well-defined regardless of the offset convention, which is the whole point of converting to cube coordinates before computing distance.
\nROWS % 2 BranchThe conversion from offset coordinates to cube depends on whether the layout is odd-r (odd rows shifted right) or even-r (even rows shifted right). The formulas from the Red Blob Games model are:
\nq = \\text{col} - \\lfloor(\\text{row} - \\text{parity}) / 2\\rfloor, r = \\text{row}, s = -q - rq = \\text{col} - \\lfloor(\\text{row} + \\text{parity}) / 2\\rfloor, r = \\text{row}, s = -q - rwhere \\text{parity} = \\text{row} \\bmod 2. The only difference between the two is a single sign flip in the q computation: +parity for even-r, -parity for odd-r. Getting that sign wrong produces plausible-looking results that silently fail on specific row parities.
Now the complication. The specification numbers rows from the bottom. The cube-coordinate formulas assume top-based numbering. Before converting, the code flips the row:
\ny' = (\\text{ROWS} - 1) - yThis flip inverts parity when \\text{ROWS} - 1 is odd (i.e., when ROWS is even):
\\operatorname{parity}(y') = \\operatorname{parity}(\\text{ROWS} - 1) \\oplus \\operatorname{parity}(y)The consequence:
\nROWS is odd, \\text{ROWS} - 1 is even, parity is preserved → the grid remains odd-r after flipping → use the odd-r formula (-\\text{parity})ROWS is even, \\text{ROWS} - 1 is odd, parity inverts → the grid becomes even-r after flipping → use the even-r formula (+\\text{parity})That is the real reason the final code branches on ROWS % 2:
if(ROWS % 2 == 0){\n i1 = (ROWS - 1) - i1;\n i2 = (ROWS - 1) - i2;\n\n int8_t parity1 = i1 % 2;\n int8_t parity2 = i2 % 2;\n int32_t q1 = j1 - ((i1 + parity1)/2); // even-r\n int32_t q2 = j2 - ((i2 + parity2)/2);\n ...\n}else{\n i1 = (ROWS - 1) - i1;\n i2 = (ROWS - 1) - i2;\n\n int8_t parity1 = i1 % 2;\n int8_t parity2 = i2 % 2;\n int32_t q1 = j1 - ((i1 - parity1)/2); // odd-r\n int32_t q2 = j2 - ((i2 - parity2)/2);\n ...\n}Both branches then compute s = -q - r and return (|q_1-q_2|+|r_1-r_2|+|s_1-s_2|)/2.
If the even-r/odd-r choice is inverted, or if the row flip is omitted, the distance function produces results that look correct on many test cases but fail systematically on whole classes of inputs where the row parity changes the expected geometry.
\nThe final representation stores no per-route cost. Each air route is just a destination index. This is possible because of an invariant that holds throughout execution.
\nInvariant: all outgoing air routes from the same source u always have the same cost, and that cost equals \\operatorname{weight}[u].
Proof by induction on the sequence of operations:
\nu: there are no existing air routes. The cost is \\lfloor \\operatorname{weight}[u] / 1 \\rfloor = \\operatorname{weight}[u].toggle_air_route): suppose all k existing outgoing air routes from u have cost \\operatorname{weight}[u]. The new route's cost is \\lfloor (k \\cdot \\operatorname{weight}[u] + \\operatorname{weight}[u]) / (k + 1) \\rfloor = \\operatorname{weight}[u].change_cost: the specification says the same update formula applies both to \\operatorname{weight}[u] and to all outgoing air routes of u. The delta \\Delta(d) = \\lfloor v \\cdot \\max(0, (r - d)/r) \\rfloor depends only on the distance from the pivot and the parameters v and r — not on the current cost. So weight and every air-route cost receive the same integer delta. After clamping to [0,100], they remain equal.The consequence: the effective traversal cost of an air move from u is always \\operatorname{weight}[u], which is the same as the ground exit cost. The code stores only destination indices:
air_route[index1][*counter_air_route] = index2;\n*counter_air_route += 1;And the shortest-path code uses currweight for both ground and air moves:
// Ground links\nnew_distance = currdistance + currweight;\n...\n// Air links\nuint32_t neighbor = air_route[i];\nnew_distance = currdistance + currweight;This invariant also simplifies change_cost: since air-route costs are implicit in weight[], the function only needs to update the weight array. The air-route cost update mandated by the specification is automatically satisfied.
change_cost in the Final Implementationchange_cost is a geometric enumeration problem, not a traversal problem.
An earlier implementation used a queue-based expansion with a visited array — start from the pivot, expand to neighbors, track distance, stop at the radius. It was correct, but profiling showed insert_neighbors was a major bottleneck: expanding neighbors iteratively through linked lists was expensive, especially for large radii.
The final implementation starts from geometry instead. Once the pivot has been converted into cube coordinates, the code iterates directly over all cube offsets inside the hexagonal region:
\nfor (int32_t q2 = -radius; q2 <= radius; q2++) {\n int32_t r2_start = max(-radius, -q2 - radius);\n int32_t r2_end = min(radius, -q2 + radius);\n for (int32_t r2 = r2_start; r2 <= r2_end; r2++) {\n int32_t qc = q1 + q2;\n int32_t rc = r1 + r2;\n ...\n uint32_t distance = hex_distance(j1, i1, j2, i2, ROWS);\n if(distance >= (uint32_t)radius) continue;\n change_cost_hex(weight, index, distance, v, radius);\n }\n}The nested loop directly enumerates all cube coordinates (q_c, r_c, s_c) such that |q_2| + |r_2| + |s_2| \\le 2 \\cdot \\text{radius} (the hexagonal region around the pivot). The r2_start/r2_end bounds enforce the cube constraint q + r + s = 0 implicitly — they restrict r_2 so that s_2 = -q_2 - r_2 stays within [-\\text{radius}, +\\text{radius}].
![\"Hex-distance](\"https://www.tommasovaccari.com/static/change-cost-radius-example-7c69f51c.webp\")
Hexagons at the same hex distance from the center receive the same update delta, so the effect propagates by hex-distance layers.
\nEach candidate is converted back to grid coordinates, bounds-checked, and then updated with the weighted delta:
\nint32_t compute_weighted_value(float v, float radius, float distance) {\n float normalized = (radius - distance) / radius;\n float clamped = fmaxf(0.0f, normalized);\n float weighted = v * clamped;\n return (int32_t)floorf(weighted);\n}The update clamps the result to [0,100]:
void change_cost_hex(uint8_t *weight, uint32_t index, uint32_t distance,\n int8_t v, int32_t radius)\n{\n int32_t new_weight = compute_weighted_value(v, radius, distance);\n int32_t sum = weight[index] + new_weight;\n if(sum < 0) sum = 0;\n if(sum > 100) sum = 100;\n weight[index] = (uint8_t) sum;\n}travel_cost and the Dijkstra ImplementationsAt the modeling level, the problem is a sparse directed graph with nonnegative integer edge weights:
\nair_route[]weight[source] (from the invariant in Section 5)[0, 100]All edge weights are nonnegative, so Dijkstra's algorithm is the correct base choice. The final code keeps two implementations behind the compile-time flag DIJKSTRA_IMPL: 0 for binary min-heap, 1 for weight-sorted buckets. The default build uses the second.
Both Dijkstra variants share the same relaxation logic. For the currently visited node, the code computes its row parity via convert1to2, then iterates over all 6 ground neighbors (using neighbors_set[parity]) and all air-route destinations. For each reachable neighbor:
// Ground links\nfor(int8_t i = 0; i < 6; i++){\n int32_t sx2 = (int32_t)x1 + neighbors_set[parity][i][0];\n int32_t sy2 = (int32_t)y1 + neighbors_set[parity][i][1];\n if(!isinrange(sy2, sx2, ROWS, COLS)) continue;\n ...\n new_distance = currdistance + currweight;\n if(new_distance < distance[neighbor_index]){\n // relax\n }\n}\n// Air links\nfor(int8_t i = 0; i < counter_air_route[hexindex_visiting]; i++){\n uint32_t neighbor = air_route[i];\n new_distance = currdistance + currweight;\n if(new_distance < distance[neighbor]){\n // relax\n }\n}Both ground and air relaxation use currdistance + currweight as the candidate distance. The invariant from Section 5 is what makes this correct for air routes: the air-route traversal cost is always equal to the source node's weight.
Both variants handle weight == 0 the same way: a hexagon with zero weight can be visited (its distance is set during relaxation from a neighbor) but it cannot be departed from — its neighbors are never relaxed through it. This directly implements the specification rule that a zero-cost hexagon \"cannot be abandoned but can be visited.\"
Both variants use early exit: once the destination is extracted from the priority structure, the algorithm terminates. Dijkstra's optimality guarantee ensures that at that point its distance is final, so there is no need to continue.
\nBoth variants also handle travel_cost bookkeeping the same way:
if((i1 == i2) && (j1 == j2)){\n *COST = 0;\n return 1;\n}\n...\nif(find_cost(hash_table, index1, index2, COST) == 1){\n if(*COST == UINT32_MAX) return -1;\n return 1;\n}\n...\n// Run Dijkstra\n*COST = distance[index2];\ninsert_cost(hash_table, index1, index2, COST);\nif(distance[index2] == UINT32_MAX) return -1;\nreturn 1;The source-equals-destination check happens before anything else. Then the cache is queried. Only if both miss does Dijkstra run. After the run, the result is cached and returned — including UINT32_MAX for unreachable destinations, so repeated misses also avoid recomputation.
DIJKSTRA_IMPL 0)The heap variant uses a standard binary min-heap with an explicit position-tracking array.
\nEach heap entry is a heap_block_t:
typedef struct heap{\n uint32_t distance;\n uint32_t hexnumber;\n}heap_block_t;The implementation detail that makes this work is min_heap_index[] (int32_t per node), which tracks every node's state with respect to the heap using three distinct values:
-1 (initialized via memset(0xFF)): the node has never been inserted into the heapINT32_MAX: the node has been popped and must not be reinserted (settled)This three-state tracking is what allows relax_neighbor to distinguish between insertion and decrease-key without a separate settled array:
static inline void relax_neighbor(heap_block_t *min_heap, uint32_t *distance,\n int32_t *min_heap_index, uint32_t *heap_size,\n uint32_t neighbor_index, uint32_t new_distance)\n{\n if(min_heap_index[neighbor_index] == -1){\n // Never seen → insert into heap\n distance[neighbor_index] = new_distance;\n min_heap[*heap_size].distance = new_distance;\n min_heap[*heap_size].hexnumber = neighbor_index;\n min_heap_index[neighbor_index] = *heap_size;\n *heap_size += 1;\n bubble_up(min_heap, min_heap_index[neighbor_index], min_heap_index);\n }else if(min_heap_index[neighbor_index] != INT32_MAX){\n // Currently in heap → decrease-key\n distance[neighbor_index] = new_distance;\n min_heap[min_heap_index[neighbor_index]].distance = new_distance;\n bubble_up(min_heap, min_heap_index[neighbor_index], min_heap_index);\n }\n // If INT32_MAX → already settled, skip\n}pop_min extracts the root, moves the last heap element to position 0, calls min_heapify to restore the heap property, and marks the popped node as INT32_MAX in min_heap_index. Both bubble_up and min_heapify update min_heap_index on every swap, so the mapping stays consistent.
The initialization allocates the heap and index arrays once in main() after init, and each Dijkstra run resets them with memset:
memset(distance, 0xFF, ROWS*COLS * sizeof(uint32_t));\nmemset(min_heap_index, 0xFF, (ROWS*COLS)*sizeof(int32_t));The complexity is O((V + E) \\log V). Since each hexagon has at most 11 outgoing edges (6 ground + 5 air), E = O(V), so the total is O(V \\log V).
DIJKSTRA_IMPL 1)The bucket variant exploits the fact that edge weights are bounded integers in [0, 100].
The idea comes from Dial's algorithm: instead of a heap, use a circular array of 101 buckets indexed by distance % 101. Since the maximum single-edge weight is 100, all active (unprocessed) entries fit within a window of 101 consecutive distance values at any point during execution. Scanning buckets in order from the current index guarantees that entries are processed in nondecreasing distance order.
Each bucket is a buckets_t:
typedef struct buck{\n uint32_t *nodes; // node indices\n uint32_t *keys; // distances at insertion time\n int32_t size; // total entries pushed\n int32_t cursor; // next entry to pop\n int32_t maxdim; // current allocation capacity\n}buckets_t;Each bucket stores two parallel arrays — nodes[] and keys[] — instead of an array of structs. This was a deliberate choice driven by the same principle that led to dismantling hex_t: keeping related data types contiguous for better cache behavior during sequential scans.
The relaxation function is much simpler than the heap variant — it just pushes a new entry into the appropriate bucket:
\nstatic inline void relax_neighbor_weight_sort(buckets_t *bucket, uint32_t key,\n uint32_t hexindex, uint32_t *tobeprocessed, uint32_t *distance)\n{\n distance[hexindex] = key;\n push(bucket, hexindex, key, tobeprocessed);\n}The caller selects the bucket as buckets[new_distance % 101]. This means the same node can appear in multiple buckets with different distances (a lazy approach). The algorithm handles this with stale-entry detection at pop time:
pop(&buckets[idx], &tobeprocessed, &hexindex_visiting, &key);\nif(hexindex_visiting == UINT32_MAX){\n idx = (idx+1) % 101;\n continue;\n}\n...\nif (key != currdistance) continue; // stale entry — distance was improved since insertion\nif (settled[hexindex_visiting]) continue; // already finalized\nif (hexindex_visiting == DST) break; // early exit — even if weight is 0\nif (currweight == 0) continue; // cannot depart from this nodeThe stale-entry check (key != currdistance) is necessary because a node's best distance may have been improved after this entry was pushed — the old entry remains in the bucket but is no longer valid. The settled[] array (uint8_t per node, zeroed at the start of each run) tracks finalized nodes. The early-exit check comes before the weight check: the destination's distance is already final when it is popped, and its exit cost is ignored by the specification, so there is no reason to continue regardless of its weight.
The bucket scanning loop advances idx = (idx+1) % 101 whenever a bucket is exhausted, cycling through the circular array. The push operation doubles the bucket's allocation when full (realloc to 2 * maxdim). The pop operation resets cursor and size to 0 when the cursor catches up to size, reclaiming the bucket for reuse.
The amortized complexity is O(V + E + C \\cdot V) where C = 101. Since E = O(V) and C is a constant, this simplifies to O(V) for this problem — eliminating the \\log V factor from heap operations.
Both satisfy the 30L thresholds. The bucket variant nearly halves the execution time because it replaces O(\\log V) heap operations with O(1) bucket insertions and amortized scanning.
The specification hints that travel_cost will be called frequently and that queries concentrate in the same zones. That made caching worth implementing.
The final strategy is strict:
\n(source, destination)init, change_cost, or toggle_air_routeUINT32_MAX, so repeated misses also avoid recomputationThere is no partial invalidation. The graph is dynamic and correctness is easy to lose with selective invalidation schemes. Full flush on mutation is simple and trustworthy.
\nThe cache is a 256-bucket hash table with separate chaining. The key is the ordered pair (src, dst): first I combine the two indices with a mixing function using large primes, then I hash the result using Knuth's multiplicative method and take the top 8 bits as the bucket index:
uint32_t combine(uint32_t src, uint32_t dst){\n return src * 73856093u ^ dst * 19349663u;\n}\n\nstatic inline uint32_t hash(uint32_t k) {\n uint64_t prod = (uint64_t)k * ALPHA_FIXEDPOINT;\n uint32_t frac = (uint32_t)prod;\n return frac >> 24;\n}ALPHA_FIXEDPOINT is 2654435769u, the fixed-point representation of Knuth's suggested constant (\\sqrt{5}-1)/2. The shift by 24 bits produces an 8-bit index in [0, 255].
Invalidation after a successful mutation flushes the entire table:
\nif(change_cost(weight, args, ROWS, COLS) == 1){\n printf(\"OK\\n\");\n free_hash_table(hash_table);\n}The same pattern is used after init and toggle_air_route.
Once the code was correct, profiling determined the optimization path. This was not guesswork — I used xctrace (Time Profiler) and Cachegrind at multiple stages.
One Cachegrind snapshot from an earlier version showed:
\n| Function | Instruction Share |
|---|---|
min_heapify | 31.78% |
Dijkstra | 21.12% |
pop_min | 14.32% |
insert_neighbors | 9.85% |
hex_distance | 6.86% |
bubble_up | 5.65% |
change_cost | 3.86% |
These bottlenecks were not all of the same kind:
\nmin_heapify + pop_min + bubble_up totaling ~52%) dominated inside repeated shortest-path queries. The \\log V cost per operation was being paid on every node extraction and every decrease-key.insert_neighbors at ~10%) was expensive in the traversal-based change_cost — it was walking linked lists to find neighbors instead of computing them directly.hex_t struct with embedded linked lists for ground links and air routes) were hurting cache locality. The _int_malloc and _int_free entries in the full Cachegrind output confirmed that dynamic allocation overhead was nontrivial.That breakdown pushed a specific sequence of redesigns:
\nchange_cost: traversal to geometric enumeration. The queue-based expansion with visited arrays was replaced with the direct cube-offset iteration shown in Section 6. The traversal version spent most of its time in insert_neighbors; the geometric version eliminated that function entirely._int_malloc and _int_free were contributing nontrivial overhead from the linked-list allocations.hex_t dismantled into parallel arrays. Weight, distance, heap index, air routes, and counters became separate flat arrays. This let the hot loop in Dijkstra touch only the data it actually needed per iteration instead of pulling an entire multi-field struct into cache.main() and reused across queries. The permanent map shrank to just weight[], counter_air_route[], and air_route[].[0, 100] made the binary heap overqualified. The 101-bucket Dial's variant eliminated the \\log V factor from priority operations, nearly halving execution time.Every optimization that worked was a removal of unnecessary state, not an addition of clever machinery.
\nThe final implementation works because the representation, the invariants, and the hot-path optimizations are aligned with the structure of the assignment:
\nweight[] because the specification guarantees a stable invariantchange_cost operates by geometric enumeration in cube space because the command is defined by hex distancetravel_cost is backed by Dijkstra with a three-state index tracker (heap variant) or stale-entry tolerance (bucket variant), with early exit on the destination(src, dst) memoization and full invalidation on mutation[0, 100] justify replacing a generic O(\\log V) heap with a constant-time bucket structureThe performance gain did not come from switching to a fancier shortest-path algorithm alone. It came mostly from making the data layout match the actual geometry, mutability, and bounded-weight structure of the problem: smaller permanent state, fewer cache misses, less pointer chasing, and less unnecessary memory traffic in the hot path.
","date_published":"2026-03-29T00:00:00.000Z","tags":["C","Graph Algorithms","Data Structures","Performance"]},{"id":"https://www.tommasovaccari.com/blog/from-bigrams-to-neural-networks","url":"https://www.tommasovaccari.com/blog/from-bigrams-to-neural-networks","title":"From Bigrams to Neural Networks: The First Step in Language Modeling","summary":"A hands-on walkthrough from count-based bigrams to a simple neural model.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"In this post, we tackle one of the most fundamental problems in natural language processing: language modeling. Our concrete objective is to build a system that generates plausible Italian names by predicting the next character given the previous one.
\nLanguage modeling assigns probabilities to token sequences by predicting the next token given its context. This simple idea underpins tools as familiar as autocomplete and as powerful as large-scale generative models.
\nWe begin with the bigram model: P(x_t \\mid x_{t-1}).\nThis model looks back exactly one step, capturing only immediate dependencies. Despite its simplicity, bigrams marked an early step beyond rule-based formal grammars.
From there, we reformulate the bigram as a neural network. This isn’t just reimplementation for its own sake: it shows how fixed statistical tables can be generalized into trainable systems, the very principle that scales into today’s large language models.
\nWith the goal defined, we’re ready to start from the basics.
\nWe start from data: the ISTAT list of Italian given names. The raw text contains diacritics, apostrophes, hyphens, spaces, and mixed case. To make the first model tractable, we project the text to a reduced alphabet and remove extremely short items. Formally, we apply a mapping
\n\\phi : \\Sigma_{\\text{raw}} \\;\\to\\; \\Sigma_{\\text{ascii}} = \\{a, \\ldots, z\\}\nby lowercasing, Unicode-normalizing, and stripping non-letters; the model is trained on \\phi(\\text{names}).
Normalization choices (trade-offs):
\nlen < 3) to reduce degenerate contexts and stabilize bigram counts.Minimal, deterministic preprocessing:
\nimport re\n\nwith open(\"italian_names.txt\", \"r\", encoding=\"utf-8\") as f:\n names = [line.strip().lower() for line in f if line.strip()]\n\ndef normalize(name):\n # keep only letters a–z\n return re.sub(r'[^a-z]', '', name)\n\n# Clean and filter\nnames = [n for n in (normalize(name) for name in names) if len(n) >= 3]\nnames = list(set(names)) #Now we have a set of uniques namesAfter preprocessing, we are left with around 15,000 names, which is plenty of data to estimate bigram statistics and to train our first simple models.
\nTokens are the atomic units of a language model. In this project, each character is a token, so the vocabulary is simply the set of unique letters in the preprocessed dataset. To handle sequences cleanly, we extend this set with two boundary markers: ! for start-of-sequence (SOS) and ? for end-of-sequence (EOS).\nFormally, we construct a bijection
\n\\text{stoi} : \\Sigma \\cup \\{\\text{SOS}, \\text{EOS}\\} \\;\\to\\; \\{0, \\ldots, V-1\\}and its inverse itos. Encoding applies stoi elementwise to a string; decoding applies itos to a list of indices.\nHere’s the code:
\n# Build character set from dataset\ncharacters = sorted(set(\"\".join(names)))\n\n# Define special tokens first \nstoi = {\"!\": 0} # SOS\nfor i, ch in enumerate(characters, start=1):\n stoi[ch] = i\nstoi[\"?\"] = len(stoi) # EOS\n\nitos = {i: ch for ch, i in stoi.items()}\nvocab_size = len(stoi)\nprint(f\"Vocabulary size: {vocab_size}\")To illustrate how encoding and decoding work with these mappings, consider the example of the name \"tommaso\":
\nname = \"!tommaso?\"\n\n# Encode: convert each character to its index\nencoded = [stoi[ch] for ch in name]\nprint(\"Encoded:\", encoded)\n\n# Decode: convert indices back to characters\ndecoded = ''.join([itos[ix] for ix in encoded])\nprint(\"Decoded:\", decoded)The output will be:
\nEncoded: [0, 20, 15, 13, 13, 1, 19, 15, 27]\nDecoded: !tommaso?This mapping turns text into a numerical sequence and back again. By inserting SOS and EOS during training, the model learns both when names start and when they should stop — essential for generating coherent results.
\nTo evaluate a model fairly, we split the dataset into disjoint parts: a training set to fit parameters and a test set to measure generalization. For this experiment, we use an 80/20 split. Larger projects often add a separate validation set for hyperparameter tuning, but two splits suffice here.\nFormally, each name is converted into a sequence of pairs (X_t, Y_t), where
X_t = x_{t-1}, \\quad Y_t = x_t, \\quad \\text{with } x_0 = \\text{SOS}, \\; x_{n+1} = \\text{EOS}.Each pair encodes the prediction task: “given the current character, predict the next one.”\nImplementation in PyTorch:
\nimport torch, random\n\ndef create_dataset(names, block_size=1):\n X, Y = [], []\n for name in names:\n prev_ix = stoi[\"!\"] # SOS\n for ch in name + \"?\": # append EOS\n ix = stoi[ch]\n X.append(prev_ix)\n Y.append(ix)\n prev_ix = ix\n return torch.tensor(X, dtype=torch.long), torch.tensor(Y, dtype=torch.long)\n\n# Train/test split \nnames_shuffled = names[:]\nrandom.shuffle(names_shuffled)\n\nsplit_idx = int(0.8 * len(names_shuffled))\nXtr, Ytr = create_dataset(names_shuffled[:split_idx])\nXtst, Ytst = create_dataset(names_shuffled[split_idx:])At this point, we have compact tensor representations of the training and test data, ready to be fed into the bigram model.\nTo make this concrete, let’s look at how the name \"tommaso\" is represented:
\nX, Y = create_dataset([\"tommaso\"])\nfor i in range(len(X)):\n print(itos[X[i].item()], \"-->\", itos[Y[i].item()])Output:
\n! --> t\nt --> o\no --> m\nm --> m\nm --> a\na --> s\ns --> o\no --> ?This example illustrates the idea: each pair (X, Y) captures a step in the sequence, mapping the current character (or \\text{SOS} at the start) to the next one. The final transition predicts \\text{EOS}, teaching the model when to stop generating.
With the training pairs in hand, we can now estimate conditional probabilities.\nThe first step is to build a V \\times V count matrix, where V is the vocabulary size and
C[i,j] = \\#\\{\\text{times token }j\\text{ follows token }i\\}.This matrix encodes the immediate statistical structure of the dataset. By maximum likelihood estimation (MLE) for a categorical distribution:
\n\\hat{P}(j \\mid i) = \\frac{C[i,j]}{\\sum_k C[i,k]}.Intuitively, the chance of seeing j after i is just its observed frequency.
# Bigram counts matrix: V x V where C[i, j] = times char j follows char i\ncounts = torch.zeros(len(itos), len(itos))\n\nfor i, j in zip(Xtr, Ytr):\n counts[i][j] += 1\n\n# Add-one smoothing to avoid zero probabilities ([additive smoothing](https://en.wikipedia.org/wiki/Additive_smoothing))\ncounts += 1Resulting counts (heatmap visualization):
\n![\"Counts\"](\"https://www.tommasovaccari.com/static/counts-7504e101.webp\")
The counts heatmap shows raw transition frequencies: brighter cells indicate more common bigrams\n(row = current character, column = next character, i.e. row i, column j corresponds to C[i,j]).
# Normalize each row to sum to 1\nprobs = counts / counts.sum(1, keepdim=True)Each row of probs is a categorical distribution over the next token, and Laplace smoothing ensures even rare contexts yield valid probabilities.\nFinal probability distribution(heatmap visualization):
\n![\"Probs\"](\"https://www.tommasovaccari.com/static/probs-6087aac8.webp\")
The probability heatmap shows per-row normalized transitions: each row sums to 1, highlighting the most likely successors for every character.
\nNow that we have probabilities, we can sample new names. The procedure is:
\ndef generate(n_samples, maxlen=15):\n for _ in range(n_samples):\n context = 0 # SOS\n s = \"\"\n for _ in range(maxlen):\n prob = probs[context] # distribution over next char\n idx = torch.multinomial(prob, 1).item()\n context = idx\n ch = itos[idx]\n if ch == \"?\":\n break\n s += ch\n print(s)Example Output:
\neli\ndafo\nada\nhaariva\nanna\nminla\ngila\ncoudMost outputs are noisy or unrealistic, which is expected: the bigram model can only capture relationships between adjacent characters. Still, some generated names look plausible, such as:
\nThis shows the model has picked up some local character patterns, but it fails to capture longer-range dependencies needed for realistic names.
\nA natural idea is to extend the context, e.g., moving from bigrams to trigrams. But this quickly becomes impractical: the size of the count matrix scales as O(V^n), growing polynomially with V and exponentially with n, and probabilities become sparse as context length increases. The true solution lies in a different direction — neural networks, which we will explore in the next section.
The bigram model was just a lookup table of observed transitions. A neural model reframes this idea as a trainable function.
\nThe setup is unchanged: given a context (initially one previous character), the model predicts the next one.\nThe difference is representation. Instead of fixed rows in a matrix, we introduce embeddings: a learnable map
\n\\text{Emb} : \\{0, \\ldots, V-1\\} \\;\\to\\; \\mathbb{R}^d,which assigns each token index a vector in a d-dimensional continuous space.
Architecture sketch:
\nAfter retrieving an embedding vector for the input token, the next step is to transform it into a representation that can be compared against every possible output class. This is done through a linear transformation:
\nz = W x + b\nwhere
\nx \\in \\mathbb{R}^d is the embedding of the current token,W \\in \\mathbb{R}^{V \\times d} is a weight matrix mapping the d-dimensional embedding space to the V output classes,b \\in \\mathbb{R}^V is a bias vector, andz \\in \\mathbb{R}^V are the resulting logits, one score for each vocabulary item.This step assigns a learnable score to each possible next token.
\nIf the model consisted only of embeddings and a single linear transformation, it would still be a linear function of the input indices. In fact, a single linear layer is mathematically equivalent to a lookup table — expressive enough for bigrams, but fundamentally limited. With one-hot input vectors, a linear transformation W x just selects a row of W. This is functionally identical to table lookup. Only nonlinearities break this equivalence.
To break this limitation, we introduce a nonlinear activation function between layers. In practice, we use the ReLU (Rectified Linear Unit):
\n\\text{ReLU}(z) = \\max(0, z)\nReLU has two critical effects:
\nTogether, the linear layer provides a weighted combination of input features, while the nonlinearity ensures the model is not just a rescaled lookup table. This combination forms the backbone of neural networks and prepares the logits for conversion into probabilities via the softmax function.
\nThe linear layer produces a vector of real-valued scores, or logits, one per token in the vocabulary.\nThese scores are not probabilities. To turn them into a distribution, we apply the softmax function:
\n\\text{softmax}(z)_j \\;=\\; \\frac{e^{z_j}}{\\sum_{k=1}^V e^{z_k}}, \\quad j = 1, \\ldots, Vwhere z \\in \\mathbb{R}^V is the logits vector and V is the vocabulary size.
Softmax has three key properties:
\nAt a statistical level, we assume the data is generated by an unknown distribution P_{\\text{data}} over classes \\{1, \\dots, V\\}.\nWe define a parametric model P_\\theta, with parameters \\theta, that assigns probabilities to these classes.\nTraining aims to choose \\theta so that P_\\theta approximates P_{\\text{data}} as closely as possible.
The principle of maximum likelihood estimation (MLE) chooses \\theta to maximize the probability of the observed data.\nGiven observations \\{y^{(1)}, \\dots, y^{(N)}\\}:
L(\\theta) = \\prod_{n=1}^N P_\\theta\\big(y^{(n)}\\big)\\ell(\\theta) = \\sum_{n=1}^N \\log P_\\theta\\big(y^{(n)}\\big)The MLE estimator is:
\n\\hat{\\theta} = \\arg \\max_\\theta \\, \\ell(\\theta).Maximizing log-likelihood is equivalent to minimizing the negative log-likelihood (NLL):
\n\\mathcal{L}_{\\text{NLL}}(\\theta) = - \\sum_{n=1}^N \\log P_\\theta\\big(y^{(n)}\\big)For classification with one-hot targets y^{(n)} and model predictions p_\\theta^{(n)}, this becomes:
\\mathcal{L}(y, p_\\theta) = - \\sum_{j=1}^V y_j \\, \\log p_{\\theta, j}This is exactly the cross-entropy between the empirical distribution y and the model distribution p_\\theta.\nMinimizing cross-entropy is equivalent to minimizing the KL divergence:
D_{\\text{KL}}(P_{\\text{data}} \\,||\\, P_\\theta),the gold standard of statistical approximation.
\nIn the count-based model, MLE had a closed form: normalize counts.\nIn a neural model, the mapping x \\mapsto p_\\theta is nonlinear, so no closed form exists.\nInstead, we compute gradients:
\\nabla_\\theta \\, \\mathcal{L}_{\\text{NLL}}(\\theta)and update parameters iteratively using optimization algorithms such as gradient descent and stochastic gradient descent (SGD), or Adam.\nThis drives P_\\theta toward the distribution that maximizes the likelihood of the training data.
Suppose our model outputs logits for V = 3 classes:
z = [2, 1, 0]\nStep 1: Apply Softmax
\n\\text{softmax}(z) = [0.67, 0.24, 0.09]\nStep 2: Define the True Label
\nLet the correct class be y^* = 2, i.e. y = [0, 1, 0].
Step 3: Compute Cross-Entropy Loss
\n\\mathcal{L}(y, p) = - \\log p_{y^*} = - \\log 0.24 \\approx 1.43\nInterpretation:
\n0.24, yielding a relatively high loss.1, the loss would approach 0.p_{y^*}, reducing loss over time.The count matrix was a fixed table; the neural model is a learnable table. With embeddings, nonlinear layers, and cross-entropy optimization, we generalize the same principle — conditional probabilities — into a form that scales.
\nWe now move from bigrams (context length 1) to a fixed-width context window of length 3. This allows the model to consider not just the immediately preceding character but a short history. The workflow has three parts:
\nConventions.\n! = start-of-sequence (SOS, index 0),\n? = end-of-sequence (EOS),\nblock_size = 3 = context length.\nAll tensors live on device.
\nThe dataset construction mirrors the bigram setup but with a sliding window of length 3:
\nblock_size = 3\n\ndef create_dataset(names):\n X, Y = [], []\n for name in names:\n context = [0] * block_size # SOS padding\n for ch in name + \"?\":\n ix = stoi[ch]\n X.append(context[:]) # copy current context\n Y.append(ix)\n context = context[1:] + [ix]\n return torch.tensor(X, dtype=torch.long), torch.tensor(Y, dtype=torch.long)\n\nXtr, Ytr = create_dataset(names[:split_idx])\nXtst, Ytst = create_dataset(names[split_idx:])The neural network extends the count-based model into a parameterized function. Its architecture:
\nvocab_size = len(itos)\nembedding_dim = 32\nhidden_dim = 128\nblock_size = 3 # context length\n\nclass Net(nn.Module):\n def __init__(self, vocab_size, embedding_dim, block_size=3, hidden_dim=128):\n super().__init__()\n self.emb = nn.Embedding(vocab_size, embedding_dim) # (V, d)\n self.fc1 = nn.Linear(block_size * embedding_dim, hidden_dim) # (3d → H)\n self.fc2 = nn.Linear(hidden_dim, vocab_size) # (H → V)\n\n def forward(self, x):\n \"\"\"\n x: (B, block_size) int64\n returns: logits (B, V)\n \"\"\"\n e = self.emb(x) # (B, block_size, d)\n e = e.view(e.size(0), -1) # flatten to (B, block_size*d) [docs](https://pytorch.org/docs/stable/generated/torch.Tensor.view.html)\n h = F.relu(self.fc1(e)) # (B, H)\n logits = self.fc2(h) # (B, V)\n return logits\n\nmodel = Net(vocab_size, embedding_dim, block_size, hidden_dim).to(device)batch_size = 256\noptimizer = torch.optim.Adam(model.parameters(), lr=3e-3, weight_decay=1.5e-4)Optimization logic:\nFor each mini-batch, compute logits → cross-entropy → scalar loss. Backpropagation computes \\nabla_\\theta \\mathcal{L}. Adam then updates parameters, inching the model closer to the distribution that maximizes likelihood of the observed data.
This is where the model learns. We run stochastic gradient descent (SGD) in mini-batches, using random subsets of training data to approximate gradients efficiently while introducing noise that improves generalization.
\nmodel.train()\nepochs = 3 # Chosen to balance training stability with runtime; increase for stronger convergence.\niters_per_epoch = 10_000 # Chosen to balance training stability with runtime; increase for stronger convergence.\nloss_trace = []\n\nN = Xtr.size(0)\n\nfor _ in range(epochs):\n for _ in range(iters_per_epoch):\n idx = torch.randint(0, N, (batch_size,), device=device) # (B,)\n Xb = Xtr[idx] # (B, 3)\n yb = Ytr[idx] # (B,)\n\n logits = model(Xb) # (B, V)\n loss = F.cross_entropy(logits, yb)\n loss_trace.append(loss.item())\n\n optimizer.zero_grad(set_to_none=True) # [docs](https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html)\n loss.backward() # backpropagation\n optimizer.step() # parameter update\n\n print(f\"train loss (last batch): {loss.item():.4f}\")Below is the training loss curve over iterations:\n
The loss curve decreases overall but exhibits some variance, indicating that further hyperparameter tuning could improve performance. The initial ‘hockey-stick’ shape might also be mitigated with a better weight initialization strategy. However, for our purposes, we can proceed with the current setup and keep the approach simple.
\nEvaluation is done on unseen data. We disable gradient tracking with torch.no_grad, switch to evaluation mode, and compute both loss and accuracy.
\n@torch.no_grad()\ndef eval_minibatch():\n model.eval()\n idx = torch.randint(0, Xtst.size(0), (4096,), device=device)\n logits = model(Xtst[idx]) # (B, V)\n y = Ytst[idx] # (B,)\n loss = F.cross_entropy(logits, y).item()\n acc = (logits.argmax(1) == y).float().mean().item() # [docs](https://pytorch.org/docs/stable/generated/torch.argmax.html)\n model.train()\n return loss, acc\n\nval_loss, val_acc = eval_minibatch()\nprint(f\"val loss: {val_loss:.4f} | val acc: {val_acc:.3f}\")With this setup, the model is no longer just a table of counts. It is a trainable system that builds its own internal representations. Even with a fixed 3-character window, we can already see the core ingredients of modern language models at work: embeddings, nonlinear layers, and iterative optimization.
\nTraining gives us losses and accuracies, but sampling turns those numbers into names. Generation is just sampling: repeatedly draw the next token from the model’s conditional distribution until EOS is reached.
\n@torch.no_grad()\ndef generate(n_samples=10, maxlen=15):\n model.eval()\n V = len(itos)\n\n for _ in range(n_samples):\n \n ctx = torch.zeros((1, block_size), dtype=torch.long, device=device) \n s = []\n\n for _ in range(maxlen):\n logits = model(ctx) # (1, V) logits\n\n # convert logits into probabilities\n probs = F.softmax(logits, dim=1) # (1, V) [docs](https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html)\n\n # sample from the distribution\n next_idx = torch.multinomial(probs, 1) # (1, 1)\n ix = next_idx.item()\n\n ch = itos[ix]\n if ch == \"?\":\n break\n s.append(ch)\n\n # slide the window forward\n ctx = torch.roll(ctx, shifts=-1, dims=1) # still (1, block_size) [docs](https://pytorch.org/docs/stable/generated/torch.roll.html)\n ctx[0, -1] = ix\n\n print(\"\".join(s))Generation proceeds by starting with SOS tokens, repeatedly sampling from the softmax distribution, and shifting the context forward until EOS or max length is reached.
\naskiem\nalerindro\nrichardin\nfulayma\nmarian\njhona\nnodou\nchris\nsaurie\nnardi\nenni\ngiadan\nserena\nbessa\nemmaSome outputs (marian, serena, emma) are plausible Italian names; others (askiem, fulayma) are hybrids that never appeared in training. This balance is exactly what we expect: the model has internalized local character regularities but lacks the longer memory needed for fully natural names. Sampling exposes both strengths and limits: the model generalizes beyond memorization but still lacks long-range structure.
\nWith both the count-based bigram model and the neural 3-gram model trained, we can now compare them quantitatively.
\nOn the held-out test set:
\n# Bigram test loss\nwith torch.no_grad():\n log_probs = torch.log(probs[Xtst, Ytst])\n loss_bigram = -log_probs.mean().item()\n\n# Neural 3-gram test loss \nval_loss = F.cross_entropy(model(Xtst), Ytst).item()\nprint(loss_bigram, val_loss)Results:
\nThe reduction of ~0.4 bits per character means the trigram model assigns systematically sharper probabilities, leading to fewer mistakes across the sequence. Even small cross-entropy reductions compound over long sequences.
\nIn short, language modeling is a trade-off between expressiveness (how much structure longer context can capture) and feasibility (how much compute and data are required). Modern models (RNNs → Transformers) push this frontier by scaling effective context to entire sequences.
\nWe set out to generate plausible Italian names via next-character prediction. From count-based bigrams to a neural trigram, we made the math explicit (MLE, softmax, cross-entropy) and showed how fixed counts become learnable functions through embeddings and nonlinear layers.
\nAt heart, a language model factorizes
\nP(x_1, \\dots, x_n) \\;=\\; \\prod_{t=1}^n P(x_t \\mid x_{<t}),with effectiveness governed by how much context it conditions on. (Chain rule)
\nLanguage modeling has always fought with context. Classical n-grams blow up combinatorially, while small neural nets manage only modest extensions. Modern models—RNNs and especially Transformers—address this by reusing and sharing parameters, scaling context from a handful of characters to entire documents.
\nNext up: attention—the mechanism that lets models capture dependencies across any distance without the n-gram blow-up. In the following post we’ll unpack Attention Is All You Need and see how it unlocks the long-range structure behind today’s transformers.
","date_published":"2025-09-08T00:00:00.000Z","tags":["NeuralNetworks","Math","Statistics","Python"]},{"id":"https://www.tommasovaccari.com/blog/breaking-sorting-barrier-sssp","url":"https://www.tommasovaccari.com/blog/breaking-sorting-barrier-sssp","title":"Breaking the Sorting Barrier in Single-Source Shortest Paths","summary":"Exploring the ideas behind the new breakthrough paper, breaking the long-standing sorting barrier for single-source shortest paths.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"Today I want to explore the breakthrough ideas presented in the paper Breaking the Sorting Barrier for Directed Single-Source Shortest Paths . Before diving into those concepts, it’s important to briefly introduce Dijkstra’s algorithm and its purpose.
\n![\"A](\"https://www.tommasovaccari.com/static/router-graph-d070f87f.webp\")
In many practical applications, you work with a directed graph G=(V,E), where V represents vertices and E the edges. Each edge is often associated with a weight (for our purposes, these must be positive—we will see later why this restriction matters). These weights can model:
The fundamental task is to find the shortest path (in terms of total weight) from a given source to all other vertices. This problem is known as the Single-Source Shortest Path (SSSP) problem.
\nSSSP remains central in many domains, including:
\nIn 1956, Edsger W. Dijkstra devised an elegant solution to this problem. Three years later, he formally published it in his influential paper “A Note on Two Problems in Connexion with Graphs” (1959).
\nThis contribution was historic: it provided a clean and efficient algorithm for a problem that seemed computationally challenging at the time. Given that computers in the 1950s were extremely limited in speed and memory, Dijkstra’s solution was not just clever but fundamental.
\nDijkstra’s original form of the algorithm became a cornerstone of computer science. Later decades brought refinements—through improved data structures, heuristic methods, and adaptations for massive graphs—that extended its reach. Yet none have surpassed the foundational impact of Dijkstra’s 1959 publication.
\nBecause of the importance of the SSSP problem, research has continued for decades to improve the time complexity of shortest-path algorithms. And now, after more than sixty years, we have reached a new frontier: breaking the long-standing sorting barrier for directed SSSP.
\nIn this section, we review how Dijkstra’s algorithm works at a high level. The goal is simple: compute the shortest path from a given source to every other vertex in the graph.
\nFormally, consider a graph G=(V,E), where V is the set of vertices and E is the set of edges. Each vertex stores a value w, representing its current best-known distance from the source, and a list of its adjacent vertices.
Initialization.
\n0.+∞.S of processed vertices (conceptually useful, though not always explicit in implementations).Q containing all vertices, keyed by their distance w.Main loop.
\nWhile Q is not empty:
Q the vertex u with the smallest w.u to S.v of u, perform a relaxation step: update v’s distance if the path through u is shorter than its current value.Relaxation means checking whether:
\nw(v) > w(u) + weight(u,v) and, if so, updating w(v).
Pseudocode:
\nDijkstra(G,source):\n Initialize-Graph(G, source)\n S = EmptyList()\n Q = List(all vertices in G)\n\n while Q is not empty:\n u = ExtractMin(Q)\n S.append(u)\n for each vertex v adjacent to u:\n Relax(u, v)\nBecause the algorithm always chooses the vertex with the smallest distance, it is classified as a greedy algorithm. Greedy methods do not always guarantee correctness in general, but in this case correctness can be formally proved using induction on the size of S.
Complexity.\nIn Dijkstra’s original implementation, the priority structure Q was simply an array. Extracting the minimum then required scanning the entire array, which costs O(V) per operation. Since every vertex must be extracted once, this contributes O(V^2)
The inner relaxation loop processes each edge at most once, costing O(E).
Thus, the overall complexity is: O(V^2 + E)
The bottleneck in Dijkstra’s original implementation is the Extract-Min operation: finding the vertex with the smallest tentative distance. Using a simple array makes this O(V) which quickly dominates runtime on larger graphs.
A major breakthrough came with the introduction of binary heaps as the priority queue for Q (widely popularized in the 1970s). With this change:
O(\\log V).O(\\log V).O(E \\log V) (often written as O((V+E)\\log V)), a dramatic improvement for sparse graphs.On the theoretical side, Fibonacci heaps (Fredman & Tarjan, 1984) pushed the asymptotics even further: O(E + V \\log V). Here, Decrease-Key runs in amortized O(1). While elegant in theory, Fibonacci heaps are rarely used in practice because of pointer-heavy structures, large constant factors, and poor cache performance. This highlights a recurring lesson in algorithms: asymptotic optimality does not always translate into real-world efficiency.
For graphs with bounded integer weights, an alternative approach is Dial’s algorithm (1969). It uses an array of buckets indexed by distance values, achieving: O(V + E + C) where C is the maximum edge weight. On road networks or latency graphs with small integer costs, this method can be extremely efficient.
Other heap structures — pairing heaps, radix heaps, relaxed heaps — have been proposed over the years, each balancing theoretical guarantees and practical constants in different ways. Yet despite decades of improvements, progress has always seemed tethered to the same barrier: the cost of sorting-like operations lurking inside priority queue management.
\nThe 2025 paper Breaking the Sorting Barrier for Directed Single-Source Shortest Paths by Ran et al. presents the first deterministic algorithm that beats the classical:
\nO(|E| + |V| \\log |V|)bound for directed SSSP with nonnegative real weights in the comparison–addition model.\nIt is important to state explicitly that the result holds in this model, since it rules out shortcuts that rely on specialized data structures or RAM-model tricks.
\nTheir result achieves:
\nO(|E| \\, \\log^{2/3} |V|).This is not a wholesale replacement for Dijkstra’s algorithm. If you require the full ordering of vertices by distance, Dijkstra remains optimal. If you only need the distance values themselves, the long-assumed “sorting barrier” is not inherent.
\nThe analysis assumes two standard technical moves: (i) a degree-reduction transformation so every vertex has constant in/out degree, and (ii) a lexicographic tie-breaking rule so that every path has a unique length. Both are common but necessary for the recursion to work cleanly.
\nThe bottleneck in Dijkstra’s algorithm is maintaining a strict total order over a frontier of tentative vertices.\nAt times this frontier can be size \\Theta(|V|), forcing \\Omega(\\log |V|) work for each extract-min.\nThat’s where the extra \\log |V| factor comes from.
The new algorithm avoids this by working in bands [b,B).\nSuppose all vertices with d(v) < b are complete, and let S be the current frontier of vertices with tentative distances in [b,B).\nDefine the covered set:
\\tilde{U} = \\{ v \\in V : d(v) < B \\ \\text{and the shortest path } s \\to v \\text{ passes through some } u \\in S \\}.Every unfinished vertex with d(v) < B lies in \\tilde{U}.
Now, depending on the size of \\tilde{U} relative to |S|, the algorithm takes one of two actions, controlled by a parameter k:
|\\tilde{U}| \\geq k \\cdot |S|, a bounded multi-source shortest path (BMSSP) call finalizes all vertices in \\tilde{U} with d(v) < B in one shot.k rounds of relaxations from S. Vertices whose shortest paths use fewer than k frontier vertices are finalized.\nThe remaining ones each select a pivot in S; the number of pivots is at most |\\tilde{U}| / k.\nThis shrinks the frontier by a factor of about k.In either case, progress is guaranteed: either a large batch of vertices is finalized, or the frontier becomes much smaller.
\nThe performance of the recursion depends on two parameters:
\nk = \\lfloor \\log^{1/3} V \\rfloor, \\qquad\nt = \\lfloor \\log^{2/3} V \\rfloor.The parameter k determines how aggressively the frontier can be reduced in the pivoting step. In the worst case, the number of pivots introduced in one round is bounded by |\\widetilde{U}|/k, which contracts the frontier by a factor of \\Theta(k). The parameter t specifies the band width [b,B), and hence the distance threshold within which the bounded multi-source shortest path (BMSSP) subroutine operates.
The recursion proceeds in layers: each layer either finalizes a large fraction of \\widetilde{U} via BMSSP or reduces the frontier by a factor of about k. Since the frontier shrinks by \\Theta(k) in a reduction step, the number of recursive levels is at most
\\frac{\\log V}{\\log k} = O\\!\\left(\\frac{\\log V}{\\log^{1/3} V}\\right) = O(\\log^{2/3} V).Within each level, every edge is relaxed at most once, and the per-vertex overhead due to frontier reduction is bounded by
\n\\frac{\\log V}{k} = O(\\log^{2/3} V).Combining these observations, the overall running time is
\nO(|E| \\log^{2/3} V).The balance between k and t is crucial: larger k accelerates frontier reduction but increases per-level cost, while larger t decreases recursion depth but inflates the band width and the work of BMSSP. The chosen values k = \\log^{1/3} V and t = \\log^{2/3} V minimize the total overhead and yield the \\log^{2/3} V factor.
| Algorithm | Time complexity | Applicability |
|---|---|---|
| Dijkstra (array) | O(V^2 + E) | Historical baseline; easy to implement but quadratic on large graphs |
| Dijkstra (binary heap) | O(E \\log V) | Standard practical choice; efficient for sparse graphs and widely deployed |
| Ran et al. (2025) | O(E \\log^{2/3} V) | First deterministic algorithm to beat the O(E \\log V) bound in the comparison–addition model |
The 2025 result is not a practical replacement for heap-based Dijkstra. Its significance lies in theory: it shows that the \\log n factor in SSSP is not inherent. The breakthrough is stated in the comparison–addition model, where algorithms are only allowed to compare and add edge weights. This model is the standard baseline for analyzing general-weight graph algorithms, since it rules out shortcuts that depend on specialized word-RAM techniques or bounded integers. Citing it makes clear that the improvement holds in the broadest, most widely accepted setting—directly comparable to Dijkstra’s classical O(E \\log V) bound.
The single-source shortest path problem has been studied for almost seventy years, from Dijkstra’s elegant O(V^2+E) algorithm to heap-based improvements and practical variants like Dial’s method. For decades, every faster deterministic algorithm seemed shackled by the same bottleneck: the cost of maintaining a sorted frontier.
Ran et al. (2025) show that this barrier is not inherent. By sidestepping the need for a global order and instead settling vertices in controlled batches, their algorithm replaces the \\log |V| overhead with \\log^{2/3} |V|. This marks the first deterministic improvement to Dijkstra’s runtime in the comparison–addition model.
The breakthrough is not just about speed—it’s about perspective. Shortest paths can be computed without ever maintaining a total order of tentative distances, reshaping how we think about the greedy structure of SSSP and opening doors to further refinements.
\nI hope you enjoyed the post and that it gave you a grasp of SSSP algorithms!
\nIn this second post, we will take another step forward. We’ll dive into the concept of linear regression, explore its connection to neural networks, and demonstrate how to build a model from scratch. This approach will help us gain a deeper understanding of the core concepts behind both linear regression and neural networks.
\nPrerequisites: PyTorch, NumPy, and Matplotlib installed (for plotting and numerical operations).
\nBefore starting, we need to define some terminology. When we want to predict a value, we call it a label or target. Each label is associated with its own features. The model makes predictions based solely on these features. Therefore, we need to refine the model to accurately predict the correct label based on the provided input features.
\nWe want to predict a target, \\hat{y}, based on a set of features grouped in a vector called X. To make the prediction, we need to find the weight vector W and the bias B that give us the most accurate predictions. This relationship can be expressed using the dot product:
\\hat{y} = W \\cdot X + bLet X be the feature matrix (rows are samples, columns are features). Then the model’s predictions for all samples can be written compactly as a vector:
\\hat{y} = X \\cdot W + bNow we need to define a loss function to evaluate the model. The most common is the squared error, where \\hat{y}_i is the prediction and y_i is the corresponding true label.
L_i(W, b)= \\frac{1}{2} (\\hat{y}_i-y_i)^2We can observe that the loss is a function of weight and bias.
\nNow we can extend the evaluation of the loss along all the predictions, averaging it to obtain an intuitive idea of how our model is performing.
\nL_i (W,B) =\\frac{1}{n} \\sum_{i=1}^{n} L_i(W,B) = \\frac{1}{n} \\sum_{i=1}^{n} \\frac{1}{2} (X^{(i)} \\cdot W + b - y^{(i)})^2We have to keep in mind our goal: we want to find W and B that give us the minimum value of loss along all the predictions.
\nThe key to optimizing a model and improving its predictions is to iteratively update the weights by modifying them in the opposite direction of the gradient of the loss function. This process is known as gradient descent. Although it may seem simple at first glance, it is the foundation of many advanced techniques in machine learning and plays a central role in training modern models.
\nTo optimize the process, we do not update the model using the entire dataset at once. Instead, we randomly select smaller subsets of the data, called batches. The model is then optimized iteratively by performing updates on these batches, which makes the process more efficient and scalable, especially for large datasets. This approach is commonly referred to as mini-batch gradient descent.
\nTo keep it simple we can break down this process into four steps:
\nIf we want to express this in formulas we have:
\n\\mathbf{w} \\gets \\mathbf{w} - \\frac{\\eta}{|\\mathcal{B}|} \\sum_{i \\in \\mathcal{B}_t} \\frac{\\partial L^{(i)} (\\mathbf{w}, b)}{\\partial \\mathbf{w}} = \\mathbf{w} - \\frac{\\eta}{|\\mathcal{B}|} \\sum_{i \\in \\mathcal{B}_t} \\left( \\mathbf{x}^{(i)} \\cdot \\mathbf{w} + b - y^{(i)} \\right) \\mathbf{x^{(i)}}.b \\gets b - \\frac{\\eta}{|\\mathcal{B}|} \\sum_{i \\in \\mathcal{B}_t} \\frac{\\partial L^{(i)} (\\mathbf{w}, b)}{\\partial b} \n= b - \\frac{\\eta}{|\\mathcal{B}|} \\sum_{i \\in \\mathcal{B}_t} \\left( \\mathbf{x}^{(i)} \\cdot \\mathbf{w} + b - y^{(i)} \\right).While linear models are not sufficiently rich to express complex relationships between features, we can introduce neural networks to obtain a more expressive model. Nevertheless, we can also view a linear model as a neural network where every input feature corresponds to a neuron with its own weight and bias.
\nTo gain a deeper understanding of how a model is created and trained, we aim to implement the classes and methods from scratch. We are following this approach because, in my opinion, using machine learning libraries like PyTorch directly does not provide a comprehensive understanding of what happens under the hood.
\n%matplotlib inline\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torchNow we are ready to implement our model for linear regression from scratch. We need:
\nclass LinearRegressionModel:\n def __init__(self, num_inputs, learning_rate, sigma=0.01):\n \"\"\"\n Initialize the model parameters.\n\n Args:\n - num_inputs (int): Number of input features.\n - learning_rate (float): Learning rate for gradient descent.\n - sigma (float): Standard deviation for initializing weights.\n \"\"\"\n self.num_inputs = num_inputs\n self.learning_rate = learning_rate\n \n # Initialize weights and bias\n self.w = torch.normal(mean=0.0, std=sigma, size=(num_inputs, 1), requires_grad=True)\n self.b = torch.zeros(1, requires_grad=True)\n Here we have created a class that contains the weights and bias. Additionally, we have introduced the hyperparameters, which are typically user-defined and used to adjust various aspects during the training phase. The weights are sampled from a normal distribution with a mean of 0 and a standard deviation of 1, while the bias is initialized to zero.
\nNow we can add the method to obtain the forward pass:
\n@add_to_class(LinearRegressionModel)\ndef forward(self, X):\n \"\"\"\n Compute the forward pass: y = Xw + b.\n\n Args:\n - X (torch.Tensor): Input tensor of shape (batch_size, num_inputs).\n \n Returns:\n - torch.Tensor: Predicted values of shape (batch_size, 1).\n \"\"\"\n return torch.matmul(X, self.w) + self.bNow we add the method to calculate the loss for a single batch of examples by using the squared loss function and then returning the mean.
\n@add_to_class(LinearRegressionModel)\ndef compute_loss(self, y_pred, y_true):\n \"\"\"\n Compute Mean Squared Error loss.\n\n Args:\n - y_pred (torch.Tensor): Predicted values.\n - y_true (torch.Tensor): True values.\n \n Returns:\n - torch.Tensor: Scalar loss value.\n \"\"\"\n return 0.5 * ((y_pred - y_true) ** 2).mean()This is the fundamental part of our model—the algorithm that allows us to improve its predictions. We are going to implement Stochastic Gradient Descent (SGD), as discussed earlier.
\nThe steps we want to follow are:
\n@add_to_class(LinearRegressionModel)\ndef update_parameters(self):\n \"\"\"\n Update the model parameters using gradient descent.\n \"\"\"\n with torch.no_grad():\n self.w -= self.learning_rate * self.w.grad\n self.b -= self.learning_rate * self.b.grad\n\n # Manually zero the gradients\n self.w.grad.zero_()\n self.b.grad.zero_()\n\ndef train_step(self, X, y, batch_size):\n \"\"\"\n Perform a single training step.\n\n Args:\n - X (torch.Tensor): Input data of shape (num_samples, num_inputs).\n - y (torch.Tensor): Target data of shape (num_samples, 1).\n - batch_size (int): Number of samples per batch.\n \n Returns:\n - float: Loss value for the batch.\n \"\"\"\n # Sample a random batch\n num_samples = X.shape[0]\n indices = torch.randint(0, num_samples, (batch_size,))\n X_batch = X[indices]\n y_batch = y[indices]\n\n # Forward pass\n y_pred = self.forward(X_batch)\n\n # Compute loss\n loss = self.compute_loss(y_pred, y_batch)\n\n # Backward pass\n loss.backward()\n\n # Update parameters\n self.update_parameters()\n\n # Return the loss value as a scalar\n return loss.item()\nThe last step to complete our model is to implement a method that allows us to train it. It would be useful to visualize how the weights, bias, and learning rate change throughout the training process, so we can also add a method to plot this information.
\n@add_to_class(LinearRegressionModel)\ndef train(self, X, y, epochs, batch_size):\n \"\"\"\n Train the model over multiple epochs.\n\n Args:\n - X (torch.Tensor): Input data of shape (num_samples, num_inputs).\n - y (torch.Tensor): Target data of shape (num_samples, 1).\n - epochs (int): Number of training epochs.\n - batch_size (int): Number of samples per batch.\n \n Returns:\n - list: List of loss values for each epoch.\n \"\"\"\n losses = []\n\n for epoch in range(epochs):\n # Perform a training step and compute the average loss for the epoch\n loss = self.train_step(X, y, batch_size)\n losses.append(loss)\n\n # Print progress every 10 epochs\n if epoch % 10 == 0:\n print(f\"Epoch {epoch}, Loss: {loss:.6f}\")\n\n return losses\n Now that the class is built from scratch, we can proceed with testing it to verify if everything works correctly. Testing the model typically involves the following steps:
\nNow, we initialize a vector to represent the weights that our model needs to learn from the data during the training process, and we follow the same process for the bias.\nThen we compute the target that we are going to use, with the features to train our model.\nWe can also adjust hyperparameters as needed.
\n# Define hyperparameters\nnum_samples = 4000\nnum_inputs = 2000\nlearning_rate = 0.01\nepochs = 2000\nbatch_size = 64\n\n# Generate synthetic data\ntrue_w = torch.randn((num_inputs, 1))\ntrue_b = torch.rand((1,1))\nX = torch.randn((num_samples, num_inputs))\ny = torch.matmul(X, true_w) + true_b + torch.randn((num_samples, 1)) * 0.01num_samples rows and num_inputs columns.X with true_w, adding true_b, and adding a small amount of random noise to simulate real-world data.Now we can look at the complete class with all the methods:
\nclass LinearRegressionModel:\n def __init__(self, num_inputs, learning_rate, sigma=0.01):\n \"\"\"\n Initialize the model parameters.\n\n Args:\n - num_inputs (int): Number of input features.\n - learning_rate (float): Learning rate for gradient descent.\n - sigma (float): Standard deviation for initializing weights.\n \"\"\"\n self.num_inputs = num_inputs\n self.learning_rate = learning_rate\n \n # Initialize weights and bias\n self.w = torch.normal(mean=0.0, std=sigma, size=(num_inputs, 1), requires_grad=True)\n self.b = torch.zeros(1, requires_grad=True)\n self.losses = []\n\n def forward(self, X):\n \"\"\"\n Compute the forward pass: y = Xw + b.\n\n Args:\n - X (torch.Tensor): Input tensor of shape (batch_size, num_inputs).\n \n Returns:\n - torch.Tensor: Predicted values of shape (batch_size, 1).\n \"\"\"\n return torch.matmul(X, self.w) + self.b\n\n def compute_loss(self, y_pred, y_true):\n \"\"\"\n Compute Mean Squared Error loss.\n\n Args:\n - y_pred (torch.Tensor): Predicted values.\n - y_true (torch.Tensor): True values.\n \n Returns:\n - torch.Tensor: Scalar loss value.\n \"\"\"\n return 0.5 * ((y_pred - y_true) ** 2).mean()\n\n def update_parameters(self):\n \"\"\"\n Update the model parameters using gradient descent.\n \"\"\"\n with torch.no_grad():\n self.w -= self.learning_rate * self.w.grad\n self.b -= self.learning_rate * self.b.grad\n\n # Manually zero the gradients\n self.w.grad.zero_()\n self.b.grad.zero_()\n\n def train_step(self, X, y, batch_size):\n \"\"\"\n Perform a single training step.\n\n Args:\n - X (torch.Tensor): Input data of shape (num_samples, num_inputs).\n - y (torch.Tensor): Target data of shape (num_samples, 1).\n - batch_size (int): Number of samples per batch.\n \n Returns:\n - float: Loss value for the batch.\n \"\"\"\n # Sample a random batch\n num_samples = X.shape[0]\n indices = torch.randint(0, num_samples, (batch_size,))\n X_batch = X[indices]\n y_batch = y[indices]\n\n # Forward pass\n y_pred = self.forward(X_batch)\n\n # Compute loss\n loss = self.compute_loss(y_pred, y_batch)\n\n # Backward pass\n loss.backward()\n\n # Update parameters\n self.update_parameters()\n\n # Return the loss value as a scalar\n return loss.item()\n\n def train(self, X, y, epochs, batch_size):\n \"\"\"\n Train the model over multiple epochs.\n\n Args:\n - X (torch.Tensor): Input data of shape (num_samples, num_inputs).\n - y (torch.Tensor): Target data of shape (num_samples, 1).\n - epochs (int): Number of training epochs.\n - batch_size (int): Number of samples per batch.\n \n Returns:\n - list: List of loss values for each epoch.\n \"\"\"\n \n\n for epoch in range(epochs):\n # Perform a training step and compute the average loss for the epoch\n loss = self.train_step(X, y, batch_size)\n self.losses.append(loss)\n\n #print(loss)\n \n print(\"Final Loss:\", loss)\n self.plot_training_results()\n\n #return losses\n\n\n def plot_training_results(self):\n plt.figure(figsize=(25, 8))\n \n # First subplot: Training loss \n plt.subplot(1, 3, 1)\n plt.plot(np.log(self.losses), label=\"Log Loss\", color=\"blue\")\n plt.axhline(np.log(self.losses[-1]), linestyle=\"--\", color=\"red\", label=\"Final Loss\")\n plt.xlabel('Epoch')\n plt.ylabel('Log Loss')\n plt.title('Training Loss Over Time')\n plt.legend()\n plt.grid(True)\n \n # Second subplot: weight comparison\n plt.subplot(1, 3, 2)\n learned_weights = self.w.detach().numpy().flatten()\n true_weights = true_w.numpy().flatten()\n \n # Create a scatter plot comparing true vs learned weights\n plt.scatter(true_weights, learned_weights, alpha=0.5, color='blue')\n \n # Add a diagonal line representing perfect prediction\n max_val = max(np.max(true_weights), np.max(learned_weights))\n min_val = min(np.min(true_weights), np.min(learned_weights))\n plt.plot([min_val, max_val], [min_val, max_val], 'r--', label='Perfect Match')\n \n \n plt.xlabel('True Weights')\n plt.ylabel('Learned Weights')\n plt.title('True vs Learned Weights')\n plt.legend()\n plt.grid(True)\n \n # Third subplot: Bias comparison \n plt.subplot(1, 3, 3)\n true_bias = float(true_b.numpy().flatten()[0])\n learned_bias = float(self.b.detach().numpy().flatten()[0])\n \n x = ['True Bias', 'Learned Bias']\n y = [true_bias, learned_bias]\n plt.bar(x, y, color=['blue', 'orange'], alpha=0.7)\n \n for i, val in enumerate(y):\n plt.text(i, val + 0.02, f'{val:.2f}', ha='center', fontsize=10)\n \n plt.ylabel('Bias Value')\n plt.title('True vs Learned Bias')\n plt.grid(axis='y')\n \n plt.tight_layout()\n plt.show()Model = LinearRegressionModel(num_inputs,learning_rate)\nModel.train(X,y,epochs,batch_size)Final Loss: 0.07868940383195877![\"Training](\"https://www.tommasovaccari.com/static/output-21-1-fb3fd6c6.webp\")
This example, while simple, highlights some key concepts: we can view a linear regression model as a neural network, and we built this model from scratch without relying on the high-level APIs provided by PyTorch. In this second part of our machine learning series, we’ve taken a step toward higher-level abstraction. In the first post, we explored how each component of a neural network is built under the hood. Here, we leveraged PyTorch’s tensor implementation but still created the model from scratch with everything we needed to train it on our data.
","date_published":"2025-01-25T00:00:00.000Z","tags":["NeuralNetworks","Python","Math"]},{"id":"https://www.tommasovaccari.com/blog/introduction-to-neural-networks","url":"https://www.tommasovaccari.com/blog/introduction-to-neural-networks","title":"Neural Networks: Introduction","summary":"Introduction to Neural Networks.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"In this post, I would like to introduce how neural networks are built under the hood. I will explore the essential components and mechanisms that enable neural networks to learn and make predictions. By the end of this article, you should have a clearer understanding of how these powerful models work under the hood.
\nBefore diving deeper into neural networks, let’s briefly recap the main goal and how it is achieved. Neural networks (NNs) are a class of machine learning models inspired by the structure and function of the human brain. The primary objective of a neural network is to learn a mapping from input data to desired output values by adjusting its parameters through a process of optimization.
\nThe central aim of a neural network is to learn how to produce accurate outputs when given specific inputs. This process involves training the network on a dataset, where each data point consists of an input and a corresponding target output. For instance, in a classification problem, the input might be an image, and the target output would be the label of the object in the image. Or maybe in a regression problem, we want to train our model to answer the question, how much?
\nDuring training, the network is provided with a set of input data and its corresponding target outputs. The network makes predictions based on the input data, which are then compared to the target outputs. This comparison is crucial for evaluating how well the network is performing.
\nTo quantify the difference between the network's predictions and the target outputs, we use a loss function (or cost function). The loss function measures the prediction error. Common loss functions include Mean Squared Error for regression tasks and Cross-Entropy Loss for classification tasks. A lower value of the loss function indicates better performance of the network.
\nThe goal of training is to minimize the loss function, which involves finding the optimal set of parameters (weights and biases) for the network. To achieve this, we use an optimization algorithm that adjusts the parameters to reduce the loss. Backpropagation is the method used to compute the gradients of the loss function with respect to the network's parameters. It involves:
\nOnce the gradients are computed, they are used by the optimization algorithm (often gradient descent or its variants) to update the network's parameters. The network parameters are adjusted in the direction that reduces the loss function. This process is iterated over multiple epochs (passes through the entire dataset) until the loss function converges to a minimum value or sufficiently small error.
\nRoughly speaking, this is the process that we want to follow to achieve our goal. With the foundational concepts established, we will now delve into the detailed steps required to build a neural network. Firstly, we need to create a class that represents data within our neural network. This class must track the origin of each data value, including the operations and values used to compute it, in order to facilitate accurate gradient calculation during backpropagation. We will call this class Value.
\nNow I will start to implement in code what we have seen in the foundational concepts about NN.
\nclass Value: \n \n def __init__(self, data, _children=(), _op=''):\n \n self.data = data\n self._previously = set(_children)\n self._operations = _op\n self.grad = 0.0\n self._backward = lambda: None\n\n def __repr__(self):\n return f\"Value(data={self.data})\"Here we initialize the value objects with its own attribute. Let's break down each attribute:
\nThe repr method is used to print the value object. So for example, if we want to generate a value object we can do as follows:
\na = Value(3)\nprint(a)\n>>> Value(data=3)class Value: \n ... # existing code\n\n def __add__(self, other):\n out = Value(self.data + other.data, (self, other), '+')\n \n def _backward():\n self.grad += 1.0 * out.grad\n other.grad += 1.0 * out.grad\n out._backward = _backward\n \n return outHere we have the code for the add method. Simply in this method, we create an out Value object featured by a new value and a set of children (the two Value objects that originated it) and then we append the operation that originated that new Value object. Then from the fundamentals of calculus, we know that the gradient of the two children due to this operation can be calculated as shown in the code. For much more detail watch gradient explained here. Let's watch how the add function behaves:
\na = Value(3)\nb = Value(4)\nc = a + b\nprint(c)\n>>> Value(data=7)\nprint(c._previously)\n>>> (Value(data=3), Value(data=4))\nprint(c._operations)\n>>> '+'Now let's look at the code for the mul method:
\nclass Value:\n ... # existing code\n\n def __mul__(self, other):\n out = Value(self.data * other.data, (self, other), '*')\n \n def _backward():\n self.grad += other.data * out.grad\n other.grad += self.data * out.grad\n out._backward = _backward\n \n return outIn the add method, we create a new Value object with a new value and a set of children, which are the two Value objects that originated this result. We also append the operation that produced this new Value object. The backward function is designed to compute the gradient according to calculus rules.
\nNote that we use the += operator when updating the .grad attribute. This is because we want to accumulate the gradient. For example, if the Value object is involved in multiple operations, the gradient should reflect all these operations, making the gradient cumulative. Generally, this is the process that we follow to create new methods for operations for this class. For the purpose of creating a neural network, we need also a function that can work for us as an activation function. In this case, we add the tanh function for this scope.
\nclass Value:\n ... # existing code\n\n def tanh(self):\n x = self.data\n t = (math.exp(2*x) - 1) / (math.exp(2*x) + 1)\n\n out = Value(t, (self,), 'tanh')\n \n def _backward():\n self.grad += (1 - t**2) * out.grad\n out._backward = _backward\n \n return outFrom calculus, we know that tanh is defined as:
\n\\tanh(x) = \\frac{e^{2x} - 1}{e^{2x} + 1}Then we can implement the backward function knowing that:
\n\\frac{d}{dx} \\tanh(x) = 1 - \\tanh^2(x)Now we implement the exponential method for being able to use tanh:
\nclass Value:\n ... # existing code\n\n def exp(self):\n x = self.data\n out = Value(math.exp(x), (self,), 'exp')\n \n def _backward():\n self.grad += out.data * out.grad \n out._backward = _backward\n \n return outWe have now implemented most of the basic operations for our Value object. The final step is to implement a backward method. This method will be responsible for calling the backward functions of each object in the proper order.
\nHere's the code for the function:
\nclass Value: \n ... # existing code\n\n def backward(self):\n \n topo = []\n visited = set()\n def build_topo(v):\n if v not in visited:\n visited.add(v)\n for child in v._previously:\n build_topo(child)\n topo.append(v)\n build_topo(self)\n \n self.grad = 1.0\n for node in reversed(topo):\n node._backward()Thanks to the implementation of the Value object, we are able to construct a computation graph that captures the entire history of how a Value object was derived. This graph is formed by following the _previously set of each Value object, which tracks all the predecessor values involved in its computation.
\nTo compute the gradients, we begin by focusing on the final Value object, which typically results from a loss function in a neural network. Since we are interested in the gradient of this final Value with respect to itself, we initialize its gradient (self.grad) to 1. This initialization signifies that the gradient of the final Value with respect to itself is 1.
\nFollowing this, we execute the backward pass by invoking the _backward function for each Value object. This process starts with the final Value and proceeds backward through the computation graph to the initial Value objects. This reverse traversal ensures that the gradient for each Value is computed correctly based on the chain rule of differentiation, allowing us to accumulate the gradients appropriately.
\nWe've now looked at the basics of how a neural network works. This simple approach helps us understand the core concepts behind its operation. Now it's time to create a very simple net.
\nNow we will look more deeply into how to create a net, step by step.
\nNow we can create an artificial neuron with two inputs. This is an example of how a neuron works:
\n# Inputs x1, x2\nx1 = Value(1.0)\nx2 = Value(2.0)\n# Weights w1, w2\nw1 = Value(6.0)\nw2 = Value(-4.0)\n# Bias of the neuron\nb = Value(9.8)\n# x1*w1 + x2*w2 + b\nx1w1 = x1 * w1\nx2w2 = x2 * w2\nx1w1x2w2 = x1w1 + x2w2\nn = x1w1x2w2 + b\n# Activation function\no = n.tanh()o is the output of our neuron. The next step is to define the class Neuron, here's the code:
\nclass Neuron:\n \n def __init__(self, nin):\n self.w = list()\n self.b = Value(random.uniform(-1, 1))\n\n for _ in range(nin):\n self.w.append(Value(random.uniform(-1, 1)))\n \n def __call__(self, x):\n # w * x + b\n activation = 0\n for wi, xi in zip(self.w, x):\n activation += wi * xi\n activation += self.b\n out = activation.tanh()\n\n return out\n \n def parameters(self): # For storing the parameters\n return self.w + [self.b]Here we define the class Neuron. In the init method, we create the weights based on the number of inputs (nin) and then we create the bias Value. When we call the Neuron, it performs the activation based on the input and produces the result. Here's an example usage:
\nx = [1, 2, 3]\na = Neuron(3)\nprint(a(x))\n>>> Value(data=something)The next step is to create a layer made of neurons, here's the code:
\nclass Layer:\n \n def __init__(self, nin, nout):\n self.neurons = list()\n\n for _ in range(nout):\n self.neurons.append(Neuron(nin))\n \n def __call__(self, x):\n outs = list()\n \n for neuron in self.neurons:\n outs.append(neuron(x))\n\n return outs[0] if len(outs) == 1 else outs\n \n def parameters(self):\n return [p for neuron in self.neurons for p in neuron.parameters()]In the init method, we initialize a layer specifying the number of inputs for each neuron (nin) and the number of outputs for the layer (nout). When we call it, the call method pushes the input to every neuron and returns the output of every neuron.
\nThe final step is to create an MLP fully connected, here's the code:
\nclass MLP:\n \n def __init__(self, nin, nouts):\n sz = [nin] + nouts\n self.layers = list()\n\n for i in range(len(nouts)):\n self.layers.append(Layer(sz[i], sz[i+1]))\n \n def __call__(self, x):\n for layer in self.layers:\n x = layer(x)\n return x\n \n def parameters(self):\n return [p for layer in self.layers for p in layer.parameters()]In the init method, as usual, we initialize the MLP by giving the number of inputs of the neurons (nin) and then a list (nouts) where we store the number of neurons for each layer. When we create the self.layer attribute, we iterate over a list that we create where we have the number of inputs and outputs required for every layer. Now let's look at how we can use it:
\nx = [1.0, 5.0, -6.0]\nnet = MLP(3, [4, 4, 1])\nnet(x)\n>>> Value(data=something)Now we are ready to train the net that we have built.
\nTo start, we can manually create a very easy dataset and the desired target:
\nxs = [\n [1.0, 2.0, -1.0],\n [6.0, -4.0, 1],\n [2, 1.0, -1.0],\n [4.0, 3.0, -2.0],\n]\nys = [2.0, -4.0, -3.0, 2.0]So when we feed our net the first list, we want to obtain as a result the first element of the ys list. Now all we need is a loss function that can tell us how good the output of our net is. When we have our loss function, we can call the backpropagation on it (remember that it is a Value object so we can do it) and then we can slightly adjust the parameters according to what minimizes the loss function. Now we can implement it:
\nfor k in range(20):\n \n # forward pass\n ypred = list()\n for x in xs:\n ypred.append(net(x))\n loss = 0\n\n for ygt, yout in zip(ys, ypred):\n loss += (yout - ygt)**2\n\n # backward pass\n for p in net.parameters():\n p.grad = 0.0\n loss.backward()\n\n # update\n for p in net.parameters():\n p.data += -0.1 * p.grad\n\n print(k, loss.data)In this piece of code, we train our net following these steps:
\nIn this first part, we built the fundamentals for understanding neural networks. We built our own API to gain a deeper understanding of how they work.\nIn the next part, we will implement a neural network for linear regression using a higher-level approach.
","date_published":"2024-08-03T00:00:00.000Z","tags":["NeuralNetworks","Python","Math"]},{"id":"https://www.tommasovaccari.com/blog/neural-networks-roadmap-and-sources","url":"https://www.tommasovaccari.com/blog/neural-networks-roadmap-and-sources","title":"Neural Networks: Roadmap and Sources","summary":"Roadmap for the Neural Networks series and main sources.","authors":[{"name":"Tommaso Vaccari"}],"content_html":"I have a strong interest in both math and coding, and neural networks sit right at the intersection of the two. They are incredibly powerful architectures that may appear almost magical at first glance — but underneath, what’s really happening is a lot of fascinating mathematics and elegant ideas.
\nSince I usually take notes while studying, I decided to turn them into written explanations. This way I can both strengthen my own understanding and retention, and hopefully spark curiosity in others who might be interested in this field.
\nIn this series, as I learn new concepts, I’ll write new blog posts that mix math and theoretical deep dives with practical coding examples. These posts are not meant to be a primary reference, but rather an accessible and engaging exploration. For this reason I am going to provide, down below, all the sources I have studied and consumed to deepen my knowledge on the topic.
\nThis series will grow step by step as I deepen my understanding of neural networks. Each post focuses on a specific building block. For now, the planned chapters are:
\nSince this series is meant as a personal exploration rather than a formal reference, I’ll list here the materials I study along the way. These sources are diverse — books, academic papers, online courses, blog posts, and videos — so that anyone interested can follow the same path and dig deeper.
\nI hope you'll find something useful in this series and I encourage you to keep exploring this field!
","date_published":"2024-08-01T00:00:00.000Z","tags":["NeuralNetworks","Math","Sources"]}]}