1 Introduction↩︎

Designing new molecules is a core task in drug discovery and materials science, and language models have become a major tool for molecular generation. From RNNs [1], [2] to transformer decoders [3]–[5], these generators have grown by orders of magnitude in parameter count. However, whether such scale is necessary and what computations underlie valid generation remain poorly understood.

One promising way to reduce parameter count without reducing computational depth is to reuse the same transformer block repeatedly. Weight-shared transformers instantiate this idea by applying a single block iteratively for K passes. This matches the computational depth of a K-layer unshared stack with one block’s worth of parameters. Such architectures are Turing-complete [6], can run iterative algorithms like gradient descent [7], and match or exceed standard transformers on algorithmic reasoning [8], language modeling [9], [10], and vision [11]. This makes them particularly suited to formal-language domains like SMILES, where validity depends on repeated local and long-range checks.

Weight sharing also makes generation interpretable by attributing behavior to computation rather than parameters. Prior work in mechanistic interpretability has decomposed circuits in standard transformers [12], [13] and in molecular transformers [14]. In these architectures, however, each layer has its own weights, so functional differences between layers can reflect differences in parameters rather than differences in computation.

Here, we show that a single 53K-parameter weight-shared transformer block, trained on ZINC-250K [15] SMILES strings (Simplified Molecular Input Line Entry System), generates novel SMILES with 95% validity. At 53K parameters, the model is over an order of magnitude smaller than the smallest prior transformer-based molecular generator we compare against [3]–[5], and outperforms an identically-trained 527K-parameter unshared GPT in validity (95.3% vs.%) using 10 times fewer parameters.

Beyond generating valid molecules, the block is mechanistically interpretable. With weight sharing, the same parameters run at every pass, so functional differences across passes reflect changes in computation, not changes in weights. SMILES validity depends on three hard constraints: brackets must match, ring digits must pair at long range, and atom valences must be chemically valid. Because each constraint has an exact symbolic check, we can measure how generation improves at every intermediate pass.

Using three analyses (error classification across passes, a combined probing/sparse-autoencoder study of internal representations, and ablation), we show that the model learns chemical grammar pass by pass: brackets are resolved by pass 2, rings by pass 4, and valence by pass 8 (Fig. 1). Sparse autoencoders and linear probing recover the same ordering in the internal representations, from individual tokens in early passes to bracket depth and ring state in middle passes (ring-state accuracy climbs from 79% at pass 1 to 98% by pass 5). Ablation localizes bracket matching to a single attention head in WS-206K, where ablating it at pass 1 raises bracket errors by 23 percentage points without changing ring or valence errors, while the smaller WS-53K distributes the same function across heads.

Together, these results show a shared block resolving constraints from local to global scope, suggesting that similar pass-wise structure may appear in other formal-language domains with hierarchical constraints (Section 4). The weight-shared block thus serves as both a compact generator for molecular design and a small testbed for iterative computation.

2 Experimental setup↩︎

This section defines the SMILES generation task, the weight-shared and baseline models, and the training and evaluation protocol.

2.0.0.1 SMILES.

SMILES is a text encoding of molecular graphs. For example, atoms can be written with symbols such as C, N, and O. Bonds can be implicit or written explicitly with symbols such as = and #. Branches use parentheses, and ring closures use paired digits. The string CC(=O)O, for instance, encodes acetic acid. A generated string is valid only if brackets balance, ring digits pair, and atom valences are chemically possible. Validity rate is the fraction of generated strings that decode to real molecules via RDKit. We use SMILES rather than SELFIES [16] because SMILES leaves these grammar and chemistry constraints for the model to learn, whereas SELFIES guarantees validity by construction.

2.0.0.2 Architecture.

All models are autoregressive transformers based on GPT [17]. We compare three variants. Standard GPT uses learned positional embeddings and Gaussian Error Linear Unit (GELU) activations with a feed-forward expansion factor of 4. Modern GPT replaces these with rotary position embeddings (RoPE; [18]) and Swish-Gated Linear Units (SwiGLU; [19]). Weight-Shared (WS) reuses a single Modern GPT block for \(D\) forward passes, equivalent to a looped transformer [6]. We call \(D\) its virtual depth. All transformer variants apply pre-norm LayerNorm (LayerNorm before each sublayer), tie input/output embeddings, use dropout 0.1, and have no bias terms. We additionally evaluate Gated Recurrent Unit (GRU) baselines at parameter counts matched to WS-53K and WS-206K. Model names combine class and parameter count, so WS-53K denotes a weight-shared block with 53K parameters and default virtual depth \(D{=}8\). GPT-527K and GPT-527K-mod denote unshared Standard and Modern GPTs, and GRU-53K denotes the GRU baseline. When a WS model uses a virtual depth other than the default \(D{=}8\), its name includes a suffix such as -D4 or -D16. WS-53K and WS-206K are the two weight-shared blocks that recur, with WS-206K depth variants for the virtual-depth sweep. Full configurations are in Appendix 8.

2.0.0.3 Training.

We train on ZINC-250K [15], the standard benchmark of 249,455 drug-like molecules for SMILES-based molecular generation [3], [4], with SMILES augmentation at \(n_\text{aug}{=}10\) [20] and a character-level tokenizer (vocab 50, max length 128). We optimize with AdamW for 100 epochs at batch size 256, using a cosine learning-rate schedule (peak \(5 \times 10^{-4}\), 200-step warmup), mixed-precision training (AMP), and gradient clipping at 1.0. Core experiments use 3 seeds (42, 43, 44); supporting ablations use 1 seed. The reported experiments used approximately 500 GPU-hours on AMD Instinct MI210/MI250 GPUs. Distillation and DPO settings are given in Appendix 14.

2.0.0.4 Evaluation.

We sample \(N{=}1{,}000\) molecules per model with temperature \(0.8\) and top-\(k{=}40\). We report RDKit validity, Fréchet ChemNet Distance (FCD; [21]), internal diversity (\(1\) minus mean pairwise Tanimoto similarity, a Jaccard-style overlap measure, computed on Morgan radius-2 1024-bit fingerprints), uniqueness, and novelty, defined as the fraction of generated molecules absent from the ZINC-250K training set. We categorize invalid SMILES by failure mode for the error classification analysis (Section 4.1). Further details are in Appendix 9.

3 Compact molecular generation↩︎

A 53K-parameter weight-shared transformer reaches 95.3% SMILES validity by sharing a single block across 8 passes, outperforming an unshared GPT with 10 times more parameters and approaching the validity of a GPT with 60 times more parameters.

3.1 Pareto frontier↩︎

We first ask whether weight sharing improves the tradeoff between parameter count and SMILES validity. We show that weight-shared models dominate the sub-megaparameter Pareto frontier. At every budget below 1M, the weight-shared curve delivers higher validity at fewer parameters (Fig. 2; full numbers in Appendix 8, Table 5). WS-53K outperforms GPT-527K by 7.7 points (95.3% vs.%) and matches its internal diversity (0.852 vs.), at 10 times fewer parameters. The generated molecules show little evidence of memorization: across 3 seeds, 0.13% of valid WS-53K samples are exact training copies, and the mean nearest-neighbor Tanimoto similarity to training molecules is 0.52, below the nearest-neighbor similarity within the training set (0.59). WS-53K also matches GPT-3.2M to within 5 points (95.3% vs.%) at 60 times fewer parameters. FCD, a distributional-similarity metric where lower is better, is modestly higher for weight-shared models (2.5–2.9 vs.–2.4), trading some distributional fidelity for parameter efficiency. Thus, weight sharing shifts the compact-generation frontier rather than merely matching smaller baselines.

3.2 Virtual depth↩︎

We vary how many times the shared block is applied, measuring whether validity improves over repeated passes. Validity rises sharply with virtual depth and plateaus at \(D{=}8\): 33% \(\rightarrow\) 91% \(\rightarrow\) 97% \(\rightarrow\) 99% \(\rightarrow\) 99% across \(D \in \{1, 2, 4, 8, 16\}\) (WS-206K, 3 seeds per cell). The \(D{=}1\) value (33.0 \(\pm\) 2.3%, 3 seeds, 1,000 samples) matches the separately-trained WS-206K-D1 row of Table 1 (33.0 \(\pm\) 2.3%). The 58-point rise from \(D{=}1\) to \(D{=}2\), followed by diminishing returns through \(D{=}8\) and a plateau at \(D{=}16\), suggests that repeated passes iteratively refine validity.

We next isolate whether the gain comes from iteration itself rather than from other block-level architectural choices. An unshared Standard GPT baseline with four attention heads reaches 91.6 \(\pm\) 0.3% validity. The matched Modern GPT adds RoPE and SwiGLU and reaches 94.4 \(\pm\) 0.2%. The WS model reuses that Modern block for 8 passes and reaches 98.0% validity, adding 3.6 points while using half the parameters of the unshared baselines (Appendix 8, Table 6), consistent with the Pareto dominance of looped architectures at larger scale [10], [22], [23].

To test whether iteration explains the gain, we compare against three alternatives. First, removing iteration from WS-206K collapses validity by 65 points. Second, larger unshared GPTs remain below WS-206K. GPT-241K reaches 88.4 \(\pm\) 0.6%, and GPT-228K-mod reaches 91.2 \(\pm\) 1.4% (3 seeds each), despite having more parameters. Third, extra training compute does not close the gap. A matched Modern GPT reaches 94.2 \(\pm\) 0.1% after 800 epochs (8 times longer, with roughly twice WS-206K’s training compute), similar to its 100-epoch value and still 3.8 points below WS-206K. Thus, the compact model succeeds by refining SMILES validity over repeated applications of the same small block.

3.3 Robustness checks↩︎

We test whether this compact-generation result survives changes in model class, SMILES augmentation, sampling settings, and post-training.

First, we compare against recurrent models at matched parameter count. GRU controls trail the weight-shared block, with GRU-53K reaching 87.7% and GRU-206K reaching 95.5% (2 seeds each, 25–30 training epochs vs. for the WS comparators; matched-epoch comparison left to future work). At matched parameter count, WS exceeds the GRU by 7.7 points at 53K and 2.5 points at 206K, narrowing with capacity in line with larger GRUs (\(>\)​1M) reaching high validity [1]. Thus, the GRU comparison supports the compact-generation result at matched parameter count, while showing that the advantage narrows as recurrent capacity increases.

Second, we reduce SMILES enumeration by training with one string representation per molecule instead of 10 randomized SMILES. Under this setting, WS-53K drops from 95.3% to 92.3%, a 3.0-point decrease, and WS-206K drops from 98.0% to 97.8%, a 0.2-point decrease within seed variance (3 seeds), so augmentation helps but less than architecture.

Third, we vary decoding temperature, top-\(k\), and top-\(p\). Across temperature (\(T \in \{0.7, 0.8, 1.0, 1.2\}\)), top-\(k\) (\(\{40, \infty\}\)), and top-\(p\) (\(\{0.95, \text{none}\}\)), WS-206K remains above the 403K unshared Modern GPT in every effective cell despite using half as many parameters (sign test \(p{<}10^{-3}\) over 12 effective cells, since at \(T{\le}0.8\) the top-\(k{=}40\) and top-\(k{=}\infty\) cells are bitwise-identical).

Finally, we test standard post-training methods. Offline knowledge distillation (KD; [24]) drops WS-206K validity by \({\sim}8\) percentage points (3 seeds); the KL on soft labels spreads probability mass onto invalid tokens (Appendix 14), and teacher quality has no effect (Appendix 15). Direct Preference Optimization [25] preserves validity to within 1 pp (WS-206K \(98.5{\rightarrow}97.6\%\); full results in Appendix 16), so the compact-generation result is robust to preference fine-tuning. Together, these robustness checks support the interpretation that compact validity comes from repeated computation through the shared block, not from the sampler, memorization, or post-training.

4 Computation across passes↩︎

We next ask how weight-shared models generate valid molecules across repeated passes. Reusing a block provides iterative computation, but it does not determine whether successive passes perform distinct computations. We find here a fixed hierarchy of increasing scope. The shared block resolves brackets by pass 2 (local pairings within a few tokens), rings by pass 4 (long-range digit pairing across tens of tokens), and valence by pass 8 (atom-level integration of full bond context).

The rest of this section establishes this hierarchy with four complementary measurements: outputs, representations, causal structure, and capacity. We first classify errors across virtual depth to test whether failures disappear in the bracket-to-ring-to-valence order (Section 4.1). We then use linear probes and sparse autoencoders to ask whether hidden states follow the same progression from token-level features to bracket depth, ring state, and atom identity (Section 4.2). We next ablate individual heads to localize bracket matching to a single attention head in WS-206K (Section 4.3). Finally, we compare model sizes to test whether smaller blocks resolve the same constraints slightly later and with less localized circuitry while preserving the order itself (Section 4.4).

4.1 Errors resolve in order↩︎

To test whether repeated passes specialize at the output level, we classify errors in invalid molecules generated at each virtual depth. This sweep isolates virtual depth by holding block size fixed and varying \(D\) in separately trained WS-206K-D\(N\) models. Later analyses hold \(D{=}8\) and compare WS-53K with WS-206K. Specialization should not only make invalid samples rarer, but it should also shift failure modes from simpler syntax errors toward harder chemical constraints. The three error classes instantiate this order: bracket errors test parenthesis matching, ring errors test long-range digit pairing, and valence errors test context-sensitive chemical consistency. We train a separate model at each depth (\(D \in \{1,2,4,8,16\}\)), which avoids reading early-pass representations through an output head calibrated for pass 8.

Table 1 shows the predicted shift in assigned failure modes. At 1 pass, many samples fail before reaching later checks. By pass 2, bracket errors have collapsed from 23.0% to 1.3%, consistent with early resolution of parenthesis matching. Among the remaining assigned failures, ring-closure errors then drop between passes 2 and 4 (1.9% to 0.8%), resolving long-range digit pairing. Valence violations are the last substantial residual class, dropping from 2.0% to 1.0% between passes 4 and 8. The main error sweep therefore supports the predicted bracket-to-ring-to-valence sequence.

The main sweep compares separately trained models, so the ordering could in principle reflect training a different model at each depth rather than computation inside one model. To control for this depth-training confound, we read the trained 8-pass model’s output head at each pass (Table 7). This intermediate readout gives lower absolute validity, as expected when a pass-8 head is applied to earlier representations, but the ordering is unchanged, with brackets the first error type to drop below 5% (pass 3), rings the second (pass 5), and valence the third (pass 6). Thus, the intermediate readout rules out training depth as the source of the ordering.

4.2 Representations organize across passes↩︎

The output-level analysis shows when each constraint is satisfied in the generated molecules, and we test whether the same hierarchy appears inside the model’s hidden states as the shared block is applied repeatedly. We use two complementary tests on the hidden states of the trained 8-pass model (pass \(k\) refers to the state after the \(k\)-th application of the block). First, linear probes ask whether bracket depth and ring state are decodable at each pass. Second, sparse autoencoders (SAEs) ask whether interpretable chemical features emerge without specifying a readout in advance. Table 2 summarizes both analyses; they show the same bracket-then-ring ordering.

4.2.0.1 Linear probing.

We train linear probes on each pass of the 8-pass model to test when bracket depth and ring state become decodable. In both models, probes decode bracket depth earlier than ring state. WS-53K reaches 98.4% bracket-depth accuracy at pass 3, before ring-state accuracy reaches 98.0% at pass 5. WS-206K reaches 99.2% bracket-depth accuracy at pass 3, before ring-state accuracy reaches 98.6% at pass 4. Although WS-53K starts 3–6 points lower than WS-206K at pass 1, both models reach comparable probe accuracy by pass 5, so the difference is timing rather than final decodability. The ordering is the same in a single trained model as in the output-level analysis.

4.2.0.2 Sparse autoencoders.

Sparse autoencoders test the same progression without training a linear probe for each property. For each pass, an SAE learns sparse features from hidden states. We then ask which features track five chemical properties: is_bracket, is_ring_digit, bracket_depth, ring_state, and is_atom (Appendix 12). The learned features separate by pass most clearly in WS-53K. Token-level detectors (is_bracket, is_ring_digit) appear early, at passes 2–3 with correlations \(r{=}0.97\) and \(0.99\). Atom identity also achieves high correlation early (pass 2 in WS-53K, \(r{=}0.84\)), consistent with atom-type information being available from early passes. Two features require integrating information across multiple tokens: bracket depth (the count of unmatched open brackets at each position) and ring state (which ring digits are currently open). Ring-state correlation reaches its peak at pass 5 (\(r{=}0.64\)). Bracket-depth correlation is bimodal: near-equal peaks at passes 2 and 8 (\(r{=}0.63\) and \(0.64\)), separated by a dip to \(r{\approx}0.49\). WS-206K shows weaker separation (Table 2; Fig. 5). Because the SAEs are not trained to predict these properties, the features appear as emergent decompositions rather than only as supervised readouts.

Moreover, feature–property correlations are stable across corpus size (Appendix 12). When we scale the SAE corpus from 2,000 to 20,000 ZINC-250K SMILES, WS-53K feature–property correlations strengthen or remain within 0.01. Taken together, error classification, supervised probes, and the SAE support the same brackets-to-rings-to-valence hierarchy.

4.3 The bracket-matching head↩︎

In both error classification and internal representations, bracket matching is the first constraint to resolve. We further test whether that first step can be traced to a specific component of the shared block. If so, ablating one attention head at one pass should selectively raise bracket errors, rather than simply lowering validity whenever any part of the block is removed. We ablate a head by zeroing its contribution at one pass, then measuring the full error profile. We repeat this intervention across all four heads and all eight passes, for a total of 32 head-pass combinations, and generate \(n{=}2{,}000\) molecules per condition for both models. Table 3 summarizes the resulting changes in bracket errors and validity.

Across the full sweep, we localize the bracket-matching computation to a single attention head in WS-206K. Ablating that head at pass 1 raises bracket errors by 23 points while ring and valence errors remain at baseline; the effect persists through passes 2–3 (increases of 17 and 5 points) and then vanishes by pass 4, and all other heads at every pass remain near baseline. The dependence is thus specific to one head rather than to early-pass computation in general. WS-53K does not concentrate bracket matching the same way: at seed 42, no single head dominates (the largest pass-1 bracket-error increase is 3 points), and the 10-point increase in the 3-seed mean is carried by different heads in different seeds (heads 3, 1, and 3 for seeds 42, 43, and 44). The smaller model therefore distributes bracket matching across heads, whereas the larger model localizes it to one.

We use two controls to test whether the effect is specific to one head rather than to early-pass ablation or activation magnitude in general. To rule out general early-pass sensitivity, we ablate all of pass 1 in WS-206K rather than the bracket head alone; this leaves bracket errors at baseline and drops validity by only 0.5 points, roughly 50 times smaller than the 23-point drop from the single-head ablation. To rule out magnitude-driven effects, we compare activation norms: in all three seeds, the bracket head has the smallest pass-1 norm, and ablating the largest-norm head does not produce a larger validity drop. Together, the controls show that bracket matching depends on this specific head.

4.4 Smaller models need more passes↩︎

To test whether the pass-by-pass hierarchy changes with capacity, we compare WS-53K and WS-206K across error classification and probing. We find that WS-53K follows the same sequence as WS-206K but reaches its resolution points about one pass later (Table 8): per-pass validity peaks at pass 8 rather than pass 7, and the ring-state probe peaks at pass 5 rather than pass 4. Bracket depth peaks at pass 3 in both models, so the shift is small and appears only in the longer-range constraints. The two models reach comparable final accuracies (95.3% vs.% validity, and 98.4% vs.% best bracket-probe accuracy). Capacity also changes how bracket matching is implemented: WS-206K localizes it to a single head (Section 4.3), whereas WS-53K distributes it across heads. Thus, lower capacity slightly delays when each constraint resolves and spreads bracket matching across more heads, but does not change the order in which the constraints resolve.

5 Discussion↩︎

This paper shows that weight sharing can make molecular generation both compact and interpretable. With only 53K parameters, a shared-block transformer generates highly valid SMILES while preserving a visible sequence of validity computation. This improves the parameter-validity tradeoff, showing that small models can recover much of the validity of substantially larger unshared transformers.

We observe the same sequence across generated errors, hidden-state probes, and ablations: bracket matching first, ring closure second, and valence last. This consistency suggests that the shared block does more than reduce parameter count. By reusing one set of weights across passes, the model exposes successive stages of validity computation.

The ordering also links the empirical result to the grammar of SMILES. Bracket matching, ring closure, and valence are not arbitrary metrics; they form a progression from local syntax to long-range structure to chemically constrained validity. This connects the observed sequence to formal-language results in which fixed-depth transformers struggle with context-free and higher languages, while iterative depth supplies the recurrence such constraints require [26], [27]. Weight sharing makes this connection visible because the same parameters are applied repeatedly, so cross-pass changes track computation through time rather than differences between learned layers.

5.0.0.1 Limitations.

We qualify the claims in three ways. First, the empirical scope is deliberately focused. We study ZINC-250K SMILES, so the ordering should be read as a claim about validity computation in this representation rather than about molecular structure independent of representation. Second, the mechanistic measurements are finite-sample estimates. The SAE analysis uses 20,000 SMILES (\({\sim}\)​770K tokens), so rare features may only emerge at larger scale, and the ablation sweep uses 2,000 molecules per condition, enough to identify the bracket-head effect but not to resolve very small effects from other heads. Third, weight sharing reduces parameters but not autoregressive inference cost, because FLOPs scale with pass count.

5.0.0.2 Broader impacts.

Compact molecular generators can lower the compute cost of exploratory molecule design and mechanism analysis. The same capability could be misused to propose molecules with undesirable biological activity, so generated molecules should not be used for synthesis prioritization without toxicity, safety, and feasibility screening.

5.0.0.3 Future directions.

These results point to broader tests of whether shared depth can expose structured computation in other scientific domains. One direction is to extend to other biological, chemical, and materials systems where similar hierarchies of constraints arise, such as RNA folding from local base pairing to long-range tertiary contacts. A second direction is to span spatial and temporal scales by letting each pass resolve a different scale, from atomic interactions to macroscale dynamics and mechanics, providing a single architecture for problems that today require stitching separate models together. A third is to design distillation objectives that preserve the iterative computation, for example by matching per-pass representations rather than only output distributions.

6 Related work↩︎

6.0.0.1 Molecular generation with language models.

SMILES generation has progressed from RNN-based models [1], [2] to transformer decoders such as MolGPT [3], GP-MoLFormer [4], and REINVENT 4 [28]. Recent scaling work studies five molecular representations from 1M to 650M parameters [5]. Its smallest model is 19 times larger than WS-53K. To our knowledge, no prior transformer work explores molecular generation below 1M parameters, the regime where reusing computation offers a direct alternative to scaling parameter count.

6.0.0.2 Weight sharing and looped transformers.

Prior work uses weight sharing to compress transformers and to turn depth into repeated computation. Compression models such as ALBERT [29], Subformer [10], MobileLLM [30], and Relaxed Recursive Transformers [9] apply shared blocks at scales 1,000 to 10,000 times larger, where shared models match unshared baselines with fewer parameters [22], [23]. DeLiCaTe [31] brings the same compression motive to SMILES property prediction. Other work uses repeated depth as computation, including the Universal Transformer [8], theoretical results on looped transformers [6], and recent looped models with iteration-wise specialization [11], [32]. We use the same architectural device for a different purpose. Because SMILES validity depends on repeated grammar checks, weight sharing makes those checks observable as computation through time.

Large looped language models make repeated computation central, but they leave individual iterations difficult to interpret. Models such as Huginn [33], [34] and Ouro [35] operate at billion-parameter scale, where different iterations can perform similar computations [36]. Natural language also lacks the exact symbolic validity checks available in SMILES. Our setting combines a sub-100K bottleneck with constraints that can be measured at every pass, so the iteration is easier to separate into steps. This complements work on thermodynamic phases in looped transformers [37], formal-language constraints and circuit complexity [38], and concurrent studies of per-step specialization in latent reasoning models [39]–[41]. Where those studies rely mainly on correlational traces, training stages, or trace decoding [42], SMILES lets us test each intermediate computation with symbolic checks and causal ablation.

6.0.0.3 Mechanistic interpretability.

Mechanistic interpretability provides tools for finding circuits and features, but standard layerwise analyses make temporal ordering hard to separate from layer identity. We combine sparse autoencoders [43], [44], causal ablation [45], and probing [46]. [14] identify branch-balancing and ring-closure circuits in a 25M-parameter molecular transformer, establishing that chemically meaningful SMILES circuits can be recovered with SAE-based analysis. Weight sharing lets us test how related constraints resolve across repeated applications of the same block, rather than only where they appear across layers with separate weights. This distinction matters because [47] find a correlational syntax-to-semantics arc across BERT layers, while subsequent work questions such arcs under controls for layer position [48]. Shared weights remove that layer-identity confound. We also use causal ablation because attention-based analysis of bracket matching can mislead [49], and we validate SAE features against symbolic chemical ground truth rather than auto-interpretability scores [50]. This per-pass analysis complements SMI-TED’s single-extraction-point SAE [51] and TIDE’s temporal SAEs for diffusion transformers [52].

7 Conclusion↩︎

SMolLM, a 53K-parameter weight-shared transformer, reaches 95.3% validity on ZINC-250K, outperforms an unshared GPT with 10 times more parameters, and shifts the sub-megaparameter Pareto frontier. The architecture is also mechanistically interpretable. Across generated errors, hidden-state probes, and ablations, the shared block resolves SMILES constraints in a fixed pass-by-pass hierarchy (brackets by pass 2, rings by pass 4, valence by pass 8), consistent from 53K to 206K parameters, with capacity changing how quickly each constraint resolves rather than the order itself. Taken together, weight sharing is not only a compression device for molecular generation but also a way to study iterative computation in domains with exact symbolic constraints.

Appendix↩︎

8 Architecture configurations↩︎

4pt

6pt

6pt

9 Training and evaluation details↩︎

Dataset. ZINC-250K [15]: 249,455 drug-like molecules. SMILES augmentation [20]: each canonical SMILES expanded to \(n_\text{aug}=10\) randomized SMILES via RDKit. Character-level tokenizer, vocabulary size 50, max sequence length 128.

Training. 100 epochs, AdamW (weight decay 0.1), batch size 256, cosine LR (peak \(5 \times 10^{-4}\), 200-step warmup), mixed-precision training (AMP), gradient clipping 1.0, no label smoothing. Distillation (Section 3.3): offline KD with softmax temperature \(\tau=2.0\) and hard/soft loss mixing coefficient \(\alpha=0.5\); teacher GPT-467K-mod (98.5%). DPO (Section 3.3): 5 epochs, lr \(=10^{-5}\), KL regularization strength \(\beta=0.1\), sample budget 10K molecules per base, pairing valid samples against invalid ones (so usable pairs are bounded by the invalid count). Full DPO results are in Appendix 16.

Evaluation. \(N=1{,}000\) molecules per model, temperature \(T=0.8\), top-\(k=40\). Metrics: validity (RDKit parsing), FCD (Fréchet ChemNet Distance, analogous to FID; [21]), internal diversity (\(1 -\) mean pairwise Tanimoto, Morgan fingerprints radius 2, 1024 bits), uniqueness, novelty.

Compute. Reported runs used single AMD Instinct MI210 or MI250 GPUs (64–128GB HBM class) on an internal cluster. Training time ranged from under 2 GPU-hours for the smallest 53K models to under 24 GPU-hours for the largest reported models; most sub-megaparameter runs completed in a few GPU-hours. The reported experiments used approximately 500 GPU-hours total, excluding preliminary and failed runs that were not used in the paper. Checkpoints and logs require less than 1GB per run.

Pass indexing. The paper uses 1-indexed pass numbers (pass 1 is the first block application; pass 8 is the last). The released code uses 0-indexed positions (pass=0..7); to map paper to code, subtract 1.

Reproducibility. Code is released at https://github.com/akhljndl/smollm (MIT). Each experiment maps to a numbered configuration in experiments/smiles_gen/configs.py and can be reproduced via python -m experiments.smiles_gen.run --exp <N>. Trained checkpoints (3 seeds per config, all 21 architectures) are on Hugging Face at https://huggingface.co/akhljndl/smollm under checkpoints/0504_paper/, and training runs are public on Weights & Biases (entity ajindal, project smollm, tag 0504_paper). Two stability notes: WS-206K’s 3-seed std is 0.33 percentage points, and run-to-run variance across independent training runs is consistent with this. The four Table 6 checkpoints were regenerated and verified within 2 percentage points of reported values at \(N{=}10{,}000\) samples.

10 Error classification across passes↩︎

Companion plot to Table 1: error-type composition across virtual depths in separately-trained WS-206K-D\(N\) models. With \(n{=}1{,}000\) samples per depth averaged over 3 seeds, the cross-pass trends are robust.

11 Head-ablation heatmap↩︎

Companion to Table 3: full head \(\times\) pass heatmap for both models. We apply a Bonferroni correction across the \(K{=}32\) head-pass ablations for each model.

12 Representation organization across passes↩︎

Panel (a) shows linear-probe accuracy for bracket depth and ring state across passes. Panel (b) shows peak SAE feature correlation for five chemical properties. Companion to Table 2.

For each pass, we train an SAE with 512 features (4 times the hidden dimension for WS-206K and 8 times for WS-53K) using a sparsity penalty on 20,000 ZINC-250K SMILES (\({\sim}\)​770K tokens). Decoder dictionary atoms are constrained to unit \(\ell_2\) norm following [43]; the constraint reduces over the input dimension of the decoder weight so that each atom (column) has unit norm, which is necessary for cross-pass comparison of feature–property correlations under L1 sparsity. Scaling from 2,000 to 20,000 ZINC-250K SMILES increases or preserves the WS-53K feature–property correlations across all properties. The exact peak pass for an individual feature varies by \(\pm 2\) passes across seeds. For the late-pass convergence analysis, we compute cosine similarity between consecutive hidden states from the same trained 8-pass WS-206K model; pass 7–8 similarity reaches 0.995 while generated validity remains 97–98%.

13 Capacity comparison across methods↩︎

The diagnostics agree that WS-53K solves the same subtasks as WS-206K, with peaks appearing at the same pass or slightly later (Table 8). Five of the six diagnostic rows shift by 0 to \(+\)​1 pass, with no negative shifts; the exception is the WS-53K SAE bracket-depth profile, which is bimodal (near-equal peaks at passes 2 and 8) rather than single-peaked, spanning a \(-2\) to \(+4\) pass range relative to WS-206K. The overall picture is rate, not target: WS-53K reaches the same milestones on a slightly delayed schedule.

14 Distillation ablation↩︎

An earlier version of this sweep used a label-smoothed teacher; smoothing made the soft-target signal nearly uniform and looked like evidence that KD helps. We re-ran the full sweep with a non-smoothed teacher (GPT-467K-mod; \(\tau=2.0\)), and the apparent benefit disappeared: every mixture of soft targets underperforms hard-only (Table 9). The curriculum and curriculum+KD rows disambiguate the two effects: the scheduled data order alone gets 96.7% validity, and adding soft targets on top drops it to 90.3%, so the residual damage comes from soft targets themselves rather than from the training schedule.

15 Teacher quality↩︎

Teacher validity spans \({\sim}1\) percentage point across the three candidates (GPT-3.2M 99.4%, GPT-467K-mod 98.5%, WS-206K self-distill 98.5%). Distillation from the two external teachers yields student validities within 0.5 points (89.9% and 90.2%) despite one teacher having 7 times more parameters than the other, while self-distillation reaches 90.4% (Table 10). All three students remain \({\sim}8\) points below the no-KD WS-206K baseline of 98.0%, so teacher quality is not the binding constraint on the KD degradation. Results are single-seed.

16 Direct Preference Optimization↩︎

We apply DPO to two weight-shared bases (WS-206K, WS-82K) and one curriculum-distilled WS-206K variant, using the settings in Appendix 8. Across all three bases, validity changes by less than 2 percentage points and FCD shifts by at most 0.16 (Table 11). All results are single-seed.

References↩︎