NATO: Fourier Spectral Regularization

NATO — Fourier Spectral Regularization (FSR) with selective high-frequency penalty, N-D FFT gradient filtering, and experimental directional regularization for PyTorch.

Features

• 📊 Fourier Spectral Penalty (FSP): Selective high-frequency penalty (Definition 3.2) with configurable τ cutoff, hypercube & radial masks
• 🔧 Low-Pass Gradient Filtering: N-D FFT-based gradient smoothing with correct frequency-domain masking
• 🚀 NATOOptimizer: Custom Adam-variant optimizer with tethered updates (experimental)
• 🎯 Kakeya Directional Penalty: Gradient direction consistency regularization (experimental)
• ⚡ GPU Accelerated: Full CUDA support for all operations

Installation

From PyPI

pip install nato-opt

From Source (Editable)

git clone https://github.com/Malhar1912/NATO.git
cd NATO
pip install -e .

Quick Start

import torch
import torch.nn as nn
from nato_opt import fourier_spectral_penalty, low_pass_filter_gradients

model = nn.Sequential(nn.Conv2d(3, 16, 3), nn.ReLU(), nn.Linear(16, 10))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Selective high-frequency penalty (Definition 3.2)
    # tau=2 penalizes only high-frequency weight components
    fsp = fourier_spectral_penalty(model, lambda_fsp=1e-6, tau=2, mask_mode="hypercube")
    total_loss = loss + fsp
    total_loss.backward()

    # Low-pass filter gradients before optimizer step (Theorem 9.1)
    low_pass_filter_gradients(model, tau_ratio=0.5)

    optimizer.step()

Full-spectrum mode (ℓ₂ weight decay, Proposition 3.2)

# tau=None (default) gives P_FSP^full = ||W||_F² — equivalent to weight decay
fsp = fourier_spectral_penalty(model, lambda_fsp=1e-6, tau=None)

API Reference

fourier_spectral_penalty

fourier_spectral_penalty(
    model,
    lambda_fsp=1e-6,
    tau=None,             # None = full-spectrum (weight decay), float = selective HF
    mask_mode="hypercube", # "hypercube" (ℓ∞, Def 2.2) or "radial" (ℓ₂, §3.4)
    include_conv=True,
    include_linear=True,
    module_whitelist=None,
    module_blacklist=None,
    device=None
) -> torch.Tensor

Compute Fourier Spectral Penalty on model weights. The returned tensor retains grad_fn so autograd can differentiate through it.

Parameters:

• model: PyTorch model
• lambda_fsp: Penalty coefficient (default: 1e-6)
• tau: Frequency cutoff. None = full-spectrum (≡ weight decay). float >= 0 = selective HF penalty.
• mask_mode: "hypercube" (ℓ∞ cutoff) or "radial" (ℓ₂ cutoff)
• include_conv: Include Conv layers (default: True)
• include_linear: Include Linear layers (default: True)
• module_whitelist: Only include these module names
• module_blacklist: Exclude these module names

Returns: Scalar penalty tensor (with grad_fn)

---

low_pass_filter_gradients

low_pass_filter_gradients(
    model,
    tau_ratio=0.5,    # ρ(τ): fraction of frequencies to keep per dimension
    skip_bias=True,
    in_place=True
)

Apply low-pass FFT filtering to gradients, smoothing high-frequency noise. Uses fftshift/ifftshift internally so the centered mask correctly selects near-DC frequencies.

Parameters:

• model: PyTorch model with computed gradients
• tau_ratio: Frequency retention ratio ρ(τ), 0 < tau_ratio ≤ 1. Lower = more filtering.

---

adjust_learning_rate

adjust_learning_rate(optimizer, epoch, ...)

Utility function for learning rate scheduling.

---

Experimental Components

Note: The following components are part of the broader DSR (Directional–Spectral Regularization) research hypothesis described in DSR_Concept_Note.md. They are not covered by the FSR paper's theoretical guarantees.

NATOOptimizer

NATOOptimizer(params, lr=1e-3, beta1=0.9, beta2=0.999,
              epsilon=1e-8, gamma=0.01, tether_interval=100, ...)

Custom Adam-variant optimizer with a tether term that penalizes drift from a periodic parameter checkpoint.

kakeya_directional_penalty

kakeya_directional_penalty(
    model,
    state,          # persistent dict for gradient history
    lambda_k=1e-4
) -> torch.Tensor

Penalizes gradients that maintain high cosine similarity with previous gradients.

Usage (must be added to loss before .backward()):

kakeya_state = {}  # persistent across steps

for inputs, targets in dataloader:
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, targets)

    # Kakeya uses stored gradients from the PREVIOUS step
    k_penalty = kakeya_directional_penalty(model, kakeya_state)

    total_loss = loss + k_penalty
    total_loss.backward()
    optimizer.step()

Requirements

• Python >= 3.8
• PyTorch
• NumPy

License

MIT License - see LICENSE for details.

Authors

• Malhar Pangarkar - malharpangarkar19@gmail.com
• Atharva Khambete - atharvakhambete1@gmail.com

Citation

If you use this in your research, please cite:

@software{nato_opt,
  title = {NATO: Fourier Spectral Regularization},
  author = {Pangarkar, Malhar and Khambete, Atharva},
  year = {2026},
  url = {https://github.com/Malhar1912/NATO}
}

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

1. Fork the repository
2. Create your feature branch (git checkout -b feature/AmazingFeature)
3. Commit your changes (git commit -m 'Add some AmazingFeature')
4. Push to the branch (git push origin feature/AmazingFeature)
5. Open a Pull Request