PaperlayerView exact revision on arXiv

Source-linked AI summary

Dense Structural Compression of Transformers via Gauge-Correct Channel Removal

Jed A. Duersch, Naïm Es-Sebbani, Nathanaël Haas, Zied Bouraoui

arXiv:2609.07264v1cs.LGstat.ML

TL;DR

Transformer inference energy is dominated by dense matrix products, while uniform structures and post-hoc methods limit access to smaller, task-suited architectures. GaugeLasso applies calibrated structural pressure during training; sustained pressure discovers representations post-hoc pruning cannot extract, with retraining success depending on task type.

  • Problem

    Dense matrix products dominate transformer inference energy, and uniform layer structures do not reflect depth-dependent roles or the smaller structures that may suffice for inference.

    Source: §1 · pp. 1–2

  • Method

    GaugeLasso uses additive symmetric group-lasso penalties and efficiency-scaled channel calibration to suppress underperforming channels during training while retaining dense tensors.

    Source: §1 · p. 2, §6 · p. 10

  • Results

    During-training pressure discovers representations post-hoc methods cannot extract; fresh-init retraining recovers comparable quality on statistical tasks but fails on exact algorithmic execution.

    Source: §6 · p. 10

  • Takeaways & Limitations

    Discovered structural profiles expose depth-dependent capacity and inform the design of subsequent architectures.

    Source: §7 · p. 10

    Experiments are limited to architectures below 150M FMA trained from scratch, and fixed attention-head counts and unsupported channel regrowth constrain the method.

    Source: §6 · p. 10

Abstract

from arXiv · show

Inference energy per token drives the cost and carbon footprint of deployed transformers. It is dominated by dense matrix products that incur fused multiply-accumulate (FMA) operations and memory traffic. To reduce these computations while retaining dense tensors for high GPU throughput, we develop a methodology from first principles to adapt structural complexity during training to maximize inference utility per unit compute. Channel penalties drive entire tensor slices to zero to enable physical removal while preserving density and the network function. The natural approach, penalizing the norm of operator components acting through each channel, is provably destabilized by gauge freedom. We resolve this pathology with GaugeLasso: additive symmetric group-lasso penalties that recover a monotone function of product-norms when the network converges to gauge balance. Our equilibrium analysis enables per-channel calibration to correctly suppress slices that under-perform in inference utility per unit compute. Under adaptive pressure, the network reorganizes into depth-dependent structural profiles that can be far smaller than the architecture required to learn the task. On polynomial long division over $\mathbb{F}_{31}$, compute compresses from 148 to 255 times with perfect accuracy. On character-level language modeling, compressed models outperform the hand-designed baseline at equal FMA. On masked autoencoding, a compression trial exposes which axes were over-provisioned and which saturated, guiding a better second design. Compaction also accelerates training monotonically as the model progresses. Post-hoc pruning with the same utility ranking cannot reach these structures, showing that sustained pressure is central to discovery of efficient models. Retraining a discovered architecture recovers baseline quality on our statistical tasks, but fails on our exact algorithmic task.

1 Introduction

GaugeLasso uses calibrated channel penalties to discover smaller dense transformer architectures during training. Sustained pressure finds structures post-hoc pruning cannot reach, with task-dependent retraining outcomes.

1 Introduction

  • Dense matrix products dominate transformer inference energy, so reducing arithmetic and memory traffic matters while preserving efficient dense GPU access patterns.Attention and MLP projections use one FMA per parameter per token, with additional attention compute growing with context length.Source: §1 · p. 1
  • Uniform layer dimensions leave depth-dependent roles unmodeled, and the structure needed to discover a solution can exceed what inference requires.The large space of heterogeneous architectures makes enumeration impractical.Source: §1 · pp. 1–2
  • On polynomial long division over F31, GaugeLasso achieves up to 255× compression (median 187×) with perfect accuracy; Figure 1 traces depth-specialized structural change.The figure shows median trends across three of seven adaptive structural axes.Source: §1 · pp. 2–3, Figure 1 · p. 3
  • Additive symmetric group-lasso penalties address gauge instability, while efficiency-scaled calibration ranks channels by inference utility per FMA.Adaptive pressure produces dense tensors through physical channel removal; step time decreases as dimensions are removed.Source: §1 · p. 2
  • Post-hoc pruning cannot compress the F31 model and yields 22% worse perplexity at equal FMA on character-level language modeling; fresh-init retraining works on statistical tasks, not exact execution.A matched CoFi comparison reaches 1.97× compression where 7.09× was requested.Source: §1 · p. 2

2 Related Work

Prior efficiency methods prune fixed representations, search predefined architectures, or use heuristic coupling management. This work positions continuous structural discovery as a distinct approach to finding minimal-compute structures for a loss target.

2 Related Work

3 Structural Regularization

Structural regularization suppresses channels according to calibrated inference utility per compute. The derivation handles transformer-specific axes, gauge-connected factors, and efficiency calibration.

3.1 Transformer Setting

  • The method targets LLaMA-style transformers with SwiGLU MLPs and RMS normalization, using a prescaled variant to avoid discontinuous rescaling when channels are suppressed.Seven per-block dimensions and two global dimensions are independently adjustable.Source: §3.1 · p. 4

3.2 Group Lasso and Gauge Balance

  • For disjoint lasso groups, a channel’s utility equals its penalty at a critical point; overlapping groups require a correction term.A structural slice collects the parameters removed when dropping that channel within a factor.Source: §3.2 · pp. 4–5
  • Additive penalties on both sides of a structural axis stabilize gauge-connected factors; at balance, the group norms recover a gauge-invariant product-norm ranking.For two factors, the penalty minimum equalizes their weighted norms, and the balanced group norm gives the channel utility.Source: §3.2 · p. 5

3.3 Efficiency-Scaled Calibration

  • Efficiency-scaled calibration sets penalties using each channel’s FMA cost and slice element counts, yielding a uniform utility-per-FMA criterion across axes at gauge balance.The resulting utility is a first-order estimate because training only approaches a critical point.Source: §3.3 · p. 5

4 Training Procedure

Training combines optimizer-coupled penalties, adaptive penalty strength, and physical channel removal to compress models into smaller dense tensors. Gauge rebalancing and a conservative loss-target solver help govern when channels are dropped.

Coupled regularization

  • Coupled regularization combines task and penalty gradients in one optimizer update so both experience the same optimizer scaling.The framework tracks momentum on the task gradient alone so ρ and the penalty gradient can respond immediately to structural changes.Source: §4 · p. 6

Adaptive ρ

  • Adaptive ρ targets loss L*: the conservative solver applies compression only after loss falls below the target and strengthens pressure as the margin grows.It sets ρ to zero when the denominator is not sufficiently negative; an aggressive variant adds predicted optimizer gains as descent credit.Source: §4 · p. 6

Trust-drop

  • Trust-drop removes channels when their gauge-invariant distance to zero meets the threshold, deleting corresponding tensor slices and entire operators when an axis empties.The threshold uses trust multiplier k; at k = 1, only channels within one optimizer step of zero are dropped.Source: §4 · p. 6

Gauge Sinkhorn

  • Gauge Sinkhorn periodically rebalances gauge-connected pairs before drop decisions; reciprocal rescaling preserves the network function exactly for bilinear couplings.Alternating over axes converges in 2–3 iterations.Source: §4 · p. 6

Training phases

  • Training uses warmup without compression, an adaptive-ρ compression phase, and optional cooldown with the architecture frozen; physical drops shrink dense tensors and compatible momentum state.The procedure requires no optimization restart after structural modifications.Source: §4 · p. 6

Training Procedure

  • Each training step applies efficiency-scaled penalty gradients, solves for ρ, and periodically rebalances and drops eligible channels, cascading away empty operators.The procedure returns the converged architecture while preserving dense tensors.Source: §4 · pp. 6–7

5 Experiments

The experiments test dense compression across polynomial division, language modeling, masked autoencoding, and retrieval, assessing task quality, ablations, reachability, and training speed. They compare compressed architectures with baselines and controls to examine when compression succeeds and which structures it discovers.

5 Experiments

  • Experiments span four tasks and ask whether compression preserves performance, matches hand-designed baselines, depends on particular components, requires sustained pressure, and speeds training.The tasks are F31 polynomial division, PG-19 character-level language modeling, CelebA masked autoencoding, and synthetic retrieval.Source: §5 · p. 7

5.1 Polynomial Long Division over FF31

  • About 70k FMA yields 187× compression (up to 255×) with perfect exact-match accuracy on F31 polynomial division.At the tighter loss target, 18/25 seeds converge versus 3/5 at the looser target, with twice the training budget.Source: §5.1 · p. 7

5.2 Language Modeling and Vision

  • At matched FMA, character-level language modeling reaches perplexity 2.61 versus the baseline’s 2.65; a 24-layer model compresses 2.2× with 2.34 ppl.In the deeper model, early layer dimensions narrow while deeper layers retain capacity.Source: §5.2 · pp. 7–8, Table 2 · p. 8
  • On CelebA, the dk=256 design reaches loss 0.213 below baseline, while D and df remain near their ceilings, identifying axes for expansion.Under the stringent L*=0.32 target, dk and dv saturate in every layer.Source: Table 2 · p. 8, §5.2 · p. 8

5.3 Retrieval and Ablations

  • The retrieval ablations find exchange most impactful: removing it costs 37% more FMA, while gauge Sinkhorn accelerates compression and the sign-step ρ-solver over-compresses.The baseline compresses to one attention block and one MLP with D=26–38; k=1 is the safest tested trust-multiplier default.Source: Table 3 · p. 8, §5.3 · p. 8

5.4 Reaching the Discovered Architectures

  • Dense compaction produces 1.6–5× training speedups as step time falls monotonically; on one comparison, 7.1× FMA compression gives 2.4× lower wall-clock time.The FMA-matched uniform baseline is faster in that comparison, at 6.9s per 100 steps versus 11.9s for the compressed architecture.Source: Figure 2 · p. 9, §5.4 · p. 9
  • Post-hoc pruning cannot compress FF31 and gives 22% worse character-LM perplexity; CoFi reaches 1.97× compression against a 7.09× target.The CoFi result may reflect its mask space, which cannot express narrower key and value dimensions within surviving heads.Source: §5.4 · pp. 8–9, Table 4 · p. 9
  • Fresh-initialization training recovers comparable quality on retrieval and character-level language modeling, but fails on FF31, where four of five seeds finish below 0.01 exact match.The fifth FF31 seed reaches 0.805, compared with 1.000 for the compressed model.Source: Table 4 · p. 9, §5.4 · p. 9

6 Discussion

The discussion presents GaugeLasso as a route to dense compression and describes the structural profiles and representations training discovers. It also frames those profiles as design guidance and states limits on the experiments and supported compression axes.

6 Discussion

  • GaugeLasso uses additive group-lasso penalties to resolve gauge pathology, recover product-norm ranking at equilibrium, and calibrate channels by inference utility per FMA.Adaptive penalty strength tracks the loss target while preserving dense tensors for optimized GPU kernels.Source: §6 · p. 10
  • Under a fixed loss target, training converges to a near-minimal-compute structural profile; across retrieval seeds, expensive axes vary less in total cost than their dimensions.All ten retrieval runs retain one attention block and one MLP, while residual width spans 23–39 channels and the cheaper MLP interior spans 1–20.Source: §6 · p. 10
  • Post-hoc pruning confirms that during-training compression discovers reorganized representations, not subsets of unconstrained solutions; fresh initialization recovers comparable quality on statistical tasks but fails on exact execution.The post-hoc metric fails to compress FF31 or yields substantially worse character-level language-model quality.Source: §6 · p. 10
  • Structural profiles identify over-provisioned axes and saturated budgets, guiding later designs; on CelebA, reallocating capacity improves loss from 0.237 to 0.213.Axes that compress substantially exceed task needs, while those at their ceiling are binding.Source: §6 · p. 10
  • The experiments are below 150M FMA and keep attention-head count fixed; head penalties are future work, and narrower rather than shorter models face launch and memory latency limits.At these scales, FMA reduction overstates realized wall-clock gain; larger pretrained models remain an extension target.Source: §6 · p. 10

7 Conclusion

Sustained structural pressure during training discovers efficient, depth-specialized architectures that post-hoc methods cannot extract from converged models. Compaction accelerates training, and discovered architectures retrain to comparable quality on statistical tasks, though not on the exact algorithmic task.

7 Conclusion

  • Sustained structural pressure discovers representations more efficient than post-hoc methods can extract from converged models.The resulting architectures exhibit depth-dependent specialization that fixed-capacity training cannot discover.Source: §7 · p. 10
  • Dense compaction accelerates training, and each trial returns a structural profile that informs the next architecture.The discovered architectures retrain to comparable quality from fresh initialization on statistical tasks, but this does not hold for the exact algorithmic task.Source: §7 · p. 10
  • Dynamically calibrated, gauge-invariant channel utilities drive dense compression without disrupting optimization.The conclusion frames inference energy per token as an economic and environmental priority for deployed AI.Source: §7 · p. 10

A Design Pathologies and Their Resolutions

The appendix traces design choices that address unstable product-norm penalties, overlapping structural slices, regrowth constraints, and resource calibration. Its remedies use additive symmetric penalties, corrected utilities, constrained function-preserving growth, and calibration that can target different per-channel costs.

A Design Pathologies and Their Resolutions

  • Additive symmetric penalties keep pressure on both factors and make gauge balance a stable minimum, unlike product-norm penalties whose imbalance can leave dangling weights.Product-norm penalties can drive one factor to zero while the other retains mass, causing the penalty gradient on the surviving factor to vanish.Source: A Design Pathologies and Their… · p. 13
  • Coupled regularization restores the equilibrium identity under step-normalized optimizers, while training-time dense slice removal lets representations reorganize into fewer dimensions.Sparse pruning leaves tensors full-sized; dense slice removal produces physically smaller tensors at full hardware utilization.Source: A Design Pathologies and Their… · pp. 13–14

B.1 Corrected Utility for Overlapping Dimensions

  • When lasso groups overlap across axes, cross-terms make single-dimension group-lasso utility an underestimate of the true removal utility.The corrected utility includes the standard group-lasso self-term plus non-negative cross-terms.Source: §B.1 · p. 14

B.2 Regrowth Constraints

  • Function-preserving growth depends on the intervening transformation: linear axes allow full basis rotation, whereas nonlinear axes allow only channel permutation.The conservative solver and trust-drop criterion prevent premature removal, so regrowth is not needed in practice.Source: §B.2 · p. 14, Table 5 · p. 15

B.3 Cross-Count Calibration

  • At gauge balance, a shared rescaling can preserve RMS equalization while normalizing channel utility against a chosen resource cost.The resource can be FMAs, memory footprint, latency, or another per-channel cost.Source: §B.3 · p. 15

C Architecture Details

The architecture appendix inventories the structural axes and the resources released by dropping their channels, then describes normalization and growth mechanisms. Experimental details include the FF31 task setup and the structural profiles found under compression.

C.1 Axis Inventory

  • The inventory covers six structural axes per block plus global D, with four independently calibrated bridge paths in each block.Attention uses dk and dv; the SwiGLU MLP uses df, while residual-stream axes include D and the input/output bridge dimensions.Source: Table 6 · p. 15, §C.1 · p. 16

C.2 Freed Resources

  • The freed-resource accounting reports per-token multiply-add savings for each dropped index, with sequence-length terms requiring a target length to evaluate.The accounting distinguishes penalized gauge-connected factors from coupled tensors that are freed but do not participate in the penalty.Source: §C.2 · p. 16, Table 7 · p. 16

C.3 Prescaled RMS Normalization

  • Prescaled RMS normalization attaches penalties, mediates residual-stream gauge, and preserves surviving-channel normalization as a channel is removed.Gauge-Sinkhorn invariance is exact on bilinear couplings but only first-order elsewhere; the approximation was not observed to matter in practice.Source: §C.3 · p. 16

C.4 Procrustes Growth

  • Procrustes QR growth uses orthogonal gauge freedom to add channels on linear axes while minimally perturbing existing channels.Weight and momentum are transformed together to preserve the optimizer’s gradient basis.Source: §C.4 · p. 17

D Experimental Details

  • The experiments used single NVIDIA H100 GPUs, with approximately 500 GPU-hours across reported experiments; structural-profile panels encode axis extents by depth and training step.The profile figures use line color for training step and show seed ranges at convergence unless a caption specifies otherwise.Source: D Experimental Details · p. 17

D.1 FF31 Polynomial Division

  • FF31 requires three rounds of polynomial long division over Z/31Z, and converged profiles retain at most three attention blocks with a recurring MLP-block-2 dip.The unpenalized model reaches 100% token accuracy and exact match by step 4k; Figure 3 reports seed-level compression profiles.Source: §D.1 · p. 17

D.2 Character-Level Language Modeling

Character-level language modeling compresses wider adaptive models toward the baseline FMA budget, while the deeper model narrows channels across depth and retains most sublayers. The CelebA trial shows that compression targets reveal which axes saturate and which retain complexity.

D.2 Character-Level Language Modeling

  • 51.9M FMA (2.2×) with 2.335 ± 0.025 perplexity was reached by the 24-layer character-level model across five seeds.The compression phase was limited to 70k steps, and further compression is expected with a longer budget.Source: §D.2 · p. 18
  • In the 24-layer model, the first attention and MLP blocks are consistently narrowed, while whole-sublayer removal is sparse across seeds.The first attention block retains dk of 6–17 out of 45 and dv of 1–4 out of 90; the first MLP retains df of 68–191 out of 1920.Source: §D.2 · p. 18, Table 9 · p. 19, Figure 4 · p. 20

D.3 CelebA Masked Autoencoder

  • The tighter CelebA target retains structural complexity that the looser target drops, while early attention is suppressed and MLP complexity stays fairly constant.In the dk=256 run, dk and dv show an ascending trend with depth; Figure 3 reports seed-level FF31 structural profiles, including 18 of 25 seeds converging to about 70k FMA.Source: Figure 3 · p. 19, §D.3 · p. 21

D.4 Retrieval Task, Ablations, and Structural Profiles

The retrieval-task ablations vary training mechanisms and report their structural profiles; post-hoc pruning uses a three-phase protocol, while the full-population results distinguish FF31 seed outcomes. The CelebA profiles also show how loss targets and capacity choices affect compression and loss.

D.5 Post-Hoc Pruning Protocol

  • Post-hoc pruning follows a three-phase pipeline using the same gauge-calibrated utility metric as during-training compression.The FF31 protocol iteratively removes the lowest-efficiency channel and stops when holdout loss exceeds L∗=0.1; the character-level LM is pruned toward a 7.6M FMA target over 10 rounds.Source: §D.5 · p. 21
  • With dk=dv=256, the CelebA model reaches loss 0.213, below the baseline loss of 0.236.The figure reports aggressive compression with higher loss at the loose target and recovery of baseline quality with all 12 layers surviving at the tighter target.Source: Figure 5 · p. 22

D.4 Retrieval Task, Ablations, and Structural Profiles

  • The retrieval ablations show that without gauge Sinkhorn the optimizer alone cannot expose channel importance quickly enough, while exchange is the most impactful mechanism.The baseline compresses D from 128 to ∼30; the sign-step solver yields loss 5.6 and accuracy 25%, while raising the trust multiplier has minimal effect on the attractor.Source: §D.4 · p. 21, Table 10 · p. 21, Figure 6 · p. 23

E.1 Full-Population Results

  • FF31 full-population reporting distinguishes seeds that never compressed from seeds that compressed but fell below the exact-match criterion.Three seeds ended at the initial 13.2M FMA; the other non-converged seeds compressed into the 52k–95k band. Per-seed values for other reported conditions are listed without exclusions.Source: §E.1 · p. 25

E.2 Gauge-Component Ablation

The retrieval ablation tests whether gauge-correct penalties, rather than structured pressure alone, drive compression while preserving task performance. The symmetric penalty is the only tested variant that compresses and keeps the task; the alternatives are evaded or over-compress.

E.2 Gauge-Component Ablation

  • Gauge Sinkhorn moves one-sided compression from 442k to 302k FMA, but accuracy falls from 0.996 to 0.368.The two one-sided arms differ in the Sinkhorn flag and share seeds, schedule, and other settings; rebalancing works against one-sided pressure.Source: §E.2 · pp. 25–26
  • The product-norm arm compresses to 150k FMA, 0.6× the control, but has accuracy 0.579 ± 0.308.The passage attributes the failure to an unbounded product-norm gradient as the other factor shrinks and reports a bimodal spread across seeds.Source: §E.2 · p. 26, Table 13 · p. 27
  • The symmetric penalty is the only tested variant that compresses and keeps the task; one-sided pressure is evaded, while product-norm pressure over-compresses.The one-sided arm reaches 0.996 accuracy without sufficient compression, whereas the product-norm arm reaches 150k FMA and accuracy 0.579 ± 0.308.Source: §E.2 · pp. 25–27, Table 13 · p. 27
  • The variant arms vary data stream with seed, while symmetric reference runs vary only initialization; the one-sided-with-Sinkhorn arm has no exactly matched control.The one-sided pair it is contrasted against is exact, and the table reports step 25k because the generations share a schedule only through that point.Source: §E.2 · p. 27

E.3 An External During-Training Method

The CoFi comparison falls short of requested compression despite maintaining quality near the teacher, while its mask families leave head width untouched. FF31 reliability also depends on reaching the loss target: lowering the learning rate advances compression onset, but does not establish a reliability gain.

E.3 An External During-Training Method

  • CoFi reaches 2.05× against 8.38× requested for the parameter target, while both arms finish within 0.1 perplexity of the teacher.The two arms use the same teacher and schedule and differ only in the measured resource objective.Source: §E.3 · p. 27, Table 14 · p. 27
  • CoFi's coarse mask families give up nothing, and head width remains untouched in all twelve sublayers.The realized compression comes from the two fine mask families; the authors bound this comparison to the published mask families, which cannot express narrowing key and value dimensions inside surviving heads.Source: §E.3 · p. 27, Table 15 · p. 28

E.4 Reliability and Schedule Sensitivity

  • At η=10^-3, 2 seeds fail to start compression and 5 fail to hold; at η=5×10^-4, all start and 5 still fail to hold.A seed that never reaches L∗ cannot compress, while another may begin compressing but fail to maintain the loss target or meet exact match.Source: §E.4 · p. 28, Table 16 · p. 28
  • Halving the learning rate lowers median compression onset from 2400 to 1700 steps, but the convergence difference is unresolved and is not read as a reliability gain.Convergence rises from 18 to 20 of 25 (p=0.74); these experiments identify initialization, data order, and learning rate as factors in the compression dynamics.Source: §E.4 · p. 28
Loading 2609.07264v1…