{"id":26134330,"url":"https://github.com/mrorigo/pytorch-fftnet","last_synced_at":"2026-04-22T02:32:26.285Z","repository":{"id":281555223,"uuid":"945477881","full_name":"mrorigo/pytorch-fftnet","owner":"mrorigo","description":"FFTNet implementation in Pytorch","archived":false,"fork":false,"pushed_at":"2025-03-09T21:24:40.000Z","size":29,"stargazers_count":2,"open_issues_count":0,"forks_count":0,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-12-17T04:30:49.581Z","etag":null,"topics":["fftnet","language-model","pytorch","self-attention"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/mrorigo.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2025-03-09T14:19:25.000Z","updated_at":"2025-11-12T02:03:15.000Z","dependencies_parsed_at":"2025-03-16T05:45:37.380Z","dependency_job_id":null,"html_url":"https://github.com/mrorigo/pytorch-fftnet","commit_stats":null,"previous_names":["mrorigo/pytorch-fftnet"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/mrorigo/pytorch-fftnet","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mrorigo%2Fpytorch-fftnet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mrorigo%2Fpytorch-fftnet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mrorigo%2Fpytorch-fftnet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mrorigo%2Fpytorch-fftnet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/mrorigo","download_url":"https://codeload.github.com/mrorigo/pytorch-fftnet/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/mrorigo%2Fpytorch-fftnet/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":32118103,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-04-22T00:31:26.853Z","status":"online","status_checked_at":"2026-04-22T02:00:05.693Z","response_time":58,"last_error":null,"robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":true,"can_crawl_api":true,"host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["fftnet","language-model","pytorch","self-attention"],"created_at":"2025-03-11T00:00:03.137Z","updated_at":"2026-04-22T02:32:26.280Z","avatar_url":"https://github.com/mrorigo.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# FFTNet Implementation\n\nThis repository contains a PyTorch implementation of FFTNet, an efficient alternative to self-attention as described in the paper [\"The FFT Strikes Back: An Efficient Alternative to Self-Attention\"](https://arxiv.org/abs/2502.18394).\n\n## Overview\n\nFFTNet replaces the traditional self-attention mechanism with an adaptive spectral filtering framework based on the Fast Fourier Transform (FFT). This approach achieves global token mixing in O(n log n) time, providing a more efficient alternative to self-attention's O(n²) complexity, especially beneficial for long sequences.\n\nThis implementation provides a modular and understandable codebase for researchers and practitioners interested in exploring and utilizing FFTNet. It includes:\n\n*   **Core FFTNet Components:** Implementations of Fourier Transform, Adaptive Spectral Filtering, modReLU activation, and Inverse Fourier Transform.\n*   **Complete FFTNet Layer \u0026 Model:** Ready-to-use FFTNet Layer and full FFTNet model for sequence modeling tasks.\n*   **Utilities:** Tools for tokenization (TikToken wrapper), dataset generation, and model training.\n*   **Example Training \u0026 Generation Scripts:** Scripts to train the model and generate text examples.\n\n## Key Components\n\nThe FFTNet architecture is built upon the following core components, each implemented as a separate PyTorch module in this repository:\n\n1.  **Fourier Transform (FT):**\n    *   Converts the input sequence from the time domain to the frequency domain using the Fast Fourier Transform.\n    *   Implemented in `fftnet/model/components.py` as `FourierTransform`.\n    *   Achieves global token mixing in O(n log n) complexity.\n\n2.  **Adaptive Spectral Filter (ASF):**\n    *   Applies a learnable filter in the frequency domain to emphasize salient frequency components dynamically.\n    *   Implemented in `fftnet/model/components.py` as `AdaptiveSpectralFilter`.\n    *   Uses an MLP to generate an adaptive filter based on the sequence context.\n\n3.  **modReLU Activation:**\n    *   Introduces non-linearity in the complex domain, enhancing model expressivity.\n    *   Applies a ReLU-like function to the magnitude of complex numbers while preserving phase information.\n    *   Implemented in `fftnet/model/components.py` as `ModReLU`.\n\n4.  **Inverse Fourier Transform (IFT):**\n    *   Transforms the filtered frequency domain representation back to the time domain.\n    *   Implemented in `fftnet/model/components.py` as `InverseFourierTransform`.\n    *   Returns the real component of the inverse transformed signal.\n\nThese components are combined within the `FFTNetLayer` in `fftnet/model/fftnet_layer.py` to form a complete FFTNet processing layer.\n\n## Repository Structure\n\n```\nfftnet/\n├── fftnet/\n│   ├── __init__.py\n│   ├── model/\n│   │   ├── __init__.py\n│   │   ├── components.py     # Core FFTNet components (FT, ASF, ModReLU, IFT)\n│   │   ├── fftnet_layer.py   # Complete FFTNet layer\n│   │   └── fftnet_model.py   # Full FFTNet model\n│   └── utils/\n│       ├── __init__.py\n│       ├── dataset.py        # Dataset classes (TextDataset, Synthetic) and dataloader creation\n│       ├── tokenizer.py      # TikToken wrapper for tokenization\n│       └── trainer.py        # Trainer class for model training and evaluation\n├── example.py                # Example script for text generation\n├── main.py                   # Main training script\n├── README.md                 # This file\n├── requirements.txt          # Project dependencies\n├── setup.py                  # Installation script\n└── .gitignore                # Git ignore file\n```\n\n## Requirements\n\nBefore running the code, ensure you have the following dependencies installed. You can install them using pip:\n\n```bash\npip install -r requirements.txt\n```\n\nThe `requirements.txt` file lists the necessary packages:\n\n-   `torch\u003e=1.10.0`: PyTorch deep learning framework.\n-   `tiktoken`:  Fast BPE tokenizer, compatible with OpenAI models.\n-   `matplotlib`: Plotting library for visualization (e.g., loss plots).\n-   `numpy`: Numerical computation library.\n-   `tqdm`: Progress bar library for training loops.\n\n## Installation\n\nYou can install the package directly from the source:\n\n```bash\npip install -e .\n```\n\n## Usage\n\n### Training FFTNet\n\nTo train the FFTNet model, use the `main.py` script. You can customize training and model parameters using command-line arguments.\n\n```bash\npython main.py --d_model 128 --num_layers 4 --num_epochs 30 --batch_size 64 --lr 1e-4 --validate_every 5 --save_model --output_dir ./outputs\n```\n\n**Available arguments:**\n\n*   `--d_model`:  Model dimension (default: 128).\n*   `--num_layers`: Number of FFTNet layers (default: 4).\n*   `--mlp_hidden_dim`: MLP hidden dimension in Adaptive Spectral Filter (default: 512).\n*   `--max_seq_length`: Maximum sequence length the model can handle (default: 128).\n*   `--dropout`: Dropout rate (default: 0.1).\n*   `--tokenizer`: Tokenizer encoding name (default: \"cl100k_base\").\n*   `--seq_length`: Sequence length for the dataset (default: 128).\n*   `--text_file`: Path to the text corpus file (default: `tinyshakespeare.txt`). The script will download `tinyshakespeare.txt` from a public source if not found.\n*   `--seed`: Random seed for reproducibility (default: 42).\n*   `--batch_size`: Batch size for training (default: 64).\n*   `--lr`: Learning rate for the optimizer (default: 1e-4).\n*   `--weight_decay`: Weight decay for the optimizer (default: 0.01).\n*   `--num_epochs`: Number of training epochs (default: 30).\n*   `--validate_every`:  Validation frequency in epochs (default: 1).\n*   `--device`:  Device to use for training (`cpu`, `cuda`, or `mps`, default: `cuda` if available).\n*   `--output_dir`: Directory to save outputs (models, logs, plots) (default: `./outputs`).\n*   `--save_model`: Flag to save the best model during training.\n*   `--resume`: Path to a checkpoint file to resume training from.\n\n**Outputs of training:**\n\n*   **Model Checkpoints:** If `--save_model` is used, the best model (based on validation loss) is saved as `best_model.pth` in the `--output_dir`. Checkpoints are also saved periodically during training, including optimizer state and training progress, to allow for resuming training.\n*   **Loss Plots:** Training and validation loss curves are saved as `loss_plot.png` in the `--output_dir`.\n*   **TensorBoard Logs:** Training progress, including losses, learning rate, and optionally histograms of gradients and parameters, are logged to TensorBoard in the `logs` subdirectory of `--output_dir`. You can visualize these logs by running `tensorboard --logdir=outputs/logs` (or your specified `output_dir`).\n*   **Test Metrics:** After training, the model is evaluated on the test set, and the test loss, perplexity, and accuracy are saved to `test_metrics.txt` in the `--output_dir`.\n\n### Text Generation Example\n\nAfter training, or using a pre-trained model, you can generate text using the `example.py` script. This script loads a trained model (if available) and generates text based on a prompt.\n\n```bash\npython example.py --prompt \"Once upon a time\" --temperature 0.7 --max_length 100\n```\n\n**Available arguments:**\n\n*   `--d_model`:  Model dimension (default: 128).\n*   `--num_layers`: Number of FFTNet layers (default: 4).\n*   `--mlp_hidden_dim`: MLP hidden dimension (default: 512).\n*   `--max_seq_length`: Maximum sequence length (default: 1024).\n*   `--temperature`: Sampling temperature - higher values create more diverse outputs (default: 1.0).\n*   `--max_length`: Maximum generation length (default: 50).\n*   `--prompt`: Text prompt to continue from (default: \"Hello\").\n*   `--device`: Device to run inference on (default: \"mps\" or available device).\n*   `--tokenizer`: Tokenizer encoding name (default: \"cl100k_base\").\n*   `--top_k`: Top-k sampling parameter (default: 0, disabled).\n*   `--model_path`: Path to model checkpoint (default: \"outputs/best_model.pth\").\n\nThe script automatically attempts to load a trained model from the specified path. If a trained model is not found, it will use an untrained model for demonstration purposes (which will produce random text).\n\n**Example Output:**\n\n```\nLoaded trained model from outputs/best_model.pth\nModel has 189,923 parameters\nGenerating text with:\n- Temperature: 0.7\n- Top-k: 40\n- Max length: 50\n\n==================================================\nPrompt: Once upon a time\nGenerated:  there lived a king and queen, who had a fair daughter. And the king said unto the kingdom, \"Whosoever shall bring me the head of the dragon shall have my daughter to wife.\"\n==================================================\n```\n\n## Model Architecture\n\nThe FFTNet model architecture is designed for efficient sequence modeling, especially for long sequences. It comprises the following key layers:\n\n1.  **Embedding Layer:**\n    *   Maps input token indices to dense vector representations.\n    *   Initialized using `nn.Embedding` from PyTorch.\n    *   Scales embeddings by √d_model as per Transformer conventions.\n    *   Uses sinusoidal positional encodings to incorporate sequence order information.\n\n2.  **FFTNet Layers (Stack):**\n    *   A stack of `FFTNetLayer` modules, where each layer consists of the core FFTNet components.\n    *   Number of layers is configurable via `num_layers` argument.\n    *   Each `FFTNetLayer` sequentially applies:\n        *   Layer Normalization (optional)\n        *   Fourier Transform\n        *   Adaptive Spectral Filter\n        *   ModReLU activation\n        *   Inverse Fourier Transform\n        *   Dropout\n        *   Residual Connection\n        *   Output Layer Normalization (optional)\n\n3.  **Output Projection:**\n    *   A linear layer (`nn.Linear`) followed by Layer Normalization to project the final hidden state to logits over the vocabulary.\n    *   Used for next-token prediction tasks.\n\nThis architecture effectively replaces the self-attention layers of a standard Transformer with the FFT-based spectral filtering framework, providing a more computationally efficient approach for processing sequential data.\n\n## References\n\n*   **The FFT Strikes Back: An Efficient Alternative to Self-Attention**\n    *   arXiv preprint: [https://arxiv.org/abs/2307.00302v2](https://arxiv.org/abs/2307.00302v2)\n    *   Author: Jacob Fein-Ashley (University of Southern California)\n\n## License\n\nThis project is released under the MIT License.\n\n## Contributing\n\nContributions are welcome! Please feel free to submit a Pull Request.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmrorigo%2Fpytorch-fftnet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmrorigo%2Fpytorch-fftnet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmrorigo%2Fpytorch-fftnet/lists"}